import pytest import torch from torch.autograd import gradcheck import kornia import kornia.testing as utils # test utils from kornia.testing import assert_close def identity_matrix(batch_size, device, dtype): r"""Create a batched homogeneous identity matrix""" return torch.eye(4, device=device, dtype=dtype).repeat(batch_size, 1, 1) # Nx4x4 def euler_angles_to_rotation_matrix(x, y, z): r"""Create a rotation matrix from x, y, z angles""" assert x.dim() == 1, x.shape assert x.shape == y.shape == z.shape ones, zeros = torch.ones_like(x), torch.zeros_like(x) # the rotation matrix for the x-axis rx_tmp = [ ones, zeros, zeros, zeros, zeros, torch.cos(x), -torch.sin(x), zeros, zeros, torch.sin(x), torch.cos(x), zeros, zeros, zeros, zeros, ones, ] rx = torch.stack(rx_tmp, dim=-1).view(-1, 4, 4) # the rotation matrix for the y-axis ry_tmp = [ torch.cos(y), zeros, torch.sin(y), zeros, zeros, ones, zeros, zeros, -torch.sin(y), zeros, torch.cos(y), zeros, zeros, zeros, zeros, ones, ] ry = torch.stack(ry_tmp, dim=-1).view(-1, 4, 4) # the rotation matrix for the z-axis rz_tmp = [ torch.cos(z), -torch.sin(z), zeros, zeros, torch.sin(z), torch.cos(z), zeros, zeros, zeros, zeros, ones, zeros, zeros, zeros, zeros, ones, ] rz = torch.stack(rz_tmp, dim=-1).view(-1, 4, 4) return torch.matmul(rz, torch.matmul(ry, rx)) # Bx4x4 class TestTransformPoints: @pytest.mark.parametrize("batch_size", [1, 2, 5]) @pytest.mark.parametrize("num_points", [2, 3, 5]) @pytest.mark.parametrize("num_dims", [2, 3]) def test_transform_points(self, batch_size, num_points, num_dims, device, dtype): # generate input data eye_size = num_dims + 1 points_src = torch.rand(batch_size, num_points, num_dims, device=device, dtype=dtype) dst_homo_src = utils.create_random_homography(batch_size, eye_size).to(device=device, dtype=dtype) dst_homo_src = dst_homo_src.to(device) # transform the points from dst to ref points_dst = kornia.geometry.linalg.transform_points(dst_homo_src, points_src) # transform the points from ref to dst src_homo_dst = torch.inverse(dst_homo_src) points_dst_to_src = kornia.geometry.linalg.transform_points(src_homo_dst, points_dst) # projected should be equal as initial assert_close(points_src, points_dst_to_src, atol=1e-4, rtol=1e-4) def test_gradcheck(self, device, dtype): # generate input data batch_size, num_points, num_dims = 2, 3, 2 eye_size = num_dims + 1 points_src = torch.rand(batch_size, num_points, num_dims, device=device, dtype=dtype) dst_homo_src = utils.create_random_homography(batch_size, eye_size).to(device=device, dtype=dtype) # evaluate function gradient points_src = utils.tensor_to_gradcheck_var(points_src) # to var dst_homo_src = utils.tensor_to_gradcheck_var(dst_homo_src) # to var assert gradcheck(kornia.geometry.transform_points, (dst_homo_src, points_src), raise_exception=True) def test_jit(self, device, dtype): points = torch.ones(1, 2, 2, device=device, dtype=dtype) transform = kornia.eye_like(3, points) op = kornia.geometry.transform_points op_script = torch.jit.script(op) actual = op_script(transform, points) expected = op(transform, points) assert_close(actual, expected, atol=1e-4, rtol=1e-4) class TestComposeTransforms: def test_smoke(self, device, dtype): batch_size = 2 trans_01 = identity_matrix(batch_size=batch_size, device=device, dtype=dtype) trans_12 = identity_matrix(batch_size=batch_size, device=device, dtype=dtype) to_check_1 = kornia.geometry.compose_transformations(trans_01, trans_12) to_check_2 = kornia.geometry.compose_transformations(trans_01[0], trans_12[0]) assert to_check_1.shape == (batch_size, 4, 4) assert to_check_2.shape == (4, 4) def test_exception(self, device, dtype): to_check_1 = torch.rand((7, 4, 4, 3), device=device, dtype=dtype) to_check_2 = torch.rand((5, 10, 10), device=device, dtype=dtype) to_check_3 = torch.rand((6, 4, 4), device=device, dtype=dtype) to_check_4 = torch.rand((4, 4), device=device, dtype=dtype) to_check_5 = torch.rand((3, 3), device=device, dtype=dtype) # Testing if exception is thrown when both inputs have shape (3, 3) with pytest.raises(ValueError): _ = kornia.geometry.compose_transformations(to_check_5, to_check_5) # Testing if exception is thrown when both inputs have shape (5, 10, 10) with pytest.raises(ValueError): _ = kornia.geometry.compose_transformations(to_check_2, to_check_2) # Testing if exception is thrown when one input has shape (6, 4, 4) # whereas the other input has shape (4, 4) with pytest.raises(ValueError): _ = kornia.geometry.compose_transformations(to_check_3, to_check_4) # Testing if exception is thrown when one input has shape (7, 4, 4, 3) # whereas the other input has shape (4, 4) with pytest.raises(ValueError): _ = kornia.geometry.compose_transformations(to_check_1, to_check_4) def test_translation_4x4(self, device, dtype): offset = 10 trans_01 = identity_matrix(batch_size=1, device=device, dtype=dtype)[0] trans_12 = identity_matrix(batch_size=1, device=device, dtype=dtype)[0] trans_12[..., :3, -1] += offset # add offset to translation vector trans_02 = kornia.geometry.linalg.compose_transformations(trans_01, trans_12) assert_close(trans_02, trans_12, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_translation_Bx4x4(self, batch_size, device, dtype): offset = 10 trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_12 = identity_matrix(batch_size, device=device, dtype=dtype) trans_12[..., :3, -1] += offset # add offset to translation vector trans_02 = kornia.geometry.linalg.compose_transformations(trans_01, trans_12) assert_close(trans_02, trans_12, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_gradcheck(self, batch_size, device, dtype): trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_12 = identity_matrix(batch_size, device=device, dtype=dtype) trans_01 = utils.tensor_to_gradcheck_var(trans_01) # to var trans_12 = utils.tensor_to_gradcheck_var(trans_12) # to var assert gradcheck(kornia.geometry.linalg.compose_transformations, (trans_01, trans_12), raise_exception=True) class TestInverseTransformation: def test_smoke(self, device, dtype): batch_size = 2 trans_01 = identity_matrix(batch_size=batch_size, device=device, dtype=dtype) to_check_1 = kornia.geometry.inverse_transformation(trans_01) to_check_2 = kornia.geometry.inverse_transformation(trans_01[0]) assert to_check_1.shape == (batch_size, 4, 4) assert to_check_2.shape == (4, 4) def test_exception(self, device, dtype): to_check_1 = torch.rand((7, 4, 4, 3), device=device, dtype=dtype) to_check_2 = torch.rand((5, 10, 10), device=device, dtype=dtype) to_check_3 = torch.rand((3, 3), device=device, dtype=dtype) # Testing if exception is thrown when the input has shape (7, 4, 4, 3) with pytest.raises(ValueError): _ = kornia.geometry.inverse_transformation(to_check_1) # Testing if exception is thrown when the input has shape (5, 10, 10) with pytest.raises(ValueError): _ = kornia.geometry.inverse_transformation(to_check_2) # Testing if exception is thrown when the input has shape (3, 3) with pytest.raises(ValueError): _ = kornia.geometry.inverse_transformation(to_check_3) def test_translation_4x4(self, device, dtype): offset = 10 trans_01 = identity_matrix(batch_size=1, device=device, dtype=dtype)[0] trans_01[..., :3, -1] += offset # add offset to translation vector trans_10 = kornia.geometry.linalg.inverse_transformation(trans_01) trans_01_hat = kornia.geometry.linalg.inverse_transformation(trans_10) assert_close(trans_01, trans_01_hat, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_translation_Bx4x4(self, batch_size, device, dtype): offset = 10 trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_01[..., :3, -1] += offset # add offset to translation vector trans_10 = kornia.geometry.linalg.inverse_transformation(trans_01) trans_01_hat = kornia.geometry.linalg.inverse_transformation(trans_10) assert_close(trans_01, trans_01_hat, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_rotation_translation_Bx4x4(self, batch_size, device, dtype): offset = 10 x, y, z = 0, 0, kornia.pi ones = torch.ones(batch_size, device=device, dtype=dtype) rmat_01 = euler_angles_to_rotation_matrix(x * ones, y * ones, z * ones) trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_01[..., :3, -1] += offset # add offset to translation vector trans_01[..., :3, :3] = rmat_01[..., :3, :3] trans_10 = kornia.geometry.linalg.inverse_transformation(trans_01) trans_01_hat = kornia.geometry.linalg.inverse_transformation(trans_10) assert_close(trans_01, trans_01_hat, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_gradcheck(self, batch_size, device, dtype): trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_01 = utils.tensor_to_gradcheck_var(trans_01) # to var assert gradcheck(kornia.geometry.linalg.inverse_transformation, (trans_01,), raise_exception=True) class TestRelativeTransformation: def test_smoke(self, device, dtype): batch_size = 2 trans_01 = identity_matrix(batch_size=batch_size, device=device, dtype=dtype) trans_02 = identity_matrix(batch_size=batch_size, device=device, dtype=dtype) to_check_1 = kornia.geometry.relative_transformation(trans_01, trans_02) to_check_2 = kornia.geometry.relative_transformation(trans_01[0], trans_02[0]) assert to_check_1.shape == (batch_size, 4, 4) assert to_check_2.shape == (4, 4) def test_exception(self, device, dtype): to_check_1 = torch.rand((7, 4, 4, 3), device=device, dtype=dtype) to_check_2 = torch.rand((5, 10, 10), device=device, dtype=dtype) to_check_3 = torch.rand((6, 4, 4), device=device, dtype=dtype) to_check_4 = torch.rand((4, 4), device=device, dtype=dtype) to_check_5 = torch.rand((3, 3), device=device, dtype=dtype) # Testing if exception is thrown when both inputs have shape (3, 3) with pytest.raises(ValueError): _ = kornia.geometry.relative_transformation(to_check_5, to_check_5) # Testing if exception is thrown when both inputs have shape (5, 10, 10) with pytest.raises(ValueError): _ = kornia.geometry.relative_transformation(to_check_2, to_check_2) # Testing if exception is thrown when one input has shape (6, 4, 4) # whereas the other input has shape (4, 4) with pytest.raises(ValueError): _ = kornia.geometry.relative_transformation(to_check_3, to_check_4) # Testing if exception is thrown when one input has shape (7, 4, 4, 3) # whereas the other input has shape (4, 4) with pytest.raises(ValueError): _ = kornia.geometry.relative_transformation(to_check_1, to_check_4) def test_translation_4x4(self, device, dtype): offset = 10.0 trans_01 = identity_matrix(batch_size=1, device=device, dtype=dtype)[0] trans_02 = identity_matrix(batch_size=1, device=device, dtype=dtype)[0] trans_02[..., :3, -1] += offset # add offset to translation vector trans_12 = kornia.geometry.linalg.relative_transformation(trans_01, trans_02) trans_02_hat = kornia.geometry.linalg.compose_transformations(trans_01, trans_12) assert_close(trans_02_hat, trans_02, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_rotation_translation_Bx4x4(self, batch_size, device, dtype): offset = 10.0 x, y, z = 0.0, 0.0, kornia.pi ones = torch.ones(batch_size, device=device, dtype=dtype) rmat_02 = euler_angles_to_rotation_matrix(x * ones, y * ones, z * ones) trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_02 = identity_matrix(batch_size, device=device, dtype=dtype) trans_02[..., :3, -1] += offset # add offset to translation vector trans_02[..., :3, :3] = rmat_02[..., :3, :3] trans_12 = kornia.geometry.linalg.relative_transformation(trans_01, trans_02) trans_02_hat = kornia.geometry.linalg.compose_transformations(trans_01, trans_12) assert_close(trans_02_hat, trans_02, atol=1e-4, rtol=1e-4) @pytest.mark.parametrize("batch_size", [1, 2, 5]) def test_gradcheck(self, batch_size, device, dtype): trans_01 = identity_matrix(batch_size, device=device, dtype=dtype) trans_02 = identity_matrix(batch_size, device=device, dtype=dtype) trans_01 = utils.tensor_to_gradcheck_var(trans_01) # to var trans_02 = utils.tensor_to_gradcheck_var(trans_02) # to var assert gradcheck(kornia.geometry.linalg.relative_transformation, (trans_01, trans_02), raise_exception=True)