From 0777883fb68180617738c4fc2f81b06571fb35af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruilong=20Li=28=E6=9D=8E=E7=91=9E=E9=BE=99=29?=
 <ruilongli94@gmail.com>
Date: Tue, 26 Jul 2022 19:34:29 -0700
Subject: [PATCH] make scene collider optional because NGP does not need it
 (#210)

---
 configs/graph_default.yaml |  1 +
 nerfactory/graphs/base.py  | 12 ++++++++++--
 nerfactory/utils/config.py |  1 +
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml
index 28c26f7e2..285414d74 100644
--- a/configs/graph_default.yaml
+++ b/configs/graph_default.yaml
@@ -81,6 +81,7 @@ data:
 
 graph:
   _target_: nerfactory.graphs.vanilla_nerf.NeRFGraph
+  enable_collider: True
   collider_config:
     _target_: nerfactory.graphs.modules.scene_colliders.NearFarCollider
     near_plane: 2.0
diff --git a/nerfactory/graphs/base.py b/nerfactory/graphs/base.py
index 9824d02eb..a163af624 100644
--- a/nerfactory/graphs/base.py
+++ b/nerfactory/graphs/base.py
@@ -81,6 +81,7 @@ class Graph(AbstractGraph):
         camera_to_world (torch.Tensor): Camera to world transformation.
         loss_coefficients (DictConfig): Loss specific weights.
         scene_bounds (SceneBounds): Bounds of target scene.
+        enable_collider (bool): Whether to create a scene collider to filter rays.
         collider_config (DictConfig): Configuration of scene collider.
         enable_density_field (bool): Whether to create a density field to filter samples.
         density_field_config (DictConfig): Configuration of density field.
@@ -92,6 +93,7 @@ class Graph(AbstractGraph):
         camera_to_world: torch.Tensor = None,
         loss_coefficients: DictConfig = None,
         scene_bounds: SceneBounds = None,
+        enable_collider: bool = True,
         collider_config: DictConfig = None,
         enable_density_field: bool = False,
         density_field_config: DictConfig = None,
@@ -102,6 +104,7 @@ class Graph(AbstractGraph):
         self.intrinsics = intrinsics
         self.camera_to_world = camera_to_world
         self.scene_bounds = scene_bounds
+        self.enable_collider = enable_collider
         self.collider_config = collider_config
         self.loss_coefficients = loss_coefficients
         self.enable_density_field = enable_density_field
@@ -130,7 +133,8 @@ class Graph(AbstractGraph):
 
     def populate_collider(self):
         """Set the scene bounds collider to use."""
-        self.collider = instantiate_from_dict_config(self.collider_config, scene_bounds=self.scene_bounds)
+        if self.enable_collider:
+            self.collider = instantiate_from_dict_config(self.collider_config, scene_bounds=self.scene_bounds)
 
     @abstractmethod
     def populate_misc_modules(self):
@@ -165,7 +169,11 @@ class Graph(AbstractGraph):
 
     def forward_after_ray_generator(self, ray_bundle: RayBundle, batch: Union[str, Dict[str, torch.tensor]] = None):
         """Run forward starting with a ray bundle."""
-        intersected_ray_bundle = self.collider(ray_bundle)
+        if self.collider is not None:
+            intersected_ray_bundle = self.collider(ray_bundle)
+        else:
+            # NOTE(ruilongli): we don't need collider for ngp
+            intersected_ray_bundle = ray_bundle
 
         if batch is None:
             # during inference, keep all rays
diff --git a/nerfactory/utils/config.py b/nerfactory/utils/config.py
index 6f1044054..aa3c06a64 100644
--- a/nerfactory/utils/config.py
+++ b/nerfactory/utils/config.py
@@ -97,6 +97,7 @@ class GraphConfig:
     """Configuration for graph instantiation"""
 
     _target_: str = MISSING
+    enable_collider: Optional[bool] = True
     collider_config: Dict[str, Any] = MISSING
     num_coarse_samples: int = MISSING
     num_importance_samples: int = MISSING