def test_batch(self, batch_size, device, dtype): B: int = batch_size center = torch.rand(B, 3, device=device, dtype=dtype) angle = torch.rand(B, 3, device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angle, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle, scales) assert P.shape == (B, 3, 4)
def get_affine_matrix3d(translations: torch.Tensor, center: torch.Tensor, scale: torch.Tensor, angles: torch.Tensor, sxy: Optional[torch.Tensor] = None, sxz: Optional[torch.Tensor] = None, syx: Optional[torch.Tensor] = None, syz: Optional[torch.Tensor] = None, szx: Optional[torch.Tensor] = None, szy: Optional[torch.Tensor] = None) -> torch.Tensor: r"""Composes 3d affine matrix from the components. Args: translations (torch.Tensor): tensor containing the translation vector (dx,dy,dz) with shape :math:`(B, 3)`. center (torch.Tensor): tensor containing the center vector (x,y,z) with shape :math:`(B, 3)`. scale (torch.Tensor): tensor containing the scale factor with shape :math:`(B)`. angle: (torch.Tensor): angle axis vector containing the rotation angles in degrees in the form of (rx, ry, rz) with shape :math:`(B, 3)`. Internally it calls Rodrigues to compute the rotation matrix from axis-angle. sxy (torch.Tensor, optional): tensor containing the shear factor in the xy-direction with shape :math:`(B)`. sxz (torch.Tensor, optional): tensor containing the shear factor in the xz-direction with shape :math:`(B)`. syx (torch.Tensor, optional): tensor containing the shear factor in the yx-direction with shape :math:`(B)`. syz (torch.Tensor, optional): tensor containing the shear factor in the yz-direction with shape :math:`(B)`. szx (torch.Tensor, optional): tensor containing the shear factor in the zx-direction with shape :math:`(B)`. szy (torch.Tensor, optional): tensor containing the shear factor in the zy-direction with shape :math:`(B)`. Returns: torch.Tensor: the 3d affine transformation matrix :math:`(B, 4, 4)`. """ transform: torch.Tensor = get_projective_transform(center, -angles, scale) transform[..., 3] += translations # tx/ty/tz # pad transform to get Bx3x3 transform_h = convert_affinematrix_to_homography3d(transform) if any([s is not None for s in [sxy, sxz, syx, syz, szx, szy]]): shear_mat = get_shear_matrix3d(center, sxy, sxz, syx, syz, szx, szy) transform_h = transform_h @ shear_mat return transform_h
def test_rotate_y(self, device, dtype): input = torch.tensor( [[[ [[0.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.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.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [2.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.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.0, 90.0, 0.0]], device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angles, device=device, dtype=dtype) P = proj.get_projective_transform(center, angles, scales) output = proj.warp_affine3d(input, P, (3, 3, 3)) assert_close(output, expected, rtol=1e-4, atol=1e-4)
def test_smoke(self, device, dtype): center = torch.rand(1, 3, device=device, dtype=dtype) angle = torch.rand(1, 3, device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angle, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle, scales) assert P.shape == (1, 3, 4)
def test_rot90z(self, device, dtype): center = torch.zeros(1, 3, device=device, dtype=dtype) angle = torch.tensor([[0., 0., 90.]], device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angle, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle, scales) P_expected = torch.tensor([ [0., -1., 0., 0.], [1., 0., 0., 0.], [0., 0., 1., 0.], ], device=device, dtype=dtype).unsqueeze(0) assert_allclose(P, P_expected, atol=1e-4, rtol=1e-4)
def test_rot90y(self, device, dtype): center = torch.zeros(1, 3, device=device, dtype=dtype) angle = torch.tensor([[0., 90., 0.]], device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angle[:, 0]) P = proj.get_projective_transform(center, angle, scales) P_expected = torch.tensor([ [0., 0., 1., 0.], [0., 1., 0., 0.], [-1., 0., 0., 0.], ], device=device, dtype=dtype) assert_allclose(P, P_expected)
def test_rot90z(self, device, dtype): center = torch.zeros(1, 3, device=device, dtype=dtype) angle = torch.tensor([[0., 0., 90.]], device=device, dtype=dtype) P = proj.get_projective_transform(center, angle) P_expected = torch.tensor([ [0., -1., 0., 0.], [1., 0., 0., 0.], [0., 0., 1., 0.], ], device=device, dtype=dtype) assert_allclose(P, P_expected)
def test_identity(self, device, dtype): center = torch.zeros(1, 3, device=device, dtype=dtype) angle = torch.zeros(1, 3, device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angle, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle, scales) P_expected = torch.tensor( [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]], device=device, dtype=dtype).unsqueeze(0) assert_close(P, P_expected, atol=1e-4, rtol=1e-4)
def test_rotate_y_large(self, device, dtype): """Rotates 90deg anti-clockwise.""" input = torch.tensor( [[ [ [[0.0, 4.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 2.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [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, 0.0, 0.0], [0.0, 6.0, 7.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 8.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.0, 0.0, 0.0]], [[4.0, 2.0, 0.0], [3.0, 1.0, 0.0], [0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.0, 0.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.0]], [[0.0, 0.0, 0.0], [0.0, 6.0, 8.0], [9.0, 0.0, 0.0]], [[0.0, 0.0, 0.0], [0.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.0, 90.0, 0.0]], device=device, dtype=dtype) scales: torch.Tensor = torch.ones_like(angles, device=device, dtype=dtype) P = proj.get_projective_transform(center, angles, scales) output = proj.warp_affine3d(input, P, (3, 3, 3)) assert_close(output, expected, rtol=1e-4, atol=1e-4)
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, device=device, dtype=dtype) P = proj.get_projective_transform(center, angles, scales) output = proj.warp_affine3d(input, P, (3, 3, 3)) assert_allclose(output, expected)
def _compute_rotation_matrix3d(yaw: torch.Tensor, pitch: torch.Tensor, roll: torch.Tensor, center: torch.Tensor) -> torch.Tensor: """Computes a pure affine rotation matrix.""" if len(yaw.shape) == len(pitch.shape) == len(roll.shape) == 0: yaw = yaw.unsqueeze(dim=0) pitch = pitch.unsqueeze(dim=0) roll = roll.unsqueeze(dim=0) if len(yaw.shape) == len(pitch.shape) == len(roll.shape) == 1: yaw = yaw.unsqueeze(dim=1) pitch = pitch.unsqueeze(dim=1) roll = roll.unsqueeze(dim=1) assert len(yaw.shape) == len(pitch.shape) == len(roll.shape) == 2, \ f"Expected yaw, pitch, roll to be (B, 1). Got {yaw.shape}, {pitch.shape}, {roll.shape}." angles: torch.Tensor = torch.cat([yaw, pitch, roll], dim=1) scales: torch.Tensor = torch.ones_like(yaw) matrix: torch.Tensor = get_projective_transform(center, angles, scales) return matrix
def get_affine_matrix3d(translations: torch.Tensor, center: torch.Tensor, scale: torch.Tensor, angles: torch.Tensor, sxy: Optional[torch.Tensor] = None, sxz: Optional[torch.Tensor] = None, syx: Optional[torch.Tensor] = None, syz: Optional[torch.Tensor] = None, szx: Optional[torch.Tensor] = None, szy: Optional[torch.Tensor] = None) -> torch.Tensor: r"""Composes affine matrix Bx4x4 from the components Returns: torch.Tensor: params to be passed to the affine transformation. """ transform: torch.Tensor = get_projective_transform(center, -angles, scale) transform[..., 3] += translations # tx/ty/tz # pad transform to get Bx3x3 transform_h = convert_affinematrix_to_homography3d(transform) shear_mat = get_shear_matrix3d(center, sxy, sxz, syx, syz, szx, szy) transform_h = transform_h @ shear_mat return transform_h
def test_batch(self, batch_size, device, dtype): B: int = batch_size center = torch.rand(B, 3, device=device, dtype=dtype) angle = torch.rand(B, 3, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle) assert P.shape == (B, 3, 4)
def test_smoke(self, device, dtype): center = torch.rand(1, 3, device=device, dtype=dtype) angle = torch.rand(1, 3, device=device, dtype=dtype) P = proj.get_projective_transform(center, angle) assert P.shape == (1, 3, 4)
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)