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)