def test_stacked_transformation_assembly(self): first_matrix = torch.tensor([[[2.0, 0.0, 1.0], [0.0, 3.0, 2.0]]]) second_matrix = torch.tensor([[[4.0, 0.0, 3.0], [0.0, 5.0, 4.0]]]) trafo = StackedAffine([first_matrix, second_matrix]) sample = {"data": torch.rand(1, 3, 25, 25)} matrix = trafo.assemble_matrix(**sample) target_matrix = matrix_to_cartesian( torch.bmm(matrix_to_homogeneous(first_matrix), matrix_to_homogeneous(second_matrix))) self.assertTrue(torch.allclose(matrix, target_matrix))
def create_rotation(rotation: AffineParamType, batchsize: int, ndim: int, degree: bool = False, device: Optional[Union[torch.device, str]] = None, dtype: Optional[Union[torch.dtype, str]] = None, image_transform: bool = True) -> torch.Tensor: """ Formats the given scale parameters to a homogeneous transformation matrix Args: rotation: the rotation factor(s). Supported are: * a single parameter (as float or int), which will be replicated for all dimensions and batch samples * a parameter per sample, which will be replicated for all dimensions * a parameter per dimension, which will be replicated for all batch samples * a parameter per sampler per dimension * None will be treated as a rotation angle of 0 batchsize: the number of samples per batch ndim : the dimensionality of the transform degree: whether the given rotation(s) are in degrees. Only valid for rotation parameters, which aren't passed as full transformation matrix. device: the device to put the resulting tensor to. Defaults to the torch default device dtype: the dtype of the resulting trensor. Defaults to the torch default dtype image_transform: bool inverts the rotation matrix to match expected behavior when applied to an image, e.g. rotation > 0 should rotate the image counter clockwise but the grid clockwise Returns: torch.Tensor: the homogeneous transformation matrix [N, NDIM + 1, NDIM + 1], N is the batch size and NDIM is the number of spatial dimensions """ if rotation is None: rotation = 0 num_rot_params = 1 if ndim == 2 else ndim rotation = expand_scalar_param(rotation, batchsize, num_rot_params).to(device=device, dtype=dtype) if degree: rotation = deg_to_rad(rotation) matrix_fn = create_rotation_2d if ndim == 2 else create_rotation_3d sin, cos = torch.sin(rotation), torch.cos(rotation) rotation_matrix = torch.stack([matrix_fn(s, c) for s, c in zip(sin, cos)]) homo_rotation_matrix = matrix_to_homogeneous(rotation_matrix) if image_transform: homo_rotation_matrix = orthogonal_inverse(homo_rotation_matrix) return homo_rotation_matrix
def assemble_matrix(self, **data) -> torch.Tensor: """ Assembles the matrix (and takes care of batching and having it on the right device and in the correct dtype and dimensionality). Args: **data: the data to be transformed. Will be used to determine batchsize, dimensionality, dtype and device Returns: torch.Tensor: the (batched) transformation matrix """ if self.matrix is None: raise ValueError("Matrix needs to be initialized or overwritten.") if not torch.is_tensor(self.matrix): self.matrix = torch.tensor(self.matrix) self.matrix = self.matrix.to(data[self.keys[0]]) batchsize = data[self.keys[0]].shape[0] ndim = len(data[self.keys[0]].shape) - 2 # channel and batch dim # batch dimension missing -> Replicate for each sample in batch if len(self.matrix.shape) == 2: self.matrix = self.matrix[None].expand(batchsize, -1, -1).clone() if self.matrix.shape == (batchsize, ndim, ndim + 1): return self.matrix elif self.matrix.shape == (batchsize, ndim, ndim): return matrix_to_homogeneous(self.matrix)[:, :-1] elif self.matrix.shape == (batchsize, ndim + 1, ndim + 1): return matrix_to_cartesian(self.matrix) raise ValueError( "Invalid Shape for affine transformation matrix. " "Got %s but expected %s" % (str(tuple(self.matrix.shape)), str((batchsize, ndim, ndim + 1))) )
def test_matrix_to_homogeneous(self): inputs = [ torch.tensor([[[1, 2], [3, 4]]]), # single sample, no trans, 2d torch.tensor([[[1, 2, 5], [3, 4, 6]]]), # single sample, trans, 2d torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]), # single sample, no trans, 3d torch.tensor([[[1, 2, 3, 10], [4, 5, 6, 11], [7, 8, 9, 12]]]), # single sample, trans, 3d torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]), # multiple samples, no trans, 2d torch.tensor([[[1, 2, 10], [3, 4, 11]], [[5, 6, 12], [7, 8, 13]]]), # multiple samples, trans, 2d # multiple samples, trans, 3d torch.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]], [[10, 11, 12], [13, 14, 15], [16, 17, 18]]]), # multiple samples, trans, 3d torch.tensor([[[1, 2, 3, 21], [4, 5, 6, 22], [7, 8, 9, 23]], [[10, 11, 12, 24], [13, 14, 15, 25], [16, 17, 18, 26]]]) ] expectations = [ torch.tensor([[[1, 2, 0], [3, 4, 0], [0, 0, 1]]]), torch.tensor([[[1, 2, 5], [3, 4, 6], [0, 0, 1]]]), torch.tensor([[[1, 2, 3, 0], [4, 5, 6, 0], [7, 8, 9, 0], [0, 0, 0, 1]]]), torch.tensor([[[1, 2, 3, 10], [4, 5, 6, 11], [7, 8, 9, 12], [0, 0, 0, 1]]]), torch.tensor([[[1, 2, 0], [3, 4, 0], [0, 0, 1]], [[5, 6, 0], [7, 8, 0], [0, 0, 1]]]), torch.tensor([[[1, 2, 10], [3, 4, 11], [0, 0, 1]], [[5, 6, 12], [7, 8, 13], [0, 0, 1]]]), torch.tensor([[[1, 2, 3, 0], [4, 5, 6, 0], [7, 8, 9, 0], [0, 0, 0, 1]], [[10, 11, 12, 0], [13, 14, 15, 0], [16, 17, 18, 0], [0, 0, 0, 1]]]), torch.tensor([[[1, 2, 3, 21], [4, 5, 6, 22], [7, 8, 9, 23], [0, 0, 0, 1]], [[10, 11, 12, 24], [13, 14, 15, 25], [16, 17, 18, 26], [0, 0, 0, 1]]]) ] for inp, exp in zip(inputs, expectations): with self.subTest(input=inp, expected=exp): self.assertTrue(torch.allclose(matrix_to_homogeneous(inp), exp))
def create_translation(offset: AffineParamType, batchsize: int, ndim: int, device: Optional[Union[torch.device, str]] = None, dtype: Optional[Union[torch.dtype, str]] = None, image_transform: bool = True) -> torch.Tensor: """ Formats the given translation parameters to a homogeneous transformation matrix Args: offset: the translation offset(s). Supported are: * a single parameter (as float or int), which will be replicated for all dimensions and batch samples * a parameter per sample, which will be replicated for all dimensions * a parameter per dimension, which will be replicated for all batch samples * a parameter per sampler per dimension * None will be treated as a translation offset of 0 batchsize: the number of samples per batch ndim: the dimensionality of the transform device: the device to put the resulting tensor to. Defaults to the torch default device dtype: the dtype of the resulting trensor. Defaults to the torch default dtype image_transform: bool inverts the translation matrix to match expected behavior when applied to an image, e.g. translation > 0 should move the image in the positive direction of an axis but the grid in the negative direction Returns: torch.Tensor: the homogeneous transformation matrix [N, NDIM + 1, NDIM + 1], N is the batch size and NDIM is the number of spatial dimensions """ if offset is None: offset = 0 offset = expand_scalar_param(offset, batchsize, ndim).to(device=device, dtype=dtype) eye_batch = get_batched_eye(batchsize=batchsize, ndim=ndim, device=device, dtype=dtype) translation_matrix = torch.stack([ torch.cat([eye, o.view(-1, 1)], dim=1) for eye, o in zip(eye_batch, offset) ]) if image_transform: translation_matrix[..., -1] = -translation_matrix[..., -1] return matrix_to_homogeneous(translation_matrix)
def create_scale(scale: AffineParamType, batchsize: int, ndim: int, device: Optional[Union[torch.device, str]] = None, dtype: Optional[Union[torch.dtype, str]] = None, image_transform: bool = True) -> torch.Tensor: """ Formats the given scale parameters to a homogeneous transformation matrix Args: scale : the scale factor(s). Supported are: * a single parameter (as float or int), which will be replicated for all dimensions and batch samples * a parameter per sample, which will be replicated for all dimensions * a parameter per dimension, which will be replicated for all batch samples * a parameter per sampler per dimension * None will be treated as a scaling factor of 1 batchsize: the number of samples per batch ndim: the dimensionality of the transform device: the device to put the resulting tensor to. Defaults to the torch default device dtype: the dtype of the resulting trensor. Defaults to the torch default dtype image_transform: inverts the scale matrix to match expected behavior when applied to an image, e.g. scale>1 increases the size of an image but decrease the size of an grid Returns: torch.Tensor: the homogeneous transformation matrix [N, NDIM + 1, NDIM + 1], N is the batch size and NDIM is the number of spatial dimensions """ if scale is None: scale = 1 scale = expand_scalar_param(scale, batchsize, ndim).to(device=device, dtype=dtype) if image_transform: scale = 1 / scale scale_matrix = torch.stack([ eye * s for eye, s in zip( get_batched_eye( batchsize=batchsize, ndim=ndim, device=device, dtype=dtype), scale) ]) return matrix_to_homogeneous(scale_matrix)
def assemble_matrix(self, **data) -> torch.Tensor: """ Handles the matrix assembly and stacking Args: **data: the data to be transformed. Will be used to determine batchsize, dimensionality, dtype and device Returns: torch.Tensor: the (batched) transformation matrix """ whole_trafo = None for trafo in self.transforms: matrix = matrix_to_homogeneous(trafo.assemble_matrix(**data)) if whole_trafo is None: whole_trafo = matrix else: whole_trafo = torch.bmm(whole_trafo, matrix) return matrix_to_cartesian(whole_trafo)
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)
def test_check_image_size(self): images = [ torch.rand(11, 2, 3, 4, 5), torch.rand(11, 2, 3, 4), torch.rand(11, 2, 3, 3), ] img_sizes = [[3, 4, 5], [3, 4], 3] scales = [ torch.tensor([2.0, 3.0, 4.0]), torch.tensor([2.0, 3.0]), torch.tensor([2.0, 3.0]) ] rots = [[45.0, 90.0, 135.0], [45.0], [45.0]] trans = [[0.0, 10.0, 20.0], [10.0, 20.0], [10.0, 20.0]] edges = [ [ [0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 5.0, 1.0], [0.0, 4.0, 0.0, 1.0], [0.0, 4.0, 5.0, 1.0], [3.0, 0.0, 0.0, 1.0], [3.0, 0.0, 5.0, 1.0], [3.0, 4.0, 0.0, 1.0], [3.0, 4.0, 5.0, 1.0], ], [[0.0, 0.0, 1.0], [0.0, 4.0, 1.0], [3.0, 0.0, 1.0], [3.0, 4.0, 1.0]], [[0.0, 0.0, 1.0], [0.0, 3.0, 1.0], [3.0, 0.0, 1.0], [3.0, 3.0, 1.0]], ] for img, size, scale, rot, tran, edge_pts in zip( images, img_sizes, scales, rots, trans, edges): ndim = scale.size(-1) with self.subTest(ndim=ndim): affine = matrix_to_homogeneous( parametrize_matrix( scale=scale, rotation=rot, translation=tran, degree=True, batchsize=1, ndim=ndim, dtype=torch.float, image_transform=False, )) edge_pts = torch.tensor(edge_pts, dtype=torch.float) img = img.to(torch.float) new_edges = torch.bmm(edge_pts.unsqueeze(0), affine.clone().permute(0, 2, 1)) img_size_zero_border = new_edges.max(dim=1)[0][0] img_size_non_zero_border = (new_edges.max(dim=1)[0] - new_edges.min(dim=1)[0])[0] fn_result_zero_border = _check_new_img_size( size, matrix_to_cartesian( affine.expand(img.size(0), -1, -1).clone()), zero_border=True, ) fn_result_non_zero_border = _check_new_img_size( size, matrix_to_cartesian( affine.expand(img.size(0), -1, -1).clone()), zero_border=False, ) self.assertTrue( torch.allclose(img_size_zero_border[:-1], fn_result_zero_border)) self.assertTrue( torch.allclose(img_size_non_zero_border[:-1], fn_result_non_zero_border))