# LICENSE HEADER MANAGED BY add-license-header # # Copyright 2018 Kornia Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # import pytest import torch from kornia.geometry.conversions import euler_from_quaternion from kornia.geometry.liegroup import So3 from kornia.geometry.quaternion import Quaternion from kornia.geometry.vector import Vector3 from testing.base import BaseTester class TestSo3(BaseTester): def _make_rand_data(self, device, dtype, batch_size, dims): shape = [] if batch_size is None else [batch_size] return torch.rand([*shape, dims], device=device, dtype=dtype) def test_smoke(self, device, dtype): q = Quaternion.from_coeffs(1.0, 0.0, 0.0, 0.0) q = q.to(device, dtype) s = So3(q) assert isinstance(s, So3) self.assert_close(s.q.data, q.data) # TODO: implement me def test_cardinality(self, device, dtype): pass # TODO: implement me def test_exception(self, device, dtype): pass # TODO: implement me def test_gradcheck(self, device): pass # TODO: implement me def test_jit(self, device, dtype): pass # TODO: implement me def test_module(self, device, dtype): pass @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_init(self, device, dtype, batch_size): q = Quaternion.identity(batch_size, device, dtype) s1 = So3(q) s2 = So3(s1.q) assert isinstance(s2, So3) self.assert_close(s1.q.data, s2.q.data) @pytest.mark.parametrize("batch_size", (1, 2, 5)) def test_getitem(self, device, dtype, batch_size): q = Quaternion.random(batch_size, device, dtype) s = So3(q) for i in range(batch_size): s1 = s[i] self.assert_close(s1.q.data, q.data[i]) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_mul(self, device, dtype, batch_size): q1 = Quaternion.identity(batch_size, device, dtype) q2 = Quaternion.random(batch_size, device, dtype) t = self._make_rand_data(device, dtype, batch_size, dims=3) s1 = So3(q1) s2 = So3(q2) self.assert_close((s1 * s2).q.data, s2.q.data) self.assert_close((s2 * s2.inverse()).q.data, s1.q.data) self.assert_close((s1 * t), t) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_mul_vec(self, device, dtype, batch_size): q1 = Quaternion.identity(batch_size, device, dtype) q2 = Quaternion.random(batch_size, device, dtype) if batch_size is None: shape = () else: shape = (batch_size,) t = Vector3.random(shape, device, dtype) s1 = So3(q1) s2 = So3(q2) self.assert_close((s1 * s2).q.data, s2.q.data) self.assert_close((s2 * s2.inverse()).q.data, s1.q.data) self.assert_close((s1 * t), t) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_unit_norm(self, device, dtype, batch_size): q1 = Quaternion.random(batch_size, device, dtype) q2 = Quaternion.random(batch_size, device, dtype) s1 = So3(q1) s2 = So3(q2) s3 = s1 * s2 s4 = s1.inverse() s5 = s2.inverse() s6 = s3.inverse() ones_vec = torch.tensor(1.0, device=device, dtype=dtype) if batch_size is None: self.assert_close(s1.q.norm(), ones_vec) return for i in range(batch_size): self.assert_close(s1[i].q.norm(), ones_vec) self.assert_close(s2[i].q.norm(), ones_vec) self.assert_close(s3[i].q.norm(), ones_vec) self.assert_close(s4[i].q.norm(), ones_vec) self.assert_close(s5[i].q.norm(), ones_vec) self.assert_close(s6[i].q.norm(), ones_vec) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_exp(self, device, dtype, batch_size): q = Quaternion.identity(batch_size, device, dtype) s = So3(q) zero_vec = 0 * self._make_rand_data(device, dtype, batch_size, dims=3) self.assert_close(s.exp(zero_vec).q.data, q.data) # exp of zero vec is identity @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_log(self, device, dtype, batch_size): q = Quaternion.identity(batch_size, device, dtype) s = So3(q) zero_vec = 0 * self._make_rand_data(device, dtype, batch_size, dims=3) self.assert_close(s.log(), zero_vec) # log of identity quat is zero vec @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_exp_log(self, device, dtype, batch_size): q = Quaternion.random(batch_size, device, dtype) s = So3(q) a = self._make_rand_data(device, dtype, batch_size, dims=3) b = s.exp(a).log() self.assert_close(b, a) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_hat(self, device, dtype, batch_size): v = torch.tensor([1, 2, 3], device=device, dtype=dtype) expected = v if batch_size is not None: v = v.repeat(batch_size, 1) hat = So3.hat(v) if batch_size is None: hat = hat[None] self.assert_close(hat.unique()[-3:], expected) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_vee(self, device, dtype, batch_size): omega = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device, dtype=dtype) expected = torch.tensor([8, 3, 4], device=device, dtype=dtype) if batch_size is not None: omega = omega.repeat(batch_size, 1, 1) expected = expected.repeat(batch_size, 1) self.assert_close(So3.vee(omega), expected) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_hat_vee(self, device, dtype, batch_size): a = self._make_rand_data(device, dtype, batch_size, dims=3) omega = So3.hat(a) b = So3.vee(omega) self.assert_close(b, a) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_matrix(self, device, dtype, batch_size): q = Quaternion.random(batch_size, device, dtype) r = So3(q).matrix() if batch_size is None: q = Quaternion(q.data[None]) r = r[None] for i in range(r.shape[0]): q1 = q[i] r1 = r[i, :, :] pvec = torch.rand(3, device=device, dtype=dtype) pquat = Quaternion(torch.cat([torch.tensor([0], device=device, dtype=dtype), pvec])) qp_ = q1 * pquat * q1.inv() rp_ = torch.matmul(r1, pvec) self.assert_close(rp_, qp_.vec) # p_ = R*p = q*p*q_inv self.assert_close(rp_.norm(), pvec.norm()) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_from_wxyz(self, device, dtype, batch_size): wxyz = self._make_rand_data(device, dtype, batch_size, dims=4) s = So3.from_wxyz(wxyz) self.assert_close(s.q.data, wxyz) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_ortho(self, device, dtype, batch_size): q = Quaternion.random(batch_size, device, dtype) b_R_a = So3(q).matrix() a_R_b = So3(q).inverse().matrix() a_R_a = (So3(q) * So3(q).inverse()).matrix() eye_mat = torch.eye(3, device=device, dtype=dtype) if batch_size is None: eye_mat = eye_mat[None] a_R_a = a_R_a[None] a_R_b = a_R_b[None] b_R_a = b_R_a[None] if batch_size is not None: eye_mat = eye_mat.repeat(batch_size, 1, 1) self.assert_close(a_R_a, eye_mat) for i in range(eye_mat.shape[0]): self.assert_close(a_R_a[i, :, :], eye_mat[i]) self.assert_close(a_R_b[i, :, :] @ b_R_a[i, :, :], eye_mat[i]) self.assert_close(b_R_a[i, :, :] @ a_R_b[i, :, :], eye_mat[i]) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_inverse(self, device, dtype, batch_size): q = Quaternion.random(batch_size, device, dtype) self.assert_close(So3(q).inverse().inverse().q.data, q.data) self.assert_close(So3(q).inverse().inverse().matrix(), So3(q).matrix()) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_rot_x(self, device, dtype, batch_size): x = self._make_rand_data(device, dtype, batch_size, dims=1).squeeze(-1) so3 = So3.rot_x(x) roll, _, _ = euler_from_quaternion(*so3.q.coeffs) self.assert_close(x, roll) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_rot_y(self, device, dtype, batch_size): y = self._make_rand_data(device, dtype, batch_size, dims=1).squeeze(-1) so3 = So3.rot_y(y) _, pitch, _ = euler_from_quaternion(*so3.q.coeffs) self.assert_close(y, pitch) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_rot_z(self, device, dtype, batch_size): z = self._make_rand_data(device, dtype, batch_size, dims=1).squeeze(-1) so3 = So3.rot_z(z) _, _, yaw = euler_from_quaternion(*so3.q.coeffs) self.assert_close(z, yaw) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_adjoint(self, device, dtype, batch_size): q1 = Quaternion.random(batch_size, device, dtype) q2 = Quaternion.random(batch_size, device, dtype) x = So3(q1) y = So3(q2) self.assert_close(x.inverse().adjoint(), x.adjoint().inverse()) self.assert_close((x * y).adjoint(), x.adjoint() @ y.adjoint()) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_random(self, device, dtype, batch_size): s = So3.random(batch_size=batch_size, device=device, dtype=dtype) s_in_s = s.inverse() * s i = So3.identity(batch_size=batch_size, device=device, dtype=dtype) self.assert_close(s_in_s.q.data, i.q.data) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_right_jacobian(self, device, dtype, batch_size): vec = self._make_rand_data(device, dtype, batch_size, dims=3) Jr = So3.right_jacobian(vec) I = torch.eye(3, device=device, dtype=dtype).expand_as(Jr) # noqa: E741 self.assert_close(vec[..., None], Jr @ vec[..., None]) self.assert_close(Jr.transpose(-1, -2) @ Jr, I, atol=0.1, rtol=0.1) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_left_jacobian(self, device, dtype, batch_size): vec = self._make_rand_data(device, dtype, batch_size, dims=3) Jl = So3.left_jacobian(vec) I = torch.eye(3, device=device, dtype=dtype).expand_as(Jl) # noqa: E741 self.assert_close(vec[..., None], Jl @ vec[..., None]) self.assert_close(Jl.transpose(-1, -2) @ Jl, I, atol=0.1, rtol=0.1) @pytest.mark.parametrize("batch_size", (None, 1, 2, 5)) def test_right_left_jacobian(self, device, dtype, batch_size): vec = self._make_rand_data(device, dtype, batch_size, dims=3) Jr = So3.right_jacobian(vec) Jl = So3.left_jacobian(vec) self.assert_close(Jl, Jr.transpose(-1, -2))