Source code for shepherd_score.alignment.utils.se3_np
"""
Functions used for SE(3) transformations. (NumPy implementation).
Namely, converting quaternions to rotation matrices, getting an SE(3) transform from SE(3)
parameters, and applying the SE(3) transformation on a set of points.
Credit to Lewis J. Martin as this was adapted from
https://github.com/ljmartin/align/blob/main/0.2%20aligning%20principal%20moments%20of%20inertia.ipynb
and PyTorch's implementations.
"""
import numpy as np
[docs]
def quaternions_to_rotation_matrix_np(quaternions: np.ndarray) -> np.ndarray:
"""
Converts quaternion to a rotation matrix.
Adapted from PyTorch3D:
https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#quaternion_to_matrix
Parameters
----------
quaternions : np.ndarray (4,)
Quaternion parameters in (r,i,j,k) order.
set.
Returns
-------
rotation_matrix : np.ndarray (3,3)
Rotation matrix converted from quaternion.
"""
# Single instance
if quaternions.shape == (4,):
r, i, j, k = quaternions
two_s = 2. / (quaternions * quaternions).sum()
rotation_matrix = np.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,
)
else:
raise ValueError(f'Input "quaternions" must be a 1D Tensor of length 4. Instead the shape given was: {quaternions.shape}')
return rotation_matrix.reshape((3, 3))
[docs]
def get_SE3_transform_np(se3_params: np.ndarray
) -> np.ndarray:
"""
Constructs an SE(3) transformtion matrix from parameters. NumPy implementation
Parameters
----------
se3_params : np.ndarray (7,)
Parameters for SE(3) transformation.
The first 4 values in the last dimension are quaternions of form (r,i,j,k)
and the last 3 values of the last dimension are the translations in (x,y,z).
Returns
-------
se3_matrix : np.ndarray (4, 4)
se3_params converted to a 4x4 SE(3) transformation matrix.
"""
if se3_params.shape == (7,):
# Extract quaternion and translation parameters
quaternion_params = se3_params[:4]
translation_params = se3_params[4:]
# Normalize quaternion to ensure unit length
quaternion_params = quaternion_params / np.linalg.norm(quaternion_params)
rotation_matrix = quaternions_to_rotation_matrix_np(quaternion_params)
# Construct SE(3) transformation matrix
se3_matrix = np.eye(4)
se3_matrix[:3, :3] = rotation_matrix
se3_matrix[:3, 3] = translation_params
else:
raise ValueError(f'Input "se3_params" must be a 1D Tensor of length 7. Instead the shape given was: {se3_params.shape}')
return se3_matrix
[docs]
def apply_SE3_transform_np(points: np.ndarray,
SE3_transform: np.ndarray
) -> np.ndarray:
"""
Takes a point cloud and transforms it according to the provided SE3 transformation matrix.
NumPy implementation.
Parameters
----------
points : np.ndarray (N, 3)
Set of coordinates representing a point cloud.
SE3_transform : np.ndarray (4, 4)
SE(3) transformation matrix.
Returns
-------
transformed_points : np.ndarray (N, 3)
Set of coordinates transformed by the corresponding SE(3) transformation.
"""
if points.shape[-1] != 3:
raise ValueError(f'"points" should have shape (N_points, 3). Instead the shape given was: {points.shape}')
if SE3_transform.shape[-2:] != (4,4):
raise ValueError(f'"SE3_transform" should have shape (4, 4). Instead the shape given was: {SE3_transform.shape}')
if len(SE3_transform.shape) != len(points.shape):
raise ValueError(f'Shapes of points and SE3_transform should be the same length. Instead {len(SE3_transform.shape)} and {len(points.shape)} were given.')
# Single instance
if len(SE3_transform.shape) == 2:
transformed_points = (SE3_transform[:3,:3] @ points.T).T + SE3_transform[:3,3]
else:
raise ValueError(f'The expected length of shape for "points" and "SE3_transform" must be 2 but {len(SE3_transform)} was given.')
return transformed_points
[docs]
def apply_SO3_transform_np(points: np.ndarray,
SE3_transform: np.ndarray
) -> np.ndarray:
"""
Takes a point cloud and ONLY ROTATES it according to the provided SE3 transformation matrix.
Supports batched and non-batched inputs.
Parameters
----------
points : np.array (N, 3)
Set of coordinates representing a point cloud.
SE3_transform : (4, 4)
SE(3) transformation matrix. If 'points' argument is batched, this one should be too.
Returns
-------
rotated_points : torch.Tensor (batch, N, 3) or (N, 3)
Set of coordinates rotated by the rotation component of the SE(3) transformation.
"""
if points.shape[-1] != 3:
raise ValueError(f'"points" should have shape (N_points, 3). Instead the shape given was: {points.shape}')
if SE3_transform.shape[-2:] != (4,4):
raise ValueError(f'"SE3_transform" should have shape (4, 4). Instead the shape given was: {SE3_transform.shape}')
if len(SE3_transform.shape) != len(points.shape):
raise ValueError(f'Shapes of points and SE3_transform should be the same length. Instead {len(SE3_transform.shape)} and {len(points.shape)} were given.')
# Single instance
if len(SE3_transform.shape) == 2:
rotated_points = (SE3_transform[:3,:3] @ points.T).T
else:
raise ValueError(f'"points" and "SE3_transform" must be a single instance. \
The expected length of shape for both should be 2 single instance but {len(SE3_transform)} was given.')
return rotated_points