Reducing flicker artifacts for nerfacto method (#750)

* pose opt default nerfacto.

* added SO3 cam pose optimization which is slightly faster (~2.2m/iter) than BARF

* added occupancy grid with placeholder until nerfacc is updated to query density from it

* added occupancy grid sampling as the level 0 in proposal hierarchy, relies on nerfacc 0.1.7 (on pypi)

* merge main

* merge main

* minor changes to sampling

* adjusted sample number

* refactor occ grid stuff, settle on decent params

* nograd proposal sampling

* param tweaks

* varying resolution proposal networks

* params

* merged main

* optimizer update

* remove occ grid

* remove double nerfacto samples

* set single jitter to true

* remove pdb

* more changes

* fix method defaults

* fix pylint

* remove occupancy fn

Co-authored-by: akristoffersen <akristoffersen@berkeley.edu>
Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>

nerfstudio/configs/method_configs.py

@@ -56,8 +56,10 @@ method_configs["nerfacto"] = Config(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
-            eval_num_rays_per_batch=8192,
-            camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3)
+            ),
         ),
         model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 14),
     ),

nerfstudio/engine/optimizers.py

@@ -51,6 +51,7 @@ class AdamOptimizerConfig(OptimizerConfig):
     """Basic optimizer config with Adam"""
 
     _target: Type = torch.optim.Adam
+    weight_decay: float = 0
 
 
 @dataclass

nerfstudio/fields/density_fields.py

@@ -55,6 +55,11 @@ class HashMLPDensityField(Field):
         hidden_dim: int = 64,
         spatial_distortion: Optional[SpatialDistortion] = None,
         use_linear=False,
+        num_levels=8,
+        max_res=1024,
+        base_res=16,
+        log2_hashmap_size=18,
+        features_per_level=2,
     ) -> None:
         super().__init__()
 
@@ -62,11 +67,6 @@ class HashMLPDensityField(Field):
         self.spatial_distortion = spatial_distortion
         self.use_linear = use_linear
 
-        num_levels = 8
-        max_res = 1024
-        base_res = 16
-        log2_hashmap_size = 18
-        features_per_level = 2
         growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
         config = {

nerfstudio/model_components/ray_samplers.py

@@ -315,7 +315,6 @@ class PDFSampler(Sampler):
             # Stratified samples between 0 and 1
             u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
             u = u.expand(size=(*cdf.shape[:-1], num_bins))
-            # u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins
             if self.single_jitter:
                 rand = torch.rand((*cdf.shape[:-1], 1), device=cdf.device) / num_bins
             else:
@@ -469,7 +468,6 @@ class VolumetricSampler(Sampler):
             cone_angle=cone_angle,
             alpha_thre=1e-2,
         )
-
         num_samples = starts.shape[0]
         if num_samples == 0:
             # create a single fake sample and update packed_info accordingly
@@ -506,7 +504,7 @@ class ProposalNetworkSampler(Sampler):
         num_proposal_samples_per_ray: Tuple[int] = (64,),
         num_nerf_samples_per_ray: int = 32,
         num_proposal_network_iterations: int = 2,
-        single_jitter: bool = True,
+        single_jitter: bool = False,
     ) -> None:
         super().__init__()
         self.num_proposal_samples_per_ray = num_proposal_samples_per_ray

nerfstudio/models/nerfacto.py

@@ -61,14 +61,21 @@ class NerfactoModelConfig(ModelConfig):
     """How far along the ray to stop sampling."""
     background_color: Literal["background", "last_sample"] = "last_sample"
     """Whether to randomize the background color."""
-    num_proposal_samples_per_ray: Tuple[int] = (64,)
+    num_proposal_samples_per_ray: Tuple[int] = (200, 64)
     """Number of samples per ray for the proposal network."""
     num_nerf_samples_per_ray: int = 64
     """Number of samples per ray for the nerf network."""
-    num_proposal_network_iterations: int = 1
+    num_proposal_iterations: int = 2
     """Number of proposal network iterations."""
     use_same_proposal_network: bool = False
     """Use the same proposal network. Otherwise use different ones."""
+    proposal_net_args_list: List[Dict] = field(
+        default_factory=lambda: [
+            {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 64},
+            {"hidden_dim": 16, "log2_hashmap_size": 16, "num_levels": 5, "max_res": 256},
+        ]
+    )
+    """Arguments for the proposal density fields."""
     interlevel_loss_mult: float = 1.0
     """Proposal loss multiplier."""
     distortion_loss_mult: float = 0.002
@@ -108,29 +115,36 @@ class NerfactoModel(Model):
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,
         )
 
+        self.density_fns = []
+        num_prop_nets = self.config.num_proposal_iterations
         # Build the proposal network(s)
         self.proposal_networks = torch.nn.ModuleList()
         if self.config.use_same_proposal_network:
-            network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction)
+            assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
+            prop_net_args = self.config.proposal_net_args_list[0]
+            network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args)
             self.proposal_networks.append(network)
-            self.density_fns = [network.density_fn for _ in range(self.config.num_proposal_network_iterations)]
+            self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
         else:
-            for _ in range(self.config.num_proposal_network_iterations):
-                network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction)
+            for i in range(num_prop_nets):
+                prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)]
+                network = HashMLPDensityField(
+                    self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args
+                )
                 self.proposal_networks.append(network)
-            self.density_fns = [network.density_fn for network in self.proposal_networks]
-
-        # Collider
-        self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
+            self.density_fns.extend([network.density_fn for network in self.proposal_networks])
 
         # Samplers
         self.proposal_sampler = ProposalNetworkSampler(
             num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray,
             num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray,
-            num_proposal_network_iterations=self.config.num_proposal_network_iterations,
+            num_proposal_network_iterations=self.config.num_proposal_iterations,
             single_jitter=self.config.use_single_jitter,
         )
 
+        # Collider
+        self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
+
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
@@ -189,7 +203,7 @@ class NerfactoModel(Model):
         outputs["weights_list"] = weights_list
         outputs["ray_samples_list"] = ray_samples_list
 
-        for i in range(self.config.num_proposal_network_iterations):
+        for i in range(self.config.num_proposal_iterations):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
 
         return outputs
@@ -242,7 +256,7 @@ class NerfactoModel(Model):
 
         images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
 
-        for i in range(self.config.num_proposal_network_iterations):
+        for i in range(self.config.num_proposal_iterations):
             key = f"prop_depth_{i}"
             prop_depth_i = colormaps.apply_depth_colormap(
                 outputs[key],

nerfstudio/utils/eval_utils.py

@@ -63,9 +63,7 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path:
     return load_path
 
 
-def eval_setup(
-    config_path: Path, double_nerfacto_nerf_samples: bool = False, eval_num_rays_per_chunk: Optional[int] = None
-) -> Tuple[cfg.Config, Pipeline, Path]:
+def eval_setup(config_path: Path, eval_num_rays_per_chunk: Optional[int] = None) -> Tuple[cfg.Config, Pipeline, Path]:
     """Shared setup for loading a saved pipeline for evaluation.
 
     Args:
@@ -81,14 +79,6 @@ def eval_setup(
     if eval_num_rays_per_chunk:
         config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk
 
-    if config.method_name == "nerfacto" and double_nerfacto_nerf_samples:
-        config.pipeline.model.num_nerf_samples_per_ray *= 2
-        console.print(
-            "[bold yellow]Warning: doubling the number of nerfacto nerf samples to "
-            f"{config.pipeline.model.num_nerf_samples_per_ray} for improved quality."
-        )
-        console.print("[bold yellow]This logic is a hack and will be removed in the future.")
-
     # load checkpoints from wherever they were saved
     # TODO: expose the ability to choose an arbitrary checkpoint
     config.trainer.load_dir = config.get_checkpoint_dir()

pyproject.toml

@@ -107,6 +107,7 @@ generated-members = ["numpy.*", "torch.*", "cv2.*", "cv.*"]
 good-names-rgxs = "^[_a-zA-Z][_a-z0-9]?$"
 ignore-paths = ["scripts/colmap2nerf.py"]
 jobs = 0
+ignored-classes = ["TensorDataclass"]
 
 disable = [
   "duplicate-code",

scripts/render.py

@@ -106,8 +106,6 @@ class RenderTrajectory:
     output_path: Path = Path("renders/output.mp4")
     # How long the video should be.
     seconds: float = 5.0
-    # A hack to double the number of samples for the nerfacto method.
-    double_nerfacto_nerf_samples: bool = True
     # Specifies number of rays per chunk during eval.
     eval_num_rays_per_chunk: Optional[int] = None
 
@@ -115,7 +113,6 @@ class RenderTrajectory:
         """Main function."""
         _, pipeline, _ = eval_setup(
             self.load_config,
-            double_nerfacto_nerf_samples=self.double_nerfacto_nerf_samples,
             eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
         )