fastdev.xform¶

Submodules¶

Package Contents¶

fastdev.xform.axis_angle_vector_to_matrix(axis_angle: torch.Tensor) → torch.Tensor[source]¶

Convert rotations given as axis/angle to rotation matrices.

Parameters:: axis_angle (torch.Tensor) – Rotations given as a vector in axis angle form, as a tensor of shape (…, 3), where the magnitude is the angle turned anticlockwise in radians around the vector’s direction.
Returns:: Rotation matrices as tensor of shape (…, 3, 3).
Return type:: torch.Tensor

fastdev.xform.axis_angle_vector_to_quaternion(axis_angle)[source]¶

Convert rotations given as axis/angle to quaternions.

Parameters:: axis_angle – Rotations given as a vector in axis angle form, as a tensor of shape (…, 3), where the magnitude is the angle turned anticlockwise in radians around the vector’s direction.
Returns:: quaternions with real part first, as tensor of shape (…, 4).

Reference: https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Unit_quaternions

fastdev.xform.compose_axis_angle_vector(axis, angle)[source]¶

fastdev.xform.matrix_to_axis_angle_vector(matrix: torch.Tensor) → torch.Tensor[source]¶

Convert rotations given as rotation matrices to axis/angle.

Parameters:

matrix (torch.Tensor) – Rotation matrices as tensor of shape (…, 3, 3).

Returns:

Rotations given as a vector in axis angle form, as a tensor: of shape (…, 3), where the magnitude is the angle turned anticlockwise in radians around the vector’s direction.

Return type:

torch.Tensor

fastdev.xform.matrix_to_euler_angles(matrix: torch.Tensor, convention: str = 'xyz') → torch.Tensor[source]¶

Convert rotations given as rotation matrices to Euler angles in radians.

Parameters:

matrix (torch.Tensor) – Rotation matrices with shape (…, 3, 3).
convention (str) – Convention string of 3/4 letters, e.g. “xyz”, “sxyz”, “rxyz”, “exyz”. If the length is 3, the extrinsic rotation is assumed. If the length is 4, the first character is “r/i” (rotating/intrinsic), or “s/e” (static / extrinsic). The remaining characters are the axis “x, y, z” in the order.

Returns:

Euler angles in radians with shape (…, 3).

Return type:

torch.Tensor

fastdev.xform.matrix_to_rotation_6d(matrix: torch.Tensor) → torch.Tensor[source]¶

Converts rotation matrices to 6D rotation representation by Zhou et al. [1] by dropping the last row. Note that 6D representation is not unique.

Parameters:: matrix (torch.Tensor) – batch of rotation matrices of size […, 3, 3]
Returns:: 6D rotation representation, of shape […, 6]
Return type:: torch.Tensor

[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. CVPR 2019. arxiv

fastdev.xform.quaternion_invert(quaternion: torch.Tensor) → torch.Tensor[source]¶

Given a quaternion representing rotation, get the quaternion representing its inverse.

Parameters:: quaternion (torch.Tensor) – Quaternions as tensor of shape (…, 4), with real part first, which must be versors (unit quaternions).
Returns:: The inverse, a tensor of quaternions of shape (…, 4).
Return type:: torch.Tensor

fastdev.xform.quaternion_multiply(a: torch.Tensor, b: torch.Tensor) → torch.Tensor[source]¶

Multiply two quaternions representing rotations, returning the quaternion representing their composition, i.e. the versor with nonnegative real part. Usual torch rules for broadcasting apply.

Parameters:

a (torch.Tensor) – Quaternions as tensor of shape (…, 4), real part first.
b (torch.Tensor) – Quaternions as tensor of shape (…, 4), real part first.

Returns:

The product of a and b, a tensor of quaternions of shape (…, 4).

Return type:

torch.Tensor

fastdev.xform.quaternion_real_to_first(quaternions)[source]¶

fastdev.xform.quaternion_real_to_last(quaternions)[source]¶

fastdev.xform.quaternion_to_axis_angle_vector(quaternions)[source]¶

Convert rotations given as quaternions to axis/angle.

Parameters:

quaternions – quaternions with real part first, as tensor of shape (…, 4).

Returns:

Rotations given as a vector in axis angle form, as a tensor: of shape (…, 3), where the magnitude is the angle turned anticlockwise in radians around the vector’s direction.

Reference: https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation#Unit_quaternions

fastdev.xform.quaternion_to_matrix(quaternions: torch.Tensor) → torch.Tensor[source]¶

Convert rotations given as quaternions to rotation matrices.

Parameters:: quaternions (Tensor) – quaternions with real part first with shape (…, 4).
Returns:: Rotation matrices as tensor of shape (…, 3, 3).
Return type:: Tensor

Reference: https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation

fastdev.xform.random_rotation_matrix(num: int | None = None, random_state: int | numpy.random.Generator | numpy.random.RandomState | None = None, return_tensors: Literal['np', 'pt'] = 'np')[source]¶

Parameters:

num (Optional[int])
random_state (Optional[Union[int, numpy.random.Generator, numpy.random.RandomState]])
return_tensors (Literal['np', 'pt'])

fastdev.xform.rotation_6d_to_matrix(d6: torch.Tensor) → torch.Tensor[source]¶

Converts 6D rotation representation by Zhou et al. [1] to rotation matrix using Gram–Schmidt orthogonalization per Section B of [1].

Parameters:: d6 (Tensor) – 6D rotation representation of shape […, 6]
Returns:: Rotation matrices of shape […, 3, 3]
Return type:: Tensor

[1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H. On the Continuity of Rotation Representations in Neural Networks. CVPR 2019. arxiv

pytorch3d implementation

fastdev.xform.split_axis_angle_vector(axis_angle)[source]¶

fastdev.xform.standardize_quaternion(quaternions: torch.Tensor) → torch.Tensor[source]¶

Convert a unit quaternion to a standard form: one in which the real part is non negative.

Parameters:: quaternions (torch.Tensor) – Quaternions with real part first, as tensor of shape (…, 4).
Returns:: Standardized quaternions as tensor of shape (…, 4).
Return type:: torch.Tensor

fastdev.xform.expand_tf_mat(tf_mat: jaxtyping.Float[torch.Tensor, ... 3 4]) → jaxtyping.Float[torch.Tensor, ... 4 4][source]¶

Expand transformation matrix of shape [… 3 4] to shape [… 4 4].

Parameters:: tf_mat (torch.Tensor) – Transformation matrix in shape [… 3 4] or [… 4 4].
Returns:: Expanded transformation matrix in shape [… 4 4].
Return type:: torch.Tensor

Examples

>>> tf_mat = torch.tensor([[0, 1, 0, 1], [0, 0, 1, 2], [1, 0, 0, 3]], dtype=torch.float32)
>>> expand_tf_mat(tf_mat)
tensor([[0., 1., 0., 1.],
        [0., 0., 1., 2.],
        [1., 0., 0., 3.],
        [0., 0., 0., 1.]])

fastdev.xform.inverse_tf_mat(rot_or_tf_mat: torch.Tensor) → torch.Tensor[source]¶

Inverse a rotation matrix or a transformation matrix. Reference_

Parameters:: rot_or_tf_mat (torch.Tensor) – Rotation matrix (in shape […, 3, 3]) or transformation matrix (in shape […, 4, 4]).
Returns:: Inversed matrix.
Return type:: torch.Tensor

Examples

>>> tf_mat = torch.tensor([[0, 1, 0, 1], [0, 0, 1, 2], [1, 0, 0, 3], [0, 0, 0, 1]], dtype=torch.float32)
>>> torch.allclose(inverse_tf_mat(tf_mat) @ tf_mat, torch.eye(4))
True
>>> rot_mat = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=torch.float32)
>>> torch.allclose(inverse_tf_mat(rot_mat) @ rot_mat, torch.eye(3))
True

fastdev.xform.project_points(pts: torch.Tensor, intr_mat: torch.Tensor, return_depth: Literal[False] = False) → torch.Tensor[source]¶

fastdev.xform.project_points(pts: torch.Tensor, intr_mat: torch.Tensor, return_depth: Literal[True]) → Tuple[torch.Tensor, torch.Tensor]

Project 3D points in the camera space to the image plane.

Parameters:

pts – 3D points, could be Nx3 or BxNx3.
intr_mat – Intrinsic matrix, could be 3x3 or Bx3x3.

Returns:

the order is uv other than xy. depth (if return_depth): depth in the camera space.

Return type:

pixels

fastdev.xform.rot_tl_to_tf_mat(rot_mat: Optional[jaxtyping.Float[torch.Tensor, ... 3 3]] = None, tl: Optional[jaxtyping.Float[torch.Tensor, ... 3]] = None) → jaxtyping.Float[torch.Tensor, ... 4 4][source]¶

Build transformation matrix with rotation matrix and translation vector.

Parameters:

rot_mat (torch.Tensor, optional) – Rotation matrix in shape [… 3 3]. Defaults to None.
tl (torch.Tensor, optional) – Translation vector in shape [… 3]. Defaults to None.

Returns:

Transformation matrix in shape [… 4 4].

Return type:

torch.Tensor

Examples

>>> rot_mat = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=torch.float32)
>>> tl = torch.tensor([1, 2, 3], dtype=torch.float32)
>>> rot_tl_to_tf_mat(rot_mat, tl)
tensor([[0., 1., 0., 1.],
        [0., 0., 1., 2.],
        [1., 0., 0., 3.],
        [0., 0., 0., 1.]])
>>> rot_tl_to_tf_mat(tl=tl)
tensor([[1., 0., 0., 1.],
        [0., 1., 0., 2.],
        [0., 0., 1., 3.],
        [0., 0., 0., 1.]])
>>> rot_tl_to_tf_mat(rot_mat=rot_mat)
tensor([[0., 1., 0., 0.],
        [0., 0., 1., 0.],
        [1., 0., 0., 0.],
        [0., 0., 0., 1.]])

fastdev.xform.swap_major(rot_or_tf_mat: torch.Tensor) → torch.Tensor[source]¶

Swap the major of a rotation matrix or a transformation matrix. Reference_

Parameters:: rot_or_tf_mat (torch.Tensor) – Rotation or transformation matrix in shape […, 3, 3] or […, 4, 4].
Returns:: Matrix with swapped major.
Return type:: torch.Tensor

fastdev.xform.to_homo(pts_3d: jaxtyping.Float[torch.Tensor, ... 3]) → jaxtyping.Float[torch.Tensor, ... 4][source]¶

Convert Cartesian 3D points to Homogeneous 4D points.

Parameters:: pts_3d (torch.Tensor) – Cartesian 3D points in shape [… 3].
Returns:: Homogeneous 4D points in shape [… 4].
Return type:: torch.Tensor

Examples

>>> pts = torch.tensor([[0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=torch.float32)
>>> to_homo(pts)
tensor([[0., 1., 0., 1.],
        [0., 0., 1., 1.],
        [1., 0., 0., 1.]])

fastdev.xform.unproject_points(pixels, depth, intr_mat)[source]¶

Unproject pixels in the image plane to 3D points in the camera space.

Parameters:

pixels – Pixels in the image plane, could be Nx2 or BxNx2. The order is uv rather than xy.
depth – Depth in the camera space, could be N, Nx1, BxN or BxNx1.
intr_mat – Intrinsic matrix, could be 3x3 or Bx3x3.

Returns:

3D points, Nx3 or BxNx3.

Return type:

pts

fastdev.xform.camera_position_from_spherical_angles(distance: float, elevation: float, azimuth: float, degrees: bool = True, device: Device = 'cpu') → torch.Tensor[source]¶

Calculate the location of the camera based on the distance away from the target point, the elevation and azimuth angles.

Parameters:

distance (float) – distance of the camera from the object.
elevation (float) –
angles. The inputs distance, elevation and azimuth can be one of the following
- Python scalar
- Torch scalar
- Torch tensor of shape (N) or (1)
azimuth (float) –
angles. The inputs distance, elevation and azimuth can be one of the following
- Python scalar
- Torch scalar
- Torch tensor of shape (N) or (1)
degrees (bool) – bool, whether the angles are specified in degrees or radians.
device (Device) – str or torch.device, device for new tensors to be placed on.

Return type:

torch.Tensor

The vectors are broadcast against each other so they all have shape (N, 1).

Returns:

(N, 3) xyz location of the camera.

Return type:

camera_position

Parameters:

distance (float)
elevation (float)
azimuth (float)
degrees (bool)
device (Device)

fastdev.xform.compose_intr_mat(fu: float, fv: float, cu: float, cv: float, skew: float = 0.0) → numpy.ndarray[source]¶

Parameters:

fu (float) – horizontal focal length (width)
fv (float) – vertical focal length (height)
cu (float) – horizontal principal point (width)
cv (float) – vertical principal point (height)
skew (float) – skew coefficient, default to 0

Return type:

numpy.ndarray

fastdev.xform.coord_conversion(src_spec: str, dst_spec: str, check_handness: bool = True, return_tensors: Literal['pt'] = 'pt') → jaxtyping.Float[torch.Tensor, 3 3][source]¶

fastdev.xform.coord_conversion(src_spec: str, dst_spec: str, check_handness: bool = True, return_tensors: Literal['np'] = 'np') → jaxtyping.Float[numpy.ndarray, 3 3]

Construct a rotation matrix based on given source and destination coordinate specifications.

Parameters:

src_spec – Source coordinate specification, e.g., “x: right, y: down, z: front” or “opencv”.
dst_spec – Destination coordinate specification, e.g., “x: right, y: up, z: back” or “opengl”.
check_handness – If True, checks if the rotation matrix preserves right-handedness.
return_tensors – Return type of the rotation matrix, either “np” for NumPy array or “pt” for PyTorch tensor.

Returns:

A 3x3 rotation matrix converting coordinates from the source to the destination specification.

Examples

>>> coord_conversion("opencv", "opengl")
array([[ 1.,  0.,  0.],
       [ 0., -1.,  0.],
       [ 0.,  0., -1.]], dtype=float32)
>>> coord_conversion("x: front, y: left, z: up", "x: left, y: up, z: front")
array([[0., 1., 0.],
       [0., 0., 1.],
       [1., 0., 0.]], dtype=float32)
>>> coord_conversion("x: right, y: down, z: front", "x: left, y: up, z: front")
array([[-1.,  0.,  0.],
       [ 0., -1.,  0.],
       [ 0.,  0.,  1.]], dtype=float32)
>>> coord_conversion("x: left, y: up, z: front", "x: front, y: left, z: up", return_tensors="pt")
tensor([[0., 0., 1.],
        [1., 0., 0.],
        [0., 1., 0.]])

fastdev.xform.look_at_rotation(camera_position: jaxtyping.Float[torch.Tensor, *batch 3], at: jaxtyping.Float[torch.Tensor, *batch 3], up: jaxtyping.Float[torch.Tensor, *batch 3]) → jaxtyping.Float[torch.Tensor, *batch 3 3][source]¶

This function takes a vector camera_position which specifies the location of the camera in world coordinates and two vectors at and up which indicate the position of the object and the up directions of the world coordinate system respectively.

The output is a rotation matrix representing the rotation from camera coordinates to world coordinates.

We use the OpenGL coordinate in this function, i.e. x -> right, y -> up, z -> backward. Hence, z_axis: pos - at, x_axis: cross(up, z_axis), y_axis: cross(z_axis, x_axis)

Note that our implementation differs from pytorch3d.

our matrix is in the OpenGL coordinate
our matrix is column-major
our matrix is the camera-to-world transformation

Parameters:

camera_position (jaxtyping.Float[torch.Tensor, *batch 3]) – position of the camera in world coordinates
at (jaxtyping.Float[torch.Tensor, *batch 3]) – position of the object in world coordinates
up (jaxtyping.Float[torch.Tensor, *batch 3]) – vector specifying the up direction in the world coordinate frame.

Returns:

rotation matrices of shape […, 3, 3]

Return type:

fastdev.xform.axis_angle_to_matrix(axis: jaxtyping.Float[torch.Tensor, ... 3], angle: jaxtyping.Float[torch.Tensor, ...]) → jaxtyping.Float[torch.Tensor, ... 3 3][source]¶

Converts axis angles to rotation matrices using Rodrigues formula.

Parameters:

axis (torch.Tensor) – axis, the shape could be […, 3].
angle (torch.Tensor) – angle, the shape could be […].

Returns:

Rotation matrices […, 3, 3].

Return type:

torch.Tensor

Example

>>> axis = torch.tensor([1.0, 0.0, 0.0])
>>> angle = torch.tensor(0.5)
>>> axis_angle_to_matrix(axis, angle)
tensor([[ 1.0000,  0.0000,  0.0000],
        [ 0.0000,  0.8776, -0.4794],
        [ 0.0000,  0.4794,  0.8776]])

fastdev.xform.euler_angles_to_matrix(euler_angles: jaxtyping.Float[torch.Tensor, ... 3], axes: _AXES = 'sxyz') → jaxtyping.Float[torch.Tensor, ... 3 3][source]¶

Converts Euler angles to rotation matrices.

Parameters:

euler_angles (torch.Tensor) – Tensor of Euler angles with shape […, 3].
axes (str) – Axis specification string, one of 24 possible sequences (e.g., “sxyz”). If only 3 characters are provided, “s” will be prefixed.

Returns:

Rotation matrices with shape […, 3, 3].

Return type:

torch.Tensor

Example

>>> euler_angles = torch.tensor([1.0, 0.5, 2.0])
>>> euler_angles_to_matrix(euler_angles, axes="sxyz")
tensor([[-0.3652, -0.6592,  0.6574],
        [ 0.7980,  0.1420,  0.5857],
        [-0.4794,  0.7385,  0.4742]])
>>> euler_angles_to_matrix(euler_angles, axes="rxyz")
tensor([[-0.3652, -0.7980,  0.4794],
        [ 0.3234, -0.5917, -0.7385],
        [ 0.8729, -0.1146,  0.4742]])

fastdev.xform.matrix_to_quaternion(rot_mat: jaxtyping.Float[torch.Tensor, ... 3 3], scalar_first: bool = True, canonical: bool = True) → jaxtyping.Float[torch.Tensor, ... 4][source]¶

Converts rotation matrices to quaternions.

Parameters:

rot_mat (torch.Tensor) – Rotation matrices with shape […, 3, 3].
scalar_first (bool)
canonical (bool)

Returns:

Quaternions with shape […, 4].

Return type:

torch.Tensor

Example

>>> rot_mat = torch.tensor([[-0.2533, -0.6075,  0.7529],
...                         [ 0.8445, -0.5185, -0.1343],
...                         [ 0.4720,  0.6017,  0.6443]])
>>> matrix_to_quaternion(rot_mat)
tensor([0.4671, 0.3940, 0.1503, 0.7772])

Note

The gradient of this function differs from the pytorch3d implementation, but it should be okay for most use cases. Ref

fastdev.xform.rotate_points(pts: jaxtyping.Float[torch.Tensor, ... n 3], tf_mat: jaxtyping.Float[torch.Tensor, ... 3 3]) → jaxtyping.Float[torch.Tensor, ... n 3][source]¶

Apply a rotation matrix on a set of 3D points.

Parameters:

pts (torch.Tensor) – 3D points in shape [… n 3].
rot_mat (torch.Tensor) – Rotation matrix in shape [… 3 3].
tf_mat (jaxtyping.Float[torch.Tensor, ... 3 3])

Returns:

Rotated points in shape [… n 3].

Return type:

torch.Tensor

fastdev.xform.transform_points(pts: jaxtyping.Float[torch.Tensor, ... n 3], tf_mat: jaxtyping.Float[torch.Tensor, ... 4 4]) → jaxtyping.Float[torch.Tensor, ... n 3][source]¶

Apply a transformation matrix on a set of 3D points.

Parameters:

pts (torch.Tensor) – 3D points, could be [… n 3]
tf_mat (torch.Tensor) – Transformation matrix, could be [… 4 4]

Returns:

Transformed pts in shape of [… n 3]

Return type:

jaxtyping.Float[torch.Tensor, … n 3]

Examples

>>> pts = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
>>> tf_mat = torch.tensor([[0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 2.0], [1.0, 0.0, 0.0, 3.0], [0.0, 0.0, 0.0, 1.0]])
>>> transform_points(pts, tf_mat)
tensor([[3., 5., 4.],
        [6., 8., 7.]])

Note

The dimension number of pts and tf_mat should be the same. The batch dimensions (…) are broadcasted (and thus must be broadcastable). We don’t adopt the shapes [… 3] and [… 4 4] because there is no real broadcasted vector-matrix multiplication in pytorch. [… 3] and [… 4 4] will be converted to [… 1 3] and [… 4 4] and apply a broadcasted matrix-matrix multiplication.