def test_forth_back(self, device, dtype): out_shape = (3, 4, 5) input = torch.rand(2, 5, 3, 4, 5, device=device, dtype=dtype) P = torch.rand(2, 3, 4, device=device, dtype=dtype) P = proj.matrix_to_homogeneous(P) P_hat = (P.inverse() @ P)[:, :3] output = proj.warp_projective(input, P_hat, out_shape) assert_allclose(output, input)
def test_forth_back(self, device, dtype): out_shape = (3, 4, 5) input = torch.rand(2, 5, 3, 4, 5, device=device, dtype=dtype) P = torch.rand(2, 3, 4, device=device, dtype=dtype) P = kornia.geometry.convert_affinematrix_to_homography3d(P) P_hat = (P.inverse() @ P)[:, :3] output = proj.warp_projective(input, P_hat, out_shape) assert_allclose(output, input, rtol=1e-4, atol=1e-4)
def apply_affine3d(input: torch.Tensor, params: Dict[str, torch.Tensor], flags: Dict[str, torch.Tensor]) -> torch.Tensor: r"""Random affine transformation of the image keeping center invariant. Args: input (torch.Tensor): Tensor to be transformed with shape (D, H, W), (C, D, H, W), (B, C, D, H, W). params (Dict[str, torch.Tensor]): - params['angles']: Degrees of rotation with the shape of :math: `(*, 3)` for yaw, pitch, roll. - params['translations']: Horizontal, vertical and depthical translations (dx,dy,dz). - params['center']: Rotation center (x,y,z). - params['scale']: Isotropic scaling params. - params['sxy']: Shear param toward x-y-axis. - params['sxz']: Shear param toward x-z-axis. - params['syx']: Shear param toward y-x-axis. - params['syz']: Shear param toward y-z-axis. - params['szx']: Shear param toward z-x-axis. - params['szy']: Shear param toward z-y-axis. flags (Dict[str, torch.Tensor]): - params['resample']: Integer tensor. NEAREST = 0, BILINEAR = 1. - params['align_corners']: Boolean tensor. Returns: torch.Tensor: The transfromed input """ if not torch.is_tensor(input): raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") input = _transform_input3d(input) _validate_input_dtype( input, accepted_dtypes=[torch.float16, torch.float32, torch.float64]) # arrange input data x_data: torch.Tensor = input.view(-1, *input.shape[-4:]) depth, height, width = x_data.shape[-3:] # concatenate transforms transform: torch.Tensor = compute_affine_transformation3d(input, params) resample_name: str = Resample(flags['resample'].item()).name.lower() align_corners: bool = cast(bool, flags['align_corners'].item()) out_data: torch.Tensor = warp_projective(x_data, transform[:, :3, :], (depth, height, width), resample_name, align_corners=align_corners) return out_data.view_as(input)
def test_rotate_x(self, device, dtype): input = torch.tensor([[[[ [0., 0., 0.], [0., 2., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 1., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ]]]], device=device, dtype=dtype) expected = torch.tensor([[[[ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 1., 0.], [0., 2., 0.], ], [ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ]]]], device=device, dtype=dtype) _, _, D, H, W = input.shape center = torch.tensor([[(W - 1) / 2, (H - 1) / 2, (D - 1) / 2]], device=device, dtype=dtype) angles = torch.tensor([[90., 0., 0.]], device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angles) P = proj.get_projective_transform(center, angles, scales) output = proj.warp_projective(input, P, (3, 3, 3)) assert_allclose(output, expected)
def affine3d(tensor: torch.Tensor, matrix: torch.Tensor, mode: str = 'bilinear', align_corners: bool = False) -> torch.Tensor: r"""Apply an affine transformation to the 3d volume. Args: tensor (torch.Tensor): The image tensor to be warped in shapes of :math:`(D, H, W)`, :math:`(C, D, H, W)` and :math:`(B, C, D, H, W)`. matrix (torch.Tensor): The 3x4 affine transformation matrix. mode (str): 'bilinear' | 'nearest' align_corners(bool): interpolation flag. Default: False. See https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.interpolate for detail Returns: torch.Tensor: The warped image. """ # warping needs data in the shape of BCDHW is_unbatched: bool = tensor.ndimension() == 4 if is_unbatched: tensor = torch.unsqueeze(tensor, dim=0) # we enforce broadcasting since by default grid_sample it does not # give support for that matrix = matrix.expand(tensor.shape[0], -1, -1) # warp the input tensor depth: int = tensor.shape[-3] height: int = tensor.shape[-2] width: int = tensor.shape[-1] warped: torch.Tensor = warp_projective(tensor, matrix, (depth, height, width), mode, align_corners=align_corners) # return in the original shape if is_unbatched: warped = torch.squeeze(warped, dim=0) return warped
def test_batch(self, batch_size, num_channels, out_shape, device, dtype): B, C = batch_size, num_channels input = torch.rand(B, C, 3, 4, 5, device=device, dtype=dtype) P = torch.rand(B, 3, 4, device=device, dtype=dtype) output = proj.warp_projective(input, P, out_shape) assert list(output.shape) == [B, C] + list(out_shape)
def test_rotate_y_large(self, device, dtype): """Rotates 90deg anti-clockwise.""" input = torch.tensor([[[[ [0., 4., 0.], [0., 3., 0.], [0., 0., 0.], ], [ [0., 2., 0.], [0., 1., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ]], [[ [0., 0., 0.], [0., 0., 0.], [0., 9., 0.], ], [ [0., 0., 0.], [0., 6., 7.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 8., 0.], [0., 0., 0.], ]]]], device=device, dtype=dtype) expected = torch.tensor( [[[[ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ], [ [4., 2., 0.], [3., 1., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ]], [[ [0., 0., 0.], [0., 7., 0.], [0., 0., 0.], ], [ [0., 0., 0.], [0., 6., 8.], [9., 0., 0.], ], [ [0., 0., 0.], [0., 0., 0.], [0., 0., 0.], ]]]], device=device, dtype=dtype) _, _, D, H, W = input.shape center = torch.tensor([[(W - 1) / 2, (H - 1) / 2, (D - 1) / 2]], device=device, dtype=dtype) angles = torch.tensor([[0., 90., 0.]], device=device, dtype=dtype) P = proj.get_projective_transform(center, angles) output = proj.warp_projective(input, P, (3, 3, 3)) assert_allclose(output, expected)
def test_smoke(self, device, dtype): input = torch.rand(1, 3, 3, 4, 5, device=device, dtype=dtype) P = torch.rand(1, 3, 4, device=device, dtype=dtype) output = proj.warp_projective(input, P, (3, 4, 5)) assert output.shape == (1, 3, 3, 4, 5)