mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
132 lines
7.0 KiB
Python
132 lines
7.0 KiB
Python
# 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,)
|