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_affine3d(input, P_hat, out_shape, flags='nearest')
assert_close(output, input, rtol=1e-4, atol=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_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_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_affine3d(input, P, out_shape)
assert list(output.shape) == [B, C] + list(out_shape)
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_affine3d(input, P, (3, 4, 5))
assert output.shape == (1, 3, 3, 4, 5)