142 lines
3.8 KiB
Python
142 lines
3.8 KiB
Python
"""SO(3) group transformations."""
|
|
|
|
import kornia.geometry.conversions as C
|
|
import torch
|
|
from torch import Tensor
|
|
from math import pi as PI
|
|
|
|
|
|
@torch.jit.script
|
|
def quat_to_mat(quat_wxyz: Tensor) -> Tensor:
|
|
"""Convert scalar first quaternion to rotation matrix.
|
|
|
|
Args:
|
|
quat_wxyz: (...,4) Scalar first quaternions.
|
|
|
|
Returns:
|
|
(...,3,3) 3D rotation matrices.
|
|
"""
|
|
return C.quaternion_to_rotation_matrix(
|
|
quat_wxyz, order=C.QuaternionCoeffOrder.WXYZ
|
|
)
|
|
|
|
|
|
# @torch.jit.script
|
|
def mat_to_quat(mat: Tensor) -> Tensor:
|
|
"""Convert rotation matrix to scalar first quaternion.
|
|
|
|
Args:
|
|
mat: (...,3,3) 3D rotation matrices.
|
|
|
|
Returns:
|
|
(...,4) Scalar first quaternions.
|
|
"""
|
|
return C.rotation_matrix_to_quaternion(
|
|
mat, order=C.QuaternionCoeffOrder.WXYZ
|
|
)
|
|
|
|
|
|
@torch.jit.script
|
|
def quat_to_xyz(
|
|
quat_wxyz: Tensor, singularity_value: float = PI / 2
|
|
) -> Tensor:
|
|
"""Convert scalar first quaternion to Tait-Bryan angles.
|
|
|
|
Reference:
|
|
https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Source_code_2
|
|
|
|
Args:
|
|
quat_wxyz: (...,4) Scalar first quaternions.
|
|
singularity_value: Value that's set at the singularities.
|
|
|
|
Returns:
|
|
(...,3) The Tait-Bryan angles --- roll, pitch, and yaw.
|
|
"""
|
|
qw = quat_wxyz[..., 0]
|
|
qx = quat_wxyz[..., 1]
|
|
qy = quat_wxyz[..., 2]
|
|
qz = quat_wxyz[..., 3]
|
|
|
|
# roll (x-axis rotation)
|
|
sinr_cosp = 2 * (qw * qx + qy * qz)
|
|
cosr_cosp = 1 - 2 * (qx * qx + qy * qy)
|
|
roll = torch.atan2(sinr_cosp, cosr_cosp)
|
|
|
|
# pitch (y-axis rotation)
|
|
pitch = 2 * (qw * qy - qz * qx)
|
|
is_out_of_range = torch.abs(pitch) >= 1
|
|
pitch[is_out_of_range] = torch.copysign(
|
|
torch.as_tensor(singularity_value), pitch[is_out_of_range]
|
|
)
|
|
pitch[~is_out_of_range] = torch.asin(pitch[~is_out_of_range])
|
|
|
|
# yaw (z-axis rotation)
|
|
siny_cosp = 2 * (qw * qz + qx * qy)
|
|
cosy_cosp = 1 - 2 * (qy * qy + qz * qz)
|
|
yaw = torch.atan2(siny_cosp, cosy_cosp)
|
|
xyz = torch.stack([roll, pitch, yaw], dim=-1)
|
|
return xyz
|
|
|
|
|
|
@torch.jit.script
|
|
def quat_to_yaw(quat_wxyz: Tensor) -> Tensor:
|
|
"""Convert scalar first quaternion to yaw (rotation about vertical axis).
|
|
|
|
Reference:
|
|
https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Source_code_2
|
|
|
|
Args:
|
|
quat_wxyz: (...,4) Scalar first quaternions.
|
|
|
|
Returns:
|
|
(...,) The rotation about the z-axis in radians.
|
|
"""
|
|
xyz = quat_to_xyz(quat_wxyz)
|
|
yaw_rad: Tensor = xyz[..., -1]
|
|
return yaw_rad
|
|
|
|
|
|
@torch.jit.script
|
|
def xyz_to_quat(xyz_rad: Tensor) -> Tensor:
|
|
"""Convert euler angles (xyz - pitch, roll, yaw) to scalar first quaternions.
|
|
|
|
Args:
|
|
xyz_rad: (...,3) Tensor of roll, pitch, and yaw in radians.
|
|
|
|
Returns:
|
|
(...,4) Scalar first quaternions (wxyz).
|
|
"""
|
|
x_rad = xyz_rad[..., 0]
|
|
y_rad = xyz_rad[..., 1]
|
|
z_rad = xyz_rad[..., 2]
|
|
|
|
cy = torch.cos(z_rad * 0.5)
|
|
sy = torch.sin(z_rad * 0.5)
|
|
cp = torch.cos(y_rad * 0.5)
|
|
sp = torch.sin(y_rad * 0.5)
|
|
cr = torch.cos(x_rad * 0.5)
|
|
sr = torch.sin(x_rad * 0.5)
|
|
|
|
qw = cr * cp * cy + sr * sp * sy
|
|
qx = sr * cp * cy - cr * sp * sy
|
|
qy = cr * sp * cy + sr * cp * sy
|
|
qz = cr * cp * sy - sr * sp * cy
|
|
quat_wxyz = torch.stack([qw, qx, qy, qz], dim=-1)
|
|
return quat_wxyz
|
|
|
|
|
|
@torch.jit.script
|
|
def yaw_to_quat(yaw_rad: Tensor) -> Tensor:
|
|
"""Convert yaw (rotation about the vertical axis) to scalar first quaternions.
|
|
|
|
Args:
|
|
yaw_rad: (...,1) Rotations about the z-axis.
|
|
|
|
Returns:
|
|
(...,4) scalar first quaternions (wxyz).
|
|
"""
|
|
xyz_rad = torch.zeros_like(yaw_rad)[..., None].repeat_interleave(3, dim=-1)
|
|
xyz_rad[..., -1] = yaw_rad
|
|
quat_wxyz: Tensor = xyz_to_quat(xyz_rad)
|
|
return quat_wxyz
|