Per-ray minimum and maximum distances in sampling method (#205)

* Fixed traversal with far_plane

* test added for traversal with near and far planes

* More correct test

* black formatting for test

* Per-ray min-max distances in sampling

* Test for sampling with min/max distances added

* black formatting

nerfacc/estimators/occ_grid.py

@@ -93,6 +93,8 @@ class OccGridEstimator(AbstractEstimator):
         alpha_fn: Optional[Callable] = None,
         near_plane: float = 0.0,
         far_plane: float = 1e10,
+        t_min: Optional[Tensor] = None,  # [n_rays]
+        t_max: Optional[Tensor] = None,  # [n_rays]
         # rendering options
         render_step_size: float = 1e-3,
         early_stop_eps: float = 1e-4,
@@ -120,6 +122,10 @@ class OccGridEstimator(AbstractEstimator):
                 You should only provide either `sigma_fn` or `alpha_fn`.
             near_plane: Optional. Near plane distance. Default: 0.0.
             far_plane: Optional. Far plane distance. Default: 1e10.
+            t_min: Optional. Per-ray minimum distance. Tensor with shape (n_rays).
+                If profided, the marching will start from maximum of t_min and near_plane.
+            t_max: Optional. Per-ray maximum distance. Tensor with shape (n_rays).
+                If profided, the marching will stop by minimum of t_max and far_plane.
             render_step_size: Step size for marching. Default: 1e-3.
             early_stop_eps: Early stop threshold for skipping invisible space. Default: 1e-4.
             alpha_thre: Alpha threshold for skipping empty space. Default: 0.0.
@@ -147,6 +153,12 @@ class OccGridEstimator(AbstractEstimator):
 
         near_planes = torch.full_like(rays_o[..., 0], fill_value=near_plane)
         far_planes = torch.full_like(rays_o[..., 0], fill_value=far_plane)
+
+        if t_min is not None:
+            near_planes = torch.clamp(near_planes, min=t_min)
+        if t_max is not None:
+            far_planes = torch.clamp(far_planes, max=t_max)
+
         if stratified:
             near_planes += torch.rand_like(near_planes) * render_step_size
         intervals, samples = traverse_grids(

tests/test_grid.py

@@ -96,7 +96,52 @@ def test_traverse_grids_with_near_far_planes():
     assert (intervals.vals <= (far_planes + step_size / 2)).all()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available, reason="No CUDA device")
+def test_sampling_with_min_max_distances():
+    from nerfacc import OccGridEstimator
+
+    torch.manual_seed(42)
+    n_rays = 64
+    levels = 4
+    resolution = 32
+    render_step_size = 0.01
+    near_plane = 0.15
+    far_plane = 0.85
+
+    rays_o = torch.rand((n_rays, 3), device=device) * 2 - 1.0
+    rays_d = torch.rand((n_rays, 3), device=device)
+    rays_d = rays_d / rays_d.norm(dim=-1, keepdim=True)
+
+    aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0], device=device)
+    binaries = (
+        torch.rand((levels, resolution, resolution, resolution), device=device)
+        > 0.5
+    )
+    t_min = torch.rand((n_rays,), device=device)
+    t_max = t_min + torch.rand((n_rays,), device=device)
+
+    grid_estimator = OccGridEstimator(
+        roi_aabb=aabb, resolution=resolution, levels=levels
+    )
+
+    grid_estimator.binaries = binaries
+
+    ray_indices, t_starts, t_ends = grid_estimator.sampling(
+        rays_o=rays_o,
+        rays_d=rays_d,
+        near_plane=near_plane,
+        far_plane=far_plane,
+        t_min=t_min,
+        t_max=t_max,
+        render_step_size=render_step_size,
+    )
+
+    assert (t_starts >= (t_min[ray_indices] - render_step_size / 2)).all()
+    assert (t_ends <= (t_max[ray_indices] + render_step_size / 2)).all()
+
+
 if __name__ == "__main__":
     test_ray_aabb_intersect()
     test_traverse_grids()
     test_traverse_grids_with_near_far_planes()
+    test_sampling_with_min_max_distances()