Spaces:
Running on Zero
Running on Zero
File size: 3,870 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """Geometric transformation functions."""
import torch
from torch import Tensor
def transform_points(points: Tensor, transform: Tensor) -> Tensor:
"""Applies transform to points.
Args:
points (Tensor): points of shape (N, D) or (B, N, D).
transform (Tensor): transforms of shape (D+1, D+1) or (B, D+1, D+1).
Returns:
Tensor: (N, D) / (B, N, D) transformed points.
Raises:
ValueError: Either points or transform have incorrect shape
"""
hom_coords = torch.cat([points, torch.ones_like(points[..., 0:1])], -1)
if len(points.shape) == 2:
if len(transform.shape) == 3:
assert (
transform.shape[0] == 1
), "Got multiple transforms for single point set!"
transform = transform.squeeze(0)
transform = transform.T
elif len(points.shape) == 3:
if len(transform.shape) == 2:
transform = transform.T.unsqueeze(0)
elif len(transform.shape) == 3:
transform = transform.permute(0, 2, 1)
else:
raise ValueError(f"Shape of transform invalid: {transform.shape}")
else:
raise ValueError(f"Shape of input points invalid: {points.shape}")
points_transformed = hom_coords @ transform
return points_transformed[..., : points.shape[-1]]
def inverse_pinhole(intrinsic_matrix: Tensor) -> Tensor:
"""Calculate inverse of pinhole projection matrix.
Args:
intrinsic_matrix (Tensor): [..., 3, 3] intrinsics or single [3, 3]
intrinsics.
Returns:
Tensor: Inverse of input intrinisics.
"""
squeeze = False
inv = intrinsic_matrix.clone()
if len(intrinsic_matrix.shape) == 2:
inv = inv.unsqueeze(0)
squeeze = True
inv[..., 0, 0] = 1.0 / inv[..., 0, 0]
inv[..., 1, 1] = 1.0 / inv[..., 1, 1]
inv[..., 0, 2] = -inv[..., 0, 2] * inv[..., 0, 0]
inv[..., 1, 2] = -inv[..., 1, 2] * inv[..., 1, 1]
if squeeze:
inv = inv.squeeze(0)
return inv
def inverse_rigid_transform(transformation: Tensor) -> Tensor:
"""Calculate inverse of rigid body transformation(s).
Args:
transformation (Tensor): [N, 4, 4] transformations or single [4, 4]
transformation.
Returns:
Tensor: Inverse of input transformation(s).
"""
squeeze = False
if len(transformation.shape) == 2:
transformation = transformation.unsqueeze(0)
squeeze = True
rotation, translation = transformation[:, :3, :3], transformation[:, :3, 3]
rot = rotation.permute(0, 2, 1)
t = -rot @ translation[:, :, None]
inv = torch.cat([torch.cat([rot, t], -1), transformation[:, 3:4]], 1)
if squeeze:
inv = inv.squeeze(0)
return inv
def get_transform_matrix(rotation: Tensor, translation: Tensor) -> Tensor:
"""Assembles 4x4 transformation from rotation / translation pair(s).
Args:
rotation (Tensor): [N, 3, 3] or [3, 3] rotation(s).
translation (Tensor): [N, 3] or [3,] translation(s).
Returns:
Tensor: [N, 4, 4] or [4, 4] transformation.
"""
squeeze = False
if len(rotation.shape) == 2:
assert len(translation.shape) == 1
rotation = rotation.unsqueeze(0)
translation = translation.unsqueeze(0)
squeeze = True
batch_size = 1
else:
assert len(rotation.shape) == 3 and len(translation.shape) == 2
assert rotation.shape[0] == translation.shape[0]
batch_size = rotation.shape[0]
assert (
rotation.shape[-2] == rotation.shape[-1] == translation.shape[-1] == 3
)
transforms = rotation.new_zeros((batch_size, 4, 4))
transforms[:, :3, :3] = rotation
transforms[:, :3, 3] = translation
transforms[:, 3, 3] = 1.0
if squeeze:
transforms = transforms.squeeze(0)
return transforms
|