Spaces:
Sleeping
Sleeping
| import functools | |
| import torch | |
| import torch.nn.functional as F | |
| def quaternion_to_matrix(quaternions): | |
| r, i, j, k = torch.unbind(quaternions, -1) | |
| two_s = 2.0 / (quaternions * quaternions).sum(-1) | |
| o = torch.stack( | |
| ( | |
| 1 - two_s * (j * j + k * k), | |
| two_s * (i * j - k * r), | |
| two_s * (i * k + j * r), | |
| two_s * (i * j + k * r), | |
| 1 - two_s * (i * i + k * k), | |
| two_s * (j * k - i * r), | |
| two_s * (i * k - j * r), | |
| two_s * (j * k + i * r), | |
| 1 - two_s * (i * i + j * j), | |
| ), | |
| -1, | |
| ) | |
| return o.reshape(quaternions.shape[:-1] + (3, 3)) | |
| def _copysign(a, b): | |
| signs_differ = (a < 0) != (b < 0) | |
| return torch.where(signs_differ, -a, a) | |
| def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor: | |
| ret = torch.zeros_like(x) | |
| positive_mask = x > 0 | |
| ret[positive_mask] = torch.sqrt(x[positive_mask]) | |
| return ret | |
| def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor: | |
| if matrix.size(-1) != 3 or matrix.size(-2) != 3: | |
| raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") | |
| batch_dim = matrix.shape[:-2] | |
| m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( | |
| matrix.reshape(*batch_dim, 9), dim=-1 | |
| ) | |
| q_abs = _sqrt_positive_part( | |
| torch.stack( | |
| [ | |
| 1.0 + m00 + m11 + m22, | |
| 1.0 + m00 - m11 - m22, | |
| 1.0 - m00 + m11 - m22, | |
| 1.0 - m00 - m11 + m22, | |
| ], | |
| dim=-1, | |
| ) | |
| ) | |
| quat_by_rijk = torch.stack( | |
| [ | |
| torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1), | |
| torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1), | |
| torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1), | |
| torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1), | |
| ], | |
| dim=-2, | |
| ) | |
| quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(q_abs.new_tensor(0.1))) | |
| return quat_candidates[ | |
| F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : | |
| ].reshape(*batch_dim, 4) | |
| def _axis_angle_rotation(axis: str, angle): | |
| cos = torch.cos(angle) | |
| sin = torch.sin(angle) | |
| one = torch.ones_like(angle) | |
| zero = torch.zeros_like(angle) | |
| if axis == "X": | |
| R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos) | |
| if axis == "Y": | |
| R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos) | |
| if axis == "Z": | |
| R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one) | |
| return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3)) | |
| def euler_angles_to_matrix(euler_angles, convention: str): | |
| if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3: | |
| raise ValueError("Invalid input euler angles.") | |
| if len(convention) != 3: | |
| raise ValueError("Convention must have 3 letters.") | |
| if convention[1] in (convention[0], convention[2]): | |
| raise ValueError(f"Invalid convention {convention}.") | |
| for letter in convention: | |
| if letter not in ("X", "Y", "Z"): | |
| raise ValueError(f"Invalid letter {letter} in convention string.") | |
| matrices = map(_axis_angle_rotation, convention, torch.unbind(euler_angles, -1)) | |
| return functools.reduce(torch.matmul, matrices) | |
| def _angle_from_tan( | |
| axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool | |
| ): | |
| i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis] | |
| if horizontal: | |
| i2, i1 = i1, i2 | |
| even = (axis + other_axis) in ["XY", "YZ", "ZX"] | |
| if horizontal == even: | |
| return torch.atan2(data[..., i1], data[..., i2]) | |
| if tait_bryan: | |
| return torch.atan2(-data[..., i2], data[..., i1]) | |
| return torch.atan2(data[..., i2], -data[..., i1]) | |
| def _index_from_letter(letter: str): | |
| if letter == "X": | |
| return 0 | |
| if letter == "Y": | |
| return 1 | |
| if letter == "Z": | |
| return 2 | |
| def matrix_to_euler_angles(matrix, convention: str): | |
| if len(convention) != 3: | |
| raise ValueError("Convention must have 3 letters.") | |
| if convention[1] in (convention[0], convention[2]): | |
| raise ValueError(f"Invalid convention {convention}.") | |
| for letter in convention: | |
| if letter not in ("X", "Y", "Z"): | |
| raise ValueError(f"Invalid letter {letter} in convention string.") | |
| if matrix.size(-1) != 3 or matrix.size(-2) != 3: | |
| raise ValueError(f"Invalid rotation matrix shape f{matrix.shape}.") | |
| i0 = _index_from_letter(convention[0]) | |
| i2 = _index_from_letter(convention[2]) | |
| tait_bryan = i0 != i2 | |
| if tait_bryan: | |
| central_angle = torch.asin( | |
| matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0) | |
| ) | |
| else: | |
| central_angle = torch.acos(matrix[..., i0, i0]) | |
| o = ( | |
| _angle_from_tan( | |
| convention[0], convention[1], matrix[..., i2], False, tait_bryan | |
| ), | |
| central_angle, | |
| _angle_from_tan( | |
| convention[2], convention[1], matrix[..., i0, :], True, tait_bryan | |
| ), | |
| ) | |
| return torch.stack(o, -1) | |
| def standardize_quaternion(quaternions): | |
| return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions) | |
| def quaternion_raw_multiply(a, b): | |
| 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) | |
| def quaternion_multiply(a, b): | |
| ab = quaternion_raw_multiply(a, b) | |
| return standardize_quaternion(ab) | |
| def quaternion_invert(quaternion): | |
| return quaternion * quaternion.new_tensor([1, -1, -1, -1]) | |
| def quaternion_apply(quaternion, point): | |
| if point.size(-1) != 3: | |
| raise ValueError(f"Points are not in 3D, f{point.shape}.") | |
| real_parts = point.new_zeros(point.shape[:-1] + (1,)) | |
| point_as_quaternion = torch.cat((real_parts, point), -1) | |
| out = quaternion_raw_multiply( | |
| quaternion_raw_multiply(quaternion, point_as_quaternion), | |
| quaternion_invert(quaternion), | |
| ) | |
| return out[..., 1:] | |
| def axis_angle_to_matrix(axis_angle): | |
| return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle)) | |
| def matrix_to_axis_angle(matrix): | |
| return quaternion_to_axis_angle(matrix_to_quaternion(matrix)) | |
| def axis_angle_to_quaternion(axis_angle): | |
| angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True) | |
| half_angles = 0.5 * angles | |
| eps = 1e-6 | |
| small_angles = angles.abs() < eps | |
| sin_half_angles_over_angles = torch.empty_like(angles) | |
| sin_half_angles_over_angles[~small_angles] = ( | |
| torch.sin(half_angles[~small_angles]) / angles[~small_angles] | |
| ) | |
| # for x small, sin(x/2) is about x/2 - (x/2)^3/6 | |
| # so sin(x/2)/x is about 1/2 - (x*x)/48 | |
| sin_half_angles_over_angles[small_angles] = ( | |
| 0.5 - (angles[small_angles] * angles[small_angles]) / 48 | |
| ) | |
| quaternions = torch.cat( | |
| [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1 | |
| ) | |
| return quaternions | |
| def quaternion_to_axis_angle(quaternions): | |
| norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True) | |
| half_angles = torch.atan2(norms, quaternions[..., :1]) | |
| angles = 2 * half_angles | |
| eps = 1e-6 | |
| small_angles = angles.abs() < eps | |
| sin_half_angles_over_angles = torch.empty_like(angles) | |
| sin_half_angles_over_angles[~small_angles] = ( | |
| torch.sin(half_angles[~small_angles]) / angles[~small_angles] | |
| ) | |
| # for x small, sin(x/2) is about x/2 - (x/2)^3/6 | |
| # so sin(x/2)/x is about 1/2 - (x*x)/48 | |
| sin_half_angles_over_angles[small_angles] = ( | |
| 0.5 - (angles[small_angles] * angles[small_angles]) / 48 | |
| ) | |
| return quaternions[..., 1:] / sin_half_angles_over_angles | |
| def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor: | |
| a1, a2 = d6[..., :3], d6[..., 3:] | |
| b1 = F.normalize(a1, dim=-1) | |
| b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1 | |
| b2 = F.normalize(b2, dim=-1) | |
| b3 = torch.cross(b1, b2, dim=-1) | |
| return torch.stack((b1, b2, b3), dim=-2) | |
| def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor: | |
| return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6) | |