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)
