tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
Code

Merge c298945bf3 into 50e0e3c70c

Browse Source
This commit is contained in:
Afloat 2026-10-02 16:34:43 +00:00 committed by GitHub
commit a6dba9285b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 141 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
nerfstudio/utils/math.py
View File

@ -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)

131
tests/utils/test_parallel_aabb_intersection.py Normal file
View File

@ -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,)