diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index 7a80b9acc2..979eb44d80 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -156,11 +156,18 @@ def intersect_aabb( t_min, t_max - two tensors of shapes N representing distance of intersection from the origin. """ - tx_min = (aabb[:3] - origins) / directions - tx_max = (aabb[3:] - origins) / directions + parallel = directions == 0 + safe_directions = torch.where(parallel, torch.ones_like(directions), directions) + tx_min = (aabb[:3] - origins) / safe_directions + tx_max = (aabb[3:] - origins) / safe_directions t_min = torch.stack((tx_min, tx_max)).amin(dim=0) t_max = torch.stack((tx_min, tx_max)).amax(dim=0) + # A parallel ray imposes no parameter bound if it lies inside the slab. + # Handle outside parallel slabs as misses without ever dividing by zero. + t_min = t_min.masked_fill(parallel, -torch.inf) + t_max = t_max.masked_fill(parallel, torch.inf) + parallel_miss = (parallel & ((origins < aabb[:3]) | (origins > aabb[3:]))).any(dim=-1) t_min = t_min.amax(dim=-1) t_max = t_max.amin(dim=-1) @@ -168,7 +175,7 @@ def intersect_aabb( t_min = torch.clamp(t_min, min=0, max=max_bound) t_max = torch.clamp(t_max, min=0, max=max_bound) - cond = t_max <= t_min + cond = (t_max <= t_min) | parallel_miss t_min = torch.where(cond, invalid_value, t_min) t_max = torch.where(cond, invalid_value, t_max) diff --git a/tests/utils/test_parallel_aabb_intersection.py b/tests/utils/test_parallel_aabb_intersection.py new file mode 100644 index 0000000000..1161a0776d --- /dev/null +++ b/tests/utils/test_parallel_aabb_intersection.py @@ -0,0 +1,131 @@ +# Copyright 2026 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# 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. + +"""Test exact slab geometry when ray directions contain zero components.""" + +import pytest +import torch + +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.data.scene_box import OrientedBox, SceneBox +from nerfstudio.utils.math import intersect_aabb, intersect_obb + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("negative_zero", [False, True]) +def test_parallel_rays_on_faces_and_edges_have_finite_intervals(dtype, negative_zero): + aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=dtype) + origins = torch.tensor([[0, 0.5, -1], [1, 0.5, -1], [0, 0, -1], [1, 1, -1]], dtype=dtype) + zero = -0.0 if negative_zero else 0.0 + directions = torch.tensor([[zero, zero, 1]], dtype=dtype).expand_as(origins) + near, far = intersect_aabb(origins, directions, aabb) + torch.testing.assert_close(near, torch.ones_like(near), rtol=0, atol=0) + torch.testing.assert_close(far, torch.full_like(far, 2), rtol=0, atol=0) + + +def test_parallel_rays_outside_slabs_are_misses(): + aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.float64) + origins = torch.tensor([[-0.5, 0.5, -1], [1.5, 0.5, -1], [0.5, -0.1, -1]], dtype=torch.float64) + directions = origins.new_tensor([[0, 0, 1]]).expand_as(origins) + near, far = intersect_aabb(origins, directions, aabb, max_bound=11, invalid_value=123) + torch.testing.assert_close(near, torch.full_like(near, 123), rtol=0, atol=0) + torch.testing.assert_close(far, torch.full_like(far, 123), rtol=0, atol=0) + + +def test_ray_starting_on_face_inside_box_has_zero_near(): + origins = torch.tensor([[0, 1, 0.5]], dtype=torch.float64) + directions = origins.new_tensor([[0, 0, 1]]) + aabb = origins.new_tensor([0, 0, 0, 1, 1, 1]) + near, far = intersect_aabb(origins, directions, aabb) + torch.testing.assert_close(near, near.new_tensor([0]), rtol=0, atol=0) + torch.testing.assert_close(far, far.new_tensor([0.5]), rtol=0, atol=0) + + +def test_interior_parallel_intersection_has_finite_analytic_gradients(): + origins = torch.tensor([[0.5, 0.5, -1]], dtype=torch.float64, requires_grad=True) + directions = torch.tensor([[0, 0, 1]], dtype=torch.float64, requires_grad=True) + aabb = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.float64, requires_grad=True) + near, far = intersect_aabb(origins, directions, aabb) + origin_grad, direction_grad, box_grad = torch.autograd.grad((near + far).sum(), (origins, directions, aabb)) + torch.testing.assert_close(origin_grad, origins.new_tensor([[0, 0, -2]])) + torch.testing.assert_close(direction_grad, directions.new_tensor([[0, 0, -3]])) + torch.testing.assert_close(box_grad, aabb.new_tensor([0, 0, 1, 0, 0, 1])) + assert torch.autograd.gradcheck(lambda o, d, b: intersect_aabb(o, d, b), (origins, directions, aabb)) + + +def test_parallel_misses_have_finite_zero_gradients(): + origins = torch.tensor([[-0.5, 0.5, -1]], dtype=torch.float64, requires_grad=True) + directions = torch.tensor([[0, 0, 1]], dtype=torch.float64, requires_grad=True) + aabb = origins.new_tensor([0, 0, 0, 1, 1, 1]) + near, far = intersect_aabb(origins, directions, aabb) + for gradient in torch.autograd.grad((near + far).sum(), (origins, directions)): + torch.testing.assert_close(gradient, torch.zeros_like(gradient), rtol=0, atol=0) + + +def test_nonzero_small_directions_are_not_treated_as_parallel(): + origins = torch.tensor([[-1e-12, 0.5, -1]], dtype=torch.float64) + directions = origins.new_tensor([[1e-12, 0, 1]]) + aabb = origins.new_tensor([0, 0, 0, 1, 1, 1]) + near, far = intersect_aabb(origins, directions, aabb) + torch.testing.assert_close(near, near.new_tensor([1]), rtol=0, atol=0) + torch.testing.assert_close(far, far.new_tensor([2]), rtol=0, atol=0) + + +def test_oriented_crop_box_handles_ray_on_parallel_face(): + rotation = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=torch.float64) + translation = rotation.new_tensor([2, 3, 4]) + obb = OrientedBox(R=rotation, T=translation, S=rotation.new_tensor([2, 2, 2])) + origins = rotation.new_tensor([[1, 0, -2]]) @ rotation.T + translation + directions = rotation.new_tensor([[0, 0, 1]]) @ rotation.T + near, far = intersect_obb(origins, directions, obb) + torch.testing.assert_close(near, near.new_tensor([1]), rtol=0, atol=0) + torch.testing.assert_close(far, far.new_tensor([3]), rtol=0, atol=0) + + +def test_camera_center_ray_has_finite_crop_bounds_on_box_face(): + camera_to_world = torch.eye(4)[:3].clone() + camera_to_world[:, 3] = torch.tensor([0, 0.5, 2]) + camera = Cameras(camera_to_worlds=camera_to_world, fx=100.0, fy=100.0, cx=0.5, cy=0.5, height=1, width=1) + crop = SceneBox(aabb=torch.tensor([[0, 0, 0], [1, 1, 1]], dtype=torch.float32)) + rays = camera.generate_rays(0, aabb_box=crop) + torch.testing.assert_close(rays.directions, rays.directions.new_tensor([[[0, 0, -1]]]), rtol=0, atol=0) + assert rays.nears is not None and rays.fars is not None + torch.testing.assert_close(rays.nears, rays.nears.new_tensor([[[1]]]), rtol=0, atol=0) + torch.testing.assert_close(rays.fars, rays.fars.new_tensor([[[2]]]), rtol=0, atol=0) + + +def test_nonparallel_rays_match_original_formula_and_gradients(): + generator = torch.Generator().manual_seed(19) + origins = torch.randn(12, 3, generator=generator, dtype=torch.float64).requires_grad_() + directions = torch.randn(12, 3, generator=generator, dtype=torch.float64).requires_grad_() + aabb = origins.new_tensor([-2, -2, -2, 2, 2, 2]) + slabs = torch.stack(((aabb[:3] - origins) / directions, (aabb[3:] - origins) / directions)) + near = slabs.amin(0).amax(-1).clamp(0, 1e10) + far = slabs.amax(0).amin(-1).clamp(0, 1e10) + miss = far <= near + expected = (torch.where(miss, 1e10, near), torch.where(miss, 1e10, far)) + actual = intersect_aabb(origins, directions, aabb) + for a, b in zip(actual, expected): + torch.testing.assert_close(a, b, rtol=0, atol=0) + actual_grads = torch.autograd.grad(sum(t.sum() for t in actual), (origins, directions)) + expected_grads = torch.autograd.grad(sum(t.sum() for t in expected), (origins, directions)) + for a, b in zip(actual_grads, expected_grads): + torch.testing.assert_close(a, b, rtol=1e-12, atol=1e-12) + + +def test_empty_ray_batch_preserves_shapes(): + origins = torch.empty(0, 3) + aabb = origins.new_tensor([0, 0, 0, 1, 1, 1]) + near, far = intersect_aabb(origins, origins, aabb) + assert near.shape == far.shape == (0,)