doubling nerfacto samples on eval (#675)

doubling samples and eval num ray options

nerfstudio/models/base_model.py

@@ -169,6 +169,7 @@ class Model(nn.Module):
         num_rays = len(camera_ray_bundle)
         outputs_lists = defaultdict(list)
         for i in range(0, num_rays, num_rays_per_chunk):
+            print(self.config.eval_num_rays_per_chunk)
             start_idx = i
             end_idx = i + num_rays_per_chunk
             ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx)

nerfstudio/utils/eval_utils.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 import os
 import sys
 from pathlib import Path
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 import yaml
@@ -63,7 +63,9 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path:
     return load_path
 
 
-def eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, 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]:
     """Shared setup for loading a saved pipeline for evaluation.
 
     Args:
@@ -76,6 +78,17 @@ def eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]:
     config = yaml.load(config_path.read_text(), Loader=yaml.Loader)
     assert isinstance(config, cfg.Config)
 
+    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()

scripts/render.py

@@ -9,7 +9,7 @@ import json
 import logging
 import sys
 from pathlib import Path
-from typing import Literal
+from typing import Literal, Optional
 
 import dcargs
 import mediapy as media
@@ -103,10 +103,18 @@ class RenderTrajectory:
     output_path: Path = Path("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
 
     def main(self) -> None:
         """Main function."""
-        _, pipeline, _ = eval_setup(self.load_config)
+        _, 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,
+        )
 
         seconds = self.seconds