From 695d9a17777ea075425f64bf72f21818cac31325 Mon Sep 17 00:00:00 2001
From: Matthew Tancik <tancik@berkeley.edu>
Date: Thu, 15 Sep 2022 17:42:42 -0700
Subject: [PATCH] Add near plane to NGP sampler (#443)

---
 nerfactory/configs/base.py               |  2 ++
 nerfactory/models/compound.py            |  4 +++-
 nerfactory/models/instant_ngp.py         |  4 +++-
 nerfactory/models/modules/ray_sampler.py | 11 ++++++-----
 4 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py
index 26a36de7d..05bf21e4d 100644
--- a/nerfactory/configs/base.py
+++ b/nerfactory/configs/base.py
@@ -407,6 +407,8 @@ class CompoundModelConfig(ModelConfig):
     field_implementation: Literal["torch", "tcnn"] = "tcnn"  # torch, tcnn, ...
     loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0})
     num_samples: int = 1024  # instead of course/fine samples
+    cone_angle: float = 0.0
+    near_plane: float = 0.5
 
 
 @dataclass
diff --git a/nerfactory/models/compound.py b/nerfactory/models/compound.py
index ae9fdcc24..488dc7633 100644
--- a/nerfactory/models/compound.py
+++ b/nerfactory/models/compound.py
@@ -115,7 +115,9 @@ class CompoundModel(Model):
 
         if self.field is None:
             raise ValueError("populate_fields() must be called before get_outputs")
-        ray_samples, packed_info, t_min, t_max = self.sampler(ray_bundle, self.field.aabb)
+        ray_samples, packed_info, t_min, t_max = self.sampler(
+            ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle, near_plane=self.config.near_plane
+        )
 
         field_outputs = self.field.forward(ray_samples)
 
diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py
index b9357ec06..e659721f8 100644
--- a/nerfactory/models/instant_ngp.py
+++ b/nerfactory/models/instant_ngp.py
@@ -63,6 +63,8 @@ class InstantNGPModelConfig(cfg.ModelConfig):
     """Number of samples in field evaluation. Defaults to 1024,"""
     cone_angle: float = 0.0
     """Should be set to 0.0 for blender scenes but 1./256 for real scenes."""
+    near_plane: float = 0.05
+    """How far along ray to start sampling."""
     randomize_background: bool = False
     """Whether to randomize the background color. Defaults to False."""
 
@@ -128,7 +130,7 @@ class NGPModel(Model):
         device = ray_bundle.origins.device
 
         ray_samples, packed_info, t_min, t_max = self.sampler(
-            ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle
+            ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle, near_plane=self.config.near_plane
         )
 
         field_outputs = self.field.forward(ray_samples)
diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py
index bcccfb4ba..9376dfe7e 100644
--- a/nerfactory/models/modules/ray_sampler.py
+++ b/nerfactory/models/modules/ray_sampler.py
@@ -360,10 +360,9 @@ class NGPSpacedSampler(Sampler):
         ray_bundle: RayBundle,
         aabb: TensorType[2, 3],
         num_samples: Optional[int] = None,
-        cone_angle: Optional[float] = 0.0,
+        cone_angle: float = 0.0,
         marching_steps: Optional[int] = 128,
-        t_min: Optional[TensorType["num_rays"]] = None,
-        t_max: Optional[TensorType["num_rays"]] = None,
+        near_plane: float = 0.0,
     ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples"], TensorType["total_samples"]]:
         """Generate ray samples in a bounding box.
 
@@ -375,6 +374,8 @@ class NGPSpacedSampler(Sampler):
             aabb: Bounding box of the scene.
             num_samples: Number of samples per ray
             cone_angle: 0.0 for nerf-synthetic and 1.0/256 for real scenes
+            marching_steps: Number of marching steps (Note: not currently used)
+            near_plane: Near plane for raymarching
 
         Returns:
             First return is ray samples in a packed way where only the valid samples are kept.
@@ -392,8 +393,8 @@ class NGPSpacedSampler(Sampler):
         rays_d = ray_bundle.directions.contiguous()
         scene_scale = (aabb[3:6] - aabb[0:3]).max()
 
-        if t_min is None or t_max is None:
-            t_min, t_max = nerfactory_cuda.ray_aabb_intersect(rays_o, rays_d, aabb)
+        t_min, t_max = nerfactory_cuda.ray_aabb_intersect(rays_o, rays_d, aabb)
+        t_min = torch.max(t_min, torch.ones_like(t_min) * near_plane)
 
         marching_steps = -1  # disable marching mode for now
         # TODO(ruilongli): * 16 is for original impl for training. Need to run