Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| import torchgeometry as tgm | |
| __all__ = [ | |
| # functional api | |
| "pi", | |
| "rad2deg", | |
| "deg2rad", | |
| "convert_points_from_homogeneous", | |
| "convert_points_to_homogeneous", | |
| "angle_axis_to_rotation_matrix", | |
| "rotation_matrix_to_angle_axis", | |
| "rotation_matrix_to_quaternion", | |
| "quaternion_to_angle_axis", | |
| "angle_axis_to_quaternion", | |
| "rtvec_to_pose", | |
| # layer api | |
| "RadToDeg", | |
| "DegToRad", | |
| "ConvertPointsFromHomogeneous", | |
| "ConvertPointsToHomogeneous", | |
| ] | |
| """Constant with number pi | |
| """ | |
| pi = torch.Tensor([3.14159265358979323846]) | |
| def rad2deg(tensor): | |
| r"""Function that converts angles from radians to degrees. | |
| See :class:`~torchgeometry.RadToDeg` for details. | |
| Args: | |
| tensor (Tensor): Tensor of arbitrary shape. | |
| Returns: | |
| Tensor: Tensor with same shape as input. | |
| Example: | |
| >>> input = tgm.pi * torch.rand(1, 3, 3) | |
| >>> output = tgm.rad2deg(input) | |
| """ | |
| if not torch.is_tensor(tensor): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}" | |
| .format(type(tensor))) | |
| return 180. * tensor / pi.to(tensor.device).type(tensor.dtype) | |
| def deg2rad(tensor): | |
| r"""Function that converts angles from degrees to radians. | |
| See :class:`~torchgeometry.DegToRad` for details. | |
| Args: | |
| tensor (Tensor): Tensor of arbitrary shape. | |
| Returns: | |
| Tensor: Tensor with same shape as input. | |
| Examples:: | |
| >>> input = 360. * torch.rand(1, 3, 3) | |
| >>> output = tgm.deg2rad(input) | |
| """ | |
| if not torch.is_tensor(tensor): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}" | |
| .format(type(tensor))) | |
| return tensor * pi.to(tensor.device).type(tensor.dtype) / 180. | |
| def convert_points_from_homogeneous(points): | |
| r"""Function that converts points from homogeneous to Euclidean space. | |
| See :class:`~torchgeometry.ConvertPointsFromHomogeneous` for details. | |
| Examples:: | |
| >>> input = torch.rand(2, 4, 3) # BxNx3 | |
| >>> output = tgm.convert_points_from_homogeneous(input) # BxNx2 | |
| """ | |
| if not torch.is_tensor(points): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}".format( | |
| type(points))) | |
| if len(points.shape) < 2: | |
| raise ValueError("Input must be at least a 2D tensor. Got {}".format( | |
| points.shape)) | |
| return points[..., :-1] / points[..., -1:] | |
| def convert_points_to_homogeneous(points): | |
| r"""Function that converts points from Euclidean to homogeneous space. | |
| See :class:`~torchgeometry.ConvertPointsToHomogeneous` for details. | |
| Examples:: | |
| >>> input = torch.rand(2, 4, 3) # BxNx3 | |
| >>> output = tgm.convert_points_to_homogeneous(input) # BxNx4 | |
| """ | |
| if not torch.is_tensor(points): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}".format( | |
| type(points))) | |
| if len(points.shape) < 2: | |
| raise ValueError("Input must be at least a 2D tensor. Got {}".format( | |
| points.shape)) | |
| return nn.functional.pad(points, (0, 1), "constant", 1.0) | |
| def angle_axis_to_rotation_matrix(angle_axis): | |
| """Convert 3d vector of axis-angle rotation to 4x4 rotation matrix | |
| Args: | |
| angle_axis (Tensor): tensor of 3d vector of axis-angle rotations. | |
| Returns: | |
| Tensor: tensor of 4x4 rotation matrices. | |
| Shape: | |
| - Input: :math:`(N, 3)` | |
| - Output: :math:`(N, 4, 4)` | |
| Example: | |
| >>> input = torch.rand(1, 3) # Nx3 | |
| >>> output = tgm.angle_axis_to_rotation_matrix(input) # Nx4x4 | |
| """ | |
| def _compute_rotation_matrix(angle_axis, theta2, eps=1e-6): | |
| # We want to be careful to only evaluate the square root if the | |
| # norm of the angle_axis vector is greater than zero. Otherwise | |
| # we get a division by zero. | |
| k_one = 1.0 | |
| theta = torch.sqrt(theta2) | |
| wxyz = angle_axis / (theta + eps) | |
| wx, wy, wz = torch.chunk(wxyz, 3, dim=1) | |
| cos_theta = torch.cos(theta) | |
| sin_theta = torch.sin(theta) | |
| r00 = cos_theta + wx * wx * (k_one - cos_theta) | |
| r10 = wz * sin_theta + wx * wy * (k_one - cos_theta) | |
| r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta) | |
| r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta | |
| r11 = cos_theta + wy * wy * (k_one - cos_theta) | |
| r21 = wx * sin_theta + wy * wz * (k_one - cos_theta) | |
| r02 = wy * sin_theta + wx * wz * (k_one - cos_theta) | |
| r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta) | |
| r22 = cos_theta + wz * wz * (k_one - cos_theta) | |
| rotation_matrix = torch.cat( | |
| [r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1) | |
| return rotation_matrix.view(-1, 3, 3) | |
| def _compute_rotation_matrix_taylor(angle_axis): | |
| rx, ry, rz = torch.chunk(angle_axis, 3, dim=1) | |
| k_one = torch.ones_like(rx) | |
| rotation_matrix = torch.cat( | |
| [k_one, -rz, ry, rz, k_one, -rx, -ry, rx, k_one], dim=1) | |
| return rotation_matrix.view(-1, 3, 3) | |
| # stolen from ceres/rotation.h | |
| _angle_axis = torch.unsqueeze(angle_axis, dim=1) | |
| theta2 = torch.matmul(_angle_axis, _angle_axis.transpose(1, 2)) | |
| theta2 = torch.squeeze(theta2, dim=1) | |
| # compute rotation matrices | |
| rotation_matrix_normal = _compute_rotation_matrix(angle_axis, theta2) | |
| rotation_matrix_taylor = _compute_rotation_matrix_taylor(angle_axis) | |
| # create mask to handle both cases | |
| eps = 1e-6 | |
| mask = (theta2 > eps).view(-1, 1, 1).to(theta2.device) | |
| mask_pos = (mask).type_as(theta2) | |
| mask_neg = (mask == False).type_as(theta2) # noqa | |
| # create output pose matrix | |
| batch_size = angle_axis.shape[0] | |
| rotation_matrix = torch.eye(4).to(angle_axis.device).type_as(angle_axis) | |
| rotation_matrix = rotation_matrix.view(1, 4, 4).repeat(batch_size, 1, 1) | |
| # fill output matrix with masked values | |
| rotation_matrix[..., :3, :3] = \ | |
| mask_pos * rotation_matrix_normal + mask_neg * rotation_matrix_taylor | |
| return rotation_matrix # Nx4x4 | |
| def rtvec_to_pose(rtvec): | |
| """ | |
| Convert axis-angle rotation and translation vector to 4x4 pose matrix | |
| Args: | |
| rtvec (Tensor): Rodrigues vector transformations | |
| Returns: | |
| Tensor: transformation matrices | |
| Shape: | |
| - Input: :math:`(N, 6)` | |
| - Output: :math:`(N, 4, 4)` | |
| Example: | |
| >>> input = torch.rand(3, 6) # Nx6 | |
| >>> output = tgm.rtvec_to_pose(input) # Nx4x4 | |
| """ | |
| assert rtvec.shape[-1] == 6, 'rtvec=[rx, ry, rz, tx, ty, tz]' | |
| pose = angle_axis_to_rotation_matrix(rtvec[..., :3]) | |
| pose[..., :3, 3] = rtvec[..., 3:] | |
| return pose | |
| def rotation_matrix_to_angle_axis(rotation_matrix): | |
| """Convert 3x4 rotation matrix to Rodrigues vector | |
| Args: | |
| rotation_matrix (Tensor): rotation matrix. | |
| Returns: | |
| Tensor: Rodrigues vector transformation. | |
| Shape: | |
| - Input: :math:`(N, 3, 4)` | |
| - Output: :math:`(N, 3)` | |
| Example: | |
| >>> input = torch.rand(2, 3, 4) # Nx4x4 | |
| >>> output = tgm.rotation_matrix_to_angle_axis(input) # Nx3 | |
| """ | |
| # todo add check that matrix is a valid rotation matrix | |
| quaternion = rotation_matrix_to_quaternion(rotation_matrix) | |
| return quaternion_to_angle_axis(quaternion) | |
| def rotation_matrix_to_quaternion(rotation_matrix, eps=1e-6): | |
| """Convert 3x4 rotation matrix to 4d quaternion vector | |
| This algorithm is based on algorithm described in | |
| https://github.com/KieranWynn/pyquaternion/blob/master/pyquaternion/quaternion.py#L201 | |
| Args: | |
| rotation_matrix (Tensor): the rotation matrix to convert. | |
| Return: | |
| Tensor: the rotation in quaternion | |
| Shape: | |
| - Input: :math:`(N, 3, 4)` | |
| - Output: :math:`(N, 4)` | |
| Example: | |
| >>> input = torch.rand(4, 3, 4) # Nx3x4 | |
| >>> output = tgm.rotation_matrix_to_quaternion(input) # Nx4 | |
| """ | |
| if not torch.is_tensor(rotation_matrix): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}".format( | |
| type(rotation_matrix))) | |
| if len(rotation_matrix.shape) > 3: | |
| raise ValueError( | |
| "Input size must be a three dimensional tensor. Got {}".format( | |
| rotation_matrix.shape)) | |
| if not rotation_matrix.shape[-2:] == (3, 4): | |
| raise ValueError( | |
| "Input size must be a N x 3 x 4 tensor. Got {}".format( | |
| rotation_matrix.shape)) | |
| rmat_t = torch.transpose(rotation_matrix, 1, 2) | |
| mask_d2 = rmat_t[:, 2, 2] < eps | |
| mask_d0_d1 = rmat_t[:, 0, 0] > rmat_t[:, 1, 1] | |
| mask_d0_nd1 = rmat_t[:, 0, 0] < -rmat_t[:, 1, 1] | |
| t0 = 1 + rmat_t[:, 0, 0] - rmat_t[:, 1, 1] - rmat_t[:, 2, 2] | |
| q0 = torch.stack([rmat_t[:, 1, 2] - rmat_t[:, 2, 1], | |
| t0, rmat_t[:, 0, 1] + rmat_t[:, 1, 0], | |
| rmat_t[:, 2, 0] + rmat_t[:, 0, 2]], -1) | |
| t0_rep = t0.repeat(4, 1).t() | |
| t1 = 1 - rmat_t[:, 0, 0] + rmat_t[:, 1, 1] - rmat_t[:, 2, 2] | |
| q1 = torch.stack([rmat_t[:, 2, 0] - rmat_t[:, 0, 2], | |
| rmat_t[:, 0, 1] + rmat_t[:, 1, 0], | |
| t1, rmat_t[:, 1, 2] + rmat_t[:, 2, 1]], -1) | |
| t1_rep = t1.repeat(4, 1).t() | |
| t2 = 1 - rmat_t[:, 0, 0] - rmat_t[:, 1, 1] + rmat_t[:, 2, 2] | |
| q2 = torch.stack([rmat_t[:, 0, 1] - rmat_t[:, 1, 0], | |
| rmat_t[:, 2, 0] + rmat_t[:, 0, 2], | |
| rmat_t[:, 1, 2] + rmat_t[:, 2, 1], t2], -1) | |
| t2_rep = t2.repeat(4, 1).t() | |
| t3 = 1 + rmat_t[:, 0, 0] + rmat_t[:, 1, 1] + rmat_t[:, 2, 2] | |
| q3 = torch.stack([t3, rmat_t[:, 1, 2] - rmat_t[:, 2, 1], | |
| rmat_t[:, 2, 0] - rmat_t[:, 0, 2], | |
| rmat_t[:, 0, 1] - rmat_t[:, 1, 0]], -1) | |
| t3_rep = t3.repeat(4, 1).t() | |
| mask_c0 = mask_d2 * mask_d0_d1 | |
| mask_c1 = mask_d2 * ~(mask_d0_d1) | |
| mask_c2 = ~(mask_d2) * mask_d0_nd1 | |
| mask_c3 = ~(mask_d2) * ~(mask_d0_nd1) | |
| mask_c0 = mask_c0.view(-1, 1).type_as(q0) | |
| mask_c1 = mask_c1.view(-1, 1).type_as(q1) | |
| mask_c2 = mask_c2.view(-1, 1).type_as(q2) | |
| mask_c3 = mask_c3.view(-1, 1).type_as(q3) | |
| q = q0 * mask_c0 + q1 * mask_c1 + q2 * mask_c2 + q3 * mask_c3 | |
| q /= torch.sqrt(t0_rep * mask_c0 + t1_rep * mask_c1 + # noqa | |
| t2_rep * mask_c2 + t3_rep * mask_c3) # noqa | |
| q *= 0.5 | |
| return q | |
| def quaternion_to_angle_axis(quaternion: torch.Tensor) -> torch.Tensor: | |
| """Convert quaternion vector to angle axis of rotation. | |
| Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h | |
| Args: | |
| quaternion (torch.Tensor): tensor with quaternions. | |
| Return: | |
| torch.Tensor: tensor with angle axis of rotation. | |
| Shape: | |
| - Input: :math:`(*, 4)` where `*` means, any number of dimensions | |
| - Output: :math:`(*, 3)` | |
| Example: | |
| >>> quaternion = torch.rand(2, 4) # Nx4 | |
| >>> angle_axis = tgm.quaternion_to_angle_axis(quaternion) # Nx3 | |
| """ | |
| if not torch.is_tensor(quaternion): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}".format( | |
| type(quaternion))) | |
| if not quaternion.shape[-1] == 4: | |
| raise ValueError("Input must be a tensor of shape Nx4 or 4. Got {}" | |
| .format(quaternion.shape)) | |
| # unpack input and compute conversion | |
| q1: torch.Tensor = quaternion[..., 1] | |
| q2: torch.Tensor = quaternion[..., 2] | |
| q3: torch.Tensor = quaternion[..., 3] | |
| sin_squared_theta: torch.Tensor = q1 * q1 + q2 * q2 + q3 * q3 | |
| sin_theta: torch.Tensor = torch.sqrt(sin_squared_theta) | |
| cos_theta: torch.Tensor = quaternion[..., 0] | |
| two_theta: torch.Tensor = 2.0 * torch.where( | |
| cos_theta < 0.0, | |
| torch.atan2(-sin_theta, -cos_theta), | |
| torch.atan2(sin_theta, cos_theta)) | |
| k_pos: torch.Tensor = two_theta / sin_theta | |
| k_neg: torch.Tensor = 2.0 * torch.ones_like(sin_theta) | |
| k: torch.Tensor = torch.where(sin_squared_theta > 0.0, k_pos, k_neg) | |
| angle_axis: torch.Tensor = torch.zeros_like(quaternion)[..., :3] | |
| angle_axis[..., 0] += q1 * k | |
| angle_axis[..., 1] += q2 * k | |
| angle_axis[..., 2] += q3 * k | |
| return angle_axis | |
| # based on: | |
| # https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py#L138 | |
| def angle_axis_to_quaternion(angle_axis: torch.Tensor) -> torch.Tensor: | |
| """Convert an angle axis to a quaternion. | |
| Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h | |
| Args: | |
| angle_axis (torch.Tensor): tensor with angle axis. | |
| Return: | |
| torch.Tensor: tensor with quaternion. | |
| Shape: | |
| - Input: :math:`(*, 3)` where `*` means, any number of dimensions | |
| - Output: :math:`(*, 4)` | |
| Example: | |
| >>> angle_axis = torch.rand(2, 4) # Nx4 | |
| >>> quaternion = tgm.angle_axis_to_quaternion(angle_axis) # Nx3 | |
| """ | |
| if not torch.is_tensor(angle_axis): | |
| raise TypeError("Input type is not a torch.Tensor. Got {}".format( | |
| type(angle_axis))) | |
| if not angle_axis.shape[-1] == 3: | |
| raise ValueError("Input must be a tensor of shape Nx3 or 3. Got {}" | |
| .format(angle_axis.shape)) | |
| # unpack input and compute conversion | |
| a0: torch.Tensor = angle_axis[..., 0:1] | |
| a1: torch.Tensor = angle_axis[..., 1:2] | |
| a2: torch.Tensor = angle_axis[..., 2:3] | |
| theta_squared: torch.Tensor = a0 * a0 + a1 * a1 + a2 * a2 | |
| theta: torch.Tensor = torch.sqrt(theta_squared) | |
| half_theta: torch.Tensor = theta * 0.5 | |
| mask: torch.Tensor = theta_squared > 0.0 | |
| ones: torch.Tensor = torch.ones_like(half_theta) | |
| k_neg: torch.Tensor = 0.5 * ones | |
| k_pos: torch.Tensor = torch.sin(half_theta) / theta | |
| k: torch.Tensor = torch.where(mask, k_pos, k_neg) | |
| w: torch.Tensor = torch.where(mask, torch.cos(half_theta), ones) | |
| quaternion: torch.Tensor = torch.zeros_like(angle_axis) | |
| quaternion[..., 0:1] += a0 * k | |
| quaternion[..., 1:2] += a1 * k | |
| quaternion[..., 2:3] += a2 * k | |
| return torch.cat([w, quaternion], dim=-1) | |
| # TODO: add below funtionalities | |
| # - pose_to_rtvec | |
| # layer api | |
| class RadToDeg(nn.Module): | |
| r"""Creates an object that converts angles from radians to degrees. | |
| Args: | |
| tensor (Tensor): Tensor of arbitrary shape. | |
| Returns: | |
| Tensor: Tensor with same shape as input. | |
| Examples:: | |
| >>> input = tgm.pi * torch.rand(1, 3, 3) | |
| >>> output = tgm.RadToDeg()(input) | |
| """ | |
| def __init__(self): | |
| super(RadToDeg, self).__init__() | |
| def forward(self, input): | |
| return rad2deg(input) | |
| class DegToRad(nn.Module): | |
| r"""Function that converts angles from degrees to radians. | |
| Args: | |
| tensor (Tensor): Tensor of arbitrary shape. | |
| Returns: | |
| Tensor: Tensor with same shape as input. | |
| Examples:: | |
| >>> input = 360. * torch.rand(1, 3, 3) | |
| >>> output = tgm.DegToRad()(input) | |
| """ | |
| def __init__(self): | |
| super(DegToRad, self).__init__() | |
| def forward(self, input): | |
| return deg2rad(input) | |
| class ConvertPointsFromHomogeneous(nn.Module): | |
| r"""Creates a transformation that converts points from homogeneous to | |
| Euclidean space. | |
| Args: | |
| points (Tensor): tensor of N-dimensional points. | |
| Returns: | |
| Tensor: tensor of N-1-dimensional points. | |
| Shape: | |
| - Input: :math:`(B, D, N)` or :math:`(D, N)` | |
| - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)` | |
| Examples:: | |
| >>> input = torch.rand(2, 4, 3) # BxNx3 | |
| >>> transform = tgm.ConvertPointsFromHomogeneous() | |
| >>> output = transform(input) # BxNx2 | |
| """ | |
| def __init__(self): | |
| super(ConvertPointsFromHomogeneous, self).__init__() | |
| def forward(self, input): | |
| return convert_points_from_homogeneous(input) | |
| class ConvertPointsToHomogeneous(nn.Module): | |
| r"""Creates a transformation to convert points from Euclidean to | |
| homogeneous space. | |
| Args: | |
| points (Tensor): tensor of N-dimensional points. | |
| Returns: | |
| Tensor: tensor of N+1-dimensional points. | |
| Shape: | |
| - Input: :math:`(B, D, N)` or :math:`(D, N)` | |
| - Output: :math:`(B, D, N + 1)` or :math:`(D, N + 1)` | |
| Examples:: | |
| >>> input = torch.rand(2, 4, 3) # BxNx3 | |
| >>> transform = tgm.ConvertPointsToHomogeneous() | |
| >>> output = transform(input) # BxNx4 | |
| """ | |
| def __init__(self): | |
| super(ConvertPointsToHomogeneous, self).__init__() | |
| def forward(self, input): | |
| return convert_points_to_homogeneous(input) | |