def test_points_to_homogeneous(self):
inputs = [
torch.tensor([[[0, 0]]]), # single element, one point, 2d
torch.tensor([[[0, 0, 0]]]), # single element, one point, 3d
torch.tensor([[[3, 3], [4, 4]]]), # single element, multiple points, 2d
torch.tensor([[[3, 3, 3], [4, 4, 4]]]), # single element, multiple points 3d
torch.tensor([[[0, 0]], [[2, 2]]]), # multiple elements, one point, 2d
torch.tensor([[[0, 0, 0]], [[2, 2, 2]]]), # multiple elements, one point, 3d
torch.tensor([[[0, 0], [1, 1]], [[2, 2], [3, 3]]]), # multiple elements, multiple points, 2d
torch.tensor([[[0, 0, 0], [1, 1, 1]], [[2, 2, 2], [3, 3, 3]]]) # multiple elements, multiple points, 2d
]
expectations = [
torch.tensor([[[0, 0, 1]]]), # single element, one point, 2d
torch.tensor([[[0, 0, 0, 1]]]), # single element, one point, 3d
torch.tensor([[[3, 3, 1], [4, 4, 1]]]), # single element, multiple points, 2d
torch.tensor([[[3, 3, 3, 1], [4, 4, 4, 1]]]), # single element, multiple points 3d
torch.tensor([[[0, 0, 1]], [[2, 2, 1]]]), # multiple elements, one point, 2d
torch.tensor([[[0, 0, 0, 1]], [[2, 2, 2, 1]]]), # multiple elements, one point, 3d
torch.tensor([[[0, 0, 1], [1, 1, 1]], [[2, 2, 1], [3, 3, 1]]]), # multiple elements, multiple points, 2d
torch.tensor([[[0, 0, 0, 1], [1, 1, 1, 1]], [[2, 2, 2, 1], [3, 3, 3, 1]]]) # multiple elements,
# multiple points, 2d
]
for inp, exp in zip(inputs, expectations):
with self.subTest(input=inp, expected=exp):
self.assertTrue(torch.allclose(points_to_homogeneous(inp), exp))
def affine_point_transform(point_batch: torch.Tensor, matrix_batch: torch.Tensor) -> torch.Tensor:
"""
Function to perform an affine transformation onto point batches
Args:
point_batch: a point batch of shape [N, NP, NDIM]
``NP`` is the number of points,
``N`` is the batch size,
``NDIM`` is the number of spatial dimensions
matrix_batch : torch.Tensor
a batch of affine matrices with shape [N, NDIM, NDIM + 1],
N is the batch size and NDIM is the number of spatial dimensions
Returns:
torch.Tensor: the batch of transformed points in cartesian coordinates)
[N, NP, NDIM] ``NP`` is the number of points, ``N`` is the
batch size, ``NDIM`` is the number of spatial dimensions
"""
point_batch = points_to_homogeneous(point_batch)
matrix_batch = matrix_to_homogeneous(matrix_batch)
transformed_points = torch.bmm(point_batch, matrix_batch.permute(0, 2, 1))
return points_to_cartesian(transformed_points)