133 lines
4.5 KiB
Python
133 lines
4.5 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
# Modified from PyTorch3D, https://github.com/facebookresearch/pytorch3d
|
|
|
|
import torch
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Quaternion Order: XYZW or say ijkr, scalar-last
|
|
|
|
Convert rotations given as quaternions to rotation matrices.
|
|
Args:
|
|
quaternions: quaternions with real part last,
|
|
as tensor of shape (..., 4).
|
|
|
|
Returns:
|
|
Rotation matrices as tensor of shape (..., 3, 3).
|
|
"""
|
|
i, j, k, r = torch.unbind(quaternions, -1)
|
|
# pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
|
|
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 mat_to_quat(matrix: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Convert rotations given as rotation matrices to quaternions.
|
|
|
|
Args:
|
|
matrix: Rotation matrices as tensor of shape (..., 3, 3).
|
|
|
|
Returns:
|
|
quaternions with real part last, as tensor of shape (..., 4).
|
|
Quaternion Order: XYZW or say ijkr, scalar-last
|
|
"""
|
|
if matrix.size(-1) != 3 or matrix.size(-2) != 3:
|
|
raise ValueError(f"Invalid rotation matrix shape {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
|
|
)
|
|
)
|
|
|
|
# we produce the desired quaternion multiplied by each of r, i, j, k
|
|
quat_by_rijk = torch.stack(
|
|
[
|
|
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
|
# `int`.
|
|
torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
|
|
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
|
# `int`.
|
|
torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
|
|
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
|
# `int`.
|
|
torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
|
|
# pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
|
|
# `int`.
|
|
torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
|
|
],
|
|
dim=-2,
|
|
)
|
|
|
|
# We floor here at 0.1 but the exact level is not important; if q_abs is small,
|
|
# the candidate won't be picked.
|
|
flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
|
|
quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
|
|
|
|
# if not for numerical problems, quat_candidates[i] should be same (up to a sign),
|
|
# forall i; we pick the best-conditioned one (with the largest denominator)
|
|
out = quat_candidates[F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :].reshape(batch_dim + (4,))
|
|
|
|
# Convert from rijk to ijkr
|
|
out = out[..., [1, 2, 3, 0]]
|
|
|
|
out = standardize_quaternion(out)
|
|
|
|
return out
|
|
|
|
|
|
def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Returns torch.sqrt(torch.max(0, x))
|
|
but with a zero subgradient where x is 0.
|
|
"""
|
|
ret = torch.zeros_like(x)
|
|
positive_mask = x > 0
|
|
if torch.is_grad_enabled():
|
|
ret[positive_mask] = torch.sqrt(x[positive_mask])
|
|
else:
|
|
ret = torch.where(positive_mask, torch.sqrt(x), ret)
|
|
return ret
|
|
|
|
|
|
def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Convert a unit quaternion to a standard form: one in which the real
|
|
part is non negative.
|
|
|
|
Args:
|
|
quaternions: Quaternions with real part last,
|
|
as tensor of shape (..., 4).
|
|
|
|
Returns:
|
|
Standardized quaternions as tensor of shape (..., 4).
|
|
"""
|
|
return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions)
|