Spaces:
Paused
Paused
| # 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. | |
| """ | |
| Code is copied from PyTorch3D library. | |
| Minimal code necessary for `s03_exp_map` function which is used in the forward | |
| kinematics function. | |
| Taken from Version: v0.7.1 | |
| https://github.com/facebookresearch/pytorch3d/releases/tag/v0.7.1 | |
| """ | |
| import torch | |
| def so3_exp_map(log_rot: torch.Tensor, eps: float = 0.0001) -> torch.Tensor: | |
| """ | |
| Convert a batch of logarithmic representations of rotation matrices `log_rot` | |
| to a batch of 3x3 rotation matrices using Rodrigues formula [1]. | |
| In the logarithmic representation, each rotation matrix is represented as | |
| a 3-dimensional vector (`log_rot`) who's l2-norm and direction correspond | |
| to the magnitude of the rotation angle and the axis of rotation respectively. | |
| The conversion has a singularity around `log(R) = 0` | |
| which is handled by clamping controlled with the `eps` argument. | |
| Args: | |
| log_rot: Batch of vectors of shape `(minibatch, 3)`. | |
| eps: A float constant handling the conversion singularity. | |
| Returns: | |
| Batch of rotation matrices of shape `(minibatch, 3, 3)`. | |
| Raises: | |
| ValueError if `log_rot` is of incorrect shape. | |
| [1] https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula | |
| """ | |
| return _so3_exp_map(log_rot, eps=eps)[0] | |
| def _so3_exp_map( | |
| log_rot: torch.Tensor, eps: float = 0.0001 | |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| A helper function that computes the so3 exponential map and, | |
| apart from the rotation matrix, also returns intermediate variables | |
| that can be re-used in other functions. | |
| """ | |
| _, dim = log_rot.shape | |
| if dim != 3: | |
| raise ValueError("Input tensor shape has to be Nx3.") | |
| nrms = (log_rot * log_rot).sum(1) | |
| # phis ... rotation angles | |
| rot_angles = torch.clamp(nrms, eps).sqrt() | |
| rot_angles_inv = 1.0 / rot_angles | |
| fac1 = rot_angles_inv * rot_angles.sin() | |
| fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos()) | |
| skews = hat(log_rot) | |
| skews_square = torch.bmm(skews, skews) | |
| R = ( | |
| # pyre-fixme[16]: `float` has no attribute `__getitem__`. | |
| fac1[:, None, None] * skews | |
| + fac2[:, None, None] * skews_square | |
| + torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None] | |
| ) | |
| return R, rot_angles, skews, skews_square | |
| def hat(v: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Compute the Hat operator [1] of a batch of 3D vectors. | |
| Args: | |
| v: Batch of vectors of shape `(minibatch , 3)`. | |
| Returns: | |
| Batch of skew-symmetric matrices of shape | |
| `(minibatch, 3 , 3)` where each matrix is of the form: | |
| `[ 0 -v_z v_y ] | |
| [ v_z 0 -v_x ] | |
| [ -v_y v_x 0 ]` | |
| Raises: | |
| ValueError if `v` is of incorrect shape. | |
| [1] https://en.wikipedia.org/wiki/Hat_operator | |
| """ | |
| N, dim = v.shape | |
| if dim != 3: | |
| raise ValueError("Input vectors have to be 3-dimensional.") | |
| h = torch.zeros((N, 3, 3), dtype=v.dtype, device=v.device) | |
| x, y, z = v.unbind(1) | |
| h[:, 0, 1] = -z | |
| h[:, 0, 2] = y | |
| h[:, 1, 0] = z | |
| h[:, 1, 2] = -x | |
| h[:, 2, 0] = -y | |
| h[:, 2, 1] = x | |
| return h | |