diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 54fc6c2d5..c9d3f494e 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -489,26 +489,34 @@ def rotation_matrix_between(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Flo """ a = a / torch.linalg.norm(a) b = b / torch.linalg.norm(b) - v = torch.linalg.cross(a, b) # Axis of rotation. + v = torch.linalg.cross(a, b) + cosine = torch.dot(a, b).clamp(-1, 1) + identity = torch.eye(3, dtype=a.dtype, device=a.device) + zero = a.new_zeros(()) - # Handle cases where `a` and `b` are parallel. - eps = 1e-6 - if torch.sum(torch.abs(v)) < eps: - x = torch.tensor([1.0, 0, 0]) if abs(a[0]) < eps else torch.tensor([0, 1.0, 0]) - v = torch.linalg.cross(a, x) + def skew(vector): + x, y, z = vector.unbind() + return torch.stack((zero, -z, y, z, zero, -x, -y, x, zero)).reshape(3, 3) - v = v / torch.linalg.norm(v) - skew_sym_mat = torch.Tensor( - [ - [0, -v[2], v[1]], - [v[2], 0, -v[0]], - [-v[1], v[0], 0], - ] - ) - theta = torch.acos(torch.clip(torch.dot(a, b), -1, 1)) + if cosine >= 0: + # This form retains small cross products even when the dot product + # rounds to one, and its denominator is bounded below by one. + skew_sym_mat = skew(v) + return identity + skew_sym_mat + (skew_sym_mat @ skew_sym_mat) / (1 + cosine) - # Rodrigues rotation formula. https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula - return torch.eye(3) + torch.sin(theta) * skew_sym_mat + (1 - torch.cos(theta)) * (skew_sym_mat @ skew_sym_mat) + sine = torch.linalg.norm(v) + if sine > 0: + # atan2 remains well-conditioned as the angle approaches pi. + theta = torch.atan2(sine, cosine) + skew_sym_mat = skew(v / sine) + return identity + torch.sin(theta) * skew_sym_mat + (1 - torch.cos(theta)) * (skew_sym_mat @ skew_sym_mat) + + # Antiparallel vectors have no unique rotation axis. Pick a well-conditioned + # perpendicular axis and use the exact half-turn formula. + basis = identity[torch.argmin(a.abs())] + axis = torch.linalg.cross(a, basis) + axis = axis / torch.linalg.norm(axis) + return 2 * axis[:, None] * axis[None, :] - identity def focus_of_attention(poses: Float[Tensor, "*num_poses 4 4"], initial_focus: Float[Tensor, "3"]) -> Float[Tensor, "3"]: diff --git a/tests/cameras/test_rotation_matrix_between.py b/tests/cameras/test_rotation_matrix_between.py new file mode 100644 index 000000000..ee7189f69 --- /dev/null +++ b/tests/cameras/test_rotation_matrix_between.py @@ -0,0 +1,82 @@ +# Licensed under the Apache License, Version 2.0. +"""SO(3) invariants for rotations between camera direction vectors.""" + +import importlib.util +from pathlib import Path + +import pytest +import torch + +path = Path(__file__).resolve().parents[2] / "nerfstudio" / "cameras" / "camera_utils.py" +spec = importlib.util.spec_from_file_location("nerfstudio_camera_utils", path) +camera_utils = importlib.util.module_from_spec(spec) +spec.loader.exec_module(camera_utils) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("angle", [0.0, 1e-8, 5e-7, 1e-5, 0.7, 2.3, torch.pi - 1e-5, torch.pi]) +def test_rotation_maps_directions_and_preserves_so3(dtype, angle): + a = torch.tensor([1.0, 0.0, 0.0], dtype=dtype) + if angle == torch.pi: + b = -a + else: + theta = torch.tensor(angle, dtype=dtype) + b = torch.stack((theta.cos(), theta.sin(), theta.new_zeros(()))) + before_a, before_b = a.clone(), b.clone() + rotation = camera_utils.rotation_matrix_between(a, b) + assert rotation.dtype == dtype + assert rotation.device == a.device + # A residual threshold alone could let every tiny-angle case pass. + # Their relative vector error must also remain below 1 percent. + tolerance = max(2e-12, abs(float(b[1])) * 0.01) + torch.testing.assert_close(rotation @ a, b, rtol=1e-6, atol=tolerance) + identity = torch.eye(3, dtype=dtype) + torch.testing.assert_close(rotation.T @ rotation, identity) + torch.testing.assert_close(torch.linalg.det(rotation), torch.tensor(1.0, dtype=dtype)) + # Minimal-angle rotations satisfy trace(R) = 1 + 2 * dot(a, b). + torch.testing.assert_close(rotation.trace(), 1 + 2 * torch.dot(a, b)) + torch.testing.assert_close(a, before_a, rtol=0, atol=0) + torch.testing.assert_close(b, before_b, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_nonaxis_parallel_and_antiparallel_vectors(dtype): + a = torch.tensor([1.0, 2.0, -3.0], dtype=dtype) + for factor in [-1.0, 1.0]: + b = factor * a + rotation = camera_utils.rotation_matrix_between(a, b) + assert rotation.dtype == dtype + torch.testing.assert_close(rotation @ a, b) + torch.testing.assert_close(rotation.T @ rotation, torch.eye(3, dtype=dtype)) + torch.testing.assert_close(torch.linalg.det(rotation), torch.tensor(1.0, dtype=dtype)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_rotation_is_invariant_to_direction_magnitudes(dtype): + a = torch.tensor([1.0, 2.0, 3.0], dtype=dtype) + b = torch.tensor([-2.0, 1.0, 4.0], dtype=dtype) + reference = camera_utils.rotation_matrix_between(a, b) + scaled = camera_utils.rotation_matrix_between(a * 1000.0, b * 0.001) + assert scaled.dtype == dtype + torch.testing.assert_close(scaled, reference) + + +@pytest.mark.parametrize("offset", [0.0, 1e-5, 0.5]) +def test_parallel_and_nearparallel_rotation_has_finite_correct_jacobian(offset): + a = torch.tensor([1.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True) + b = torch.tensor([1.0, offset, 0.0], dtype=torch.float64, requires_grad=True) + rotation = camera_utils.rotation_matrix_between(a, b) + assert torch.isfinite(torch.autograd.grad(rotation.sum(), (a, b), retain_graph=True)[0]).all() + assert torch.autograd.gradcheck(camera_utils.rotation_matrix_between, (a, b), eps=1e-6, atol=2e-5, rtol=1e-3) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_direction_mapping_has_normalization_jacobian(dtype): + a = torch.tensor([1.0, 0.0, 0.0], dtype=dtype) + b = torch.tensor([0.8, 0.6, 0.0], dtype=dtype, requires_grad=True) + jacobian = torch.autograd.functional.jacobian(lambda value: camera_utils.rotation_matrix_between(a, value) @ a, b) + unit_b = b / b.norm() + # Since R(a, b) @ unit(a) = unit(b), its Jacobian is independently + # determined by normalization, including perturbations out of the plane. + expected = (torch.eye(3, dtype=dtype) - unit_b[:, None] * unit_b[None, :]) / b.norm() + torch.testing.assert_close(jacobian, expected)