Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,140 Bytes
4aa5114 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
def quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
"""
From Pytorch3d
Multiply two quaternions.
Usual torch rules for broadcasting apply.
Args:
a: Quaternions as tensor of shape (..., 4), real part first.
b: Quaternions as tensor of shape (..., 4), real part first.
Returns:
The product of a and b, a tensor of quaternions shape (..., 4).
"""
aw, ax, ay, az = torch.unbind(a, -1)
bw, bx, by, bz = torch.unbind(b, -1)
ow = aw * bw - ax * bx - ay * by - az * bz
ox = aw * bx + ax * bw + ay * bz - az * by
oy = aw * by - ax * bz + ay * bw + az * bx
oz = aw * bz + ax * by - ay * bx + az * bw
return torch.stack((ow, ox, oy, oz), -1)
# Written by Stan Szymanowicz 2023
def matrix_to_quaternion(M: torch.Tensor) -> torch.Tensor:
"""
Matrix-to-quaternion conversion method. Equation taken from
https://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/index.htm
Args:
M: rotation matrices, (3 x 3)
Returns:
q: quaternion of shape (4)
"""
tr = 1 + M[ 0, 0] + M[ 1, 1] + M[ 2, 2]
if tr > 0:
r = torch.sqrt(tr) / 2.0
x = ( M[ 2, 1] - M[ 1, 2] ) / ( 4 * r )
y = ( M[ 0, 2] - M[ 2, 0] ) / ( 4 * r )
z = ( M[ 1, 0] - M[ 0, 1] ) / ( 4 * r )
elif ( M[ 0, 0] > M[ 1, 1]) and (M[ 0, 0] > M[ 2, 2]):
S = torch.sqrt(1.0 + M[ 0, 0] - M[ 1, 1] - M[ 2, 2]) * 2 # S=4*qx
r = (M[ 2, 1] - M[ 1, 2]) / S
x = 0.25 * S
y = (M[ 0, 1] + M[ 1, 0]) / S
z = (M[ 0, 2] + M[ 2, 0]) / S
elif M[ 1, 1] > M[ 2, 2]:
S = torch.sqrt(1.0 + M[ 1, 1] - M[ 0, 0] - M[ 2, 2]) * 2 # S=4*qy
r = (M[ 0, 2] - M[ 2, 0]) / S
x = (M[ 0, 1] + M[ 1, 0]) / S
y = 0.25 * S
z = (M[ 1, 2] + M[ 2, 1]) / S
else:
S = torch.sqrt(1.0 + M[ 2, 2] - M[ 0, 0] - M[ 1, 1]) * 2 # S=4*qz
r = (M[ 1, 0] - M[ 0, 1]) / S
x = (M[ 0, 2] + M[ 2, 0]) / S
y = (M[ 1, 2] + M[ 2, 1]) / S
z = 0.25 * S
return torch.stack([r, x, y, z], dim=-1) |