cleanup

nerfstudio/models/nerfacto.py

@@ -137,8 +137,7 @@ class NerfactoModel(Model):
         """Set the fields and modules."""
         super().populate_modules()
 
-        # TODO(ethan): flag to enable scene contraction or not
-
+        # flag to enable scene contraction or not
         if self.config.use_scene_contraction:
             if self.config.scene_contraction_norm == "inf":
                 order = float("inf")
@@ -169,8 +168,7 @@ class NerfactoModel(Model):
         # Build the proposal network(s)
         self.proposal_networks = torch.nn.ModuleList()
         if self.config.use_same_proposal_network:
-            # assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
-            # using the last defined proposal network
+            # use the last defined proposal network
             prop_net_args = 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)

nerfstudio/pipelines/base_pipeline.py

@@ -340,7 +340,6 @@ class VanillaPipeline(Pipeline):
         param_groups = self.datamanager.get_param_groups()
         optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=1e-3, eps=1e-15)
 
-        # CONSOLE.print("Optimizing camera poses...")
         # rescale the image to 1/4 resolution
         self.datamanager.eval_ray_generator.cameras.rescale_output_resolution(self.config.eval_image_scale_factor)
         self.train()

projects/nerfstudio_paper/.gitignore

@@ -0,0 +1,3 @@
+camera_paths/
+notes/
+renders/

projects/nerfstudio_paper/README.md

@@ -0,0 +1,25 @@
+# Nerfstudio paper
+
+> Below we outline the steps to run the experiments from the paper.
+
+1. Download the MIP-NeRF 360 data and the Nerfstudio Dataset.
+
+```bash
+ns-download-data mipnerf360
+```
+
+```bash
+ns-download-data nerfstudio --capture-name nerfstudio-dataset
+```
+
+2. Run Nerfacto on the MIP-NeRF 360 data.
+
+```bash
+python projects/nerfstudio_paper/benchmark_nerfstudio_paper.py nerfacto-on-mipnerf360 --dry-run
+```
+
+3. Run Nerfacto ablation experiments.
+
+```bash
+python projects/nerfstudio_paper/benchmark_nerfstudio_paper.py nerfacto-ablations --dry-run
+```

projects/nerfstudio_paper/benchmark_nerfstudio_paper.py

@@ -0,0 +1,212 @@
+"""
+Benchmarking script for nerfstudio paper.
+
+- nerfacto and instant-ngp methods on mipnerf360 data
+- nerfacto ablations
+"""
+
+import threading
+import time
+from pathlib import Path
+from typing import Union
+from dataclasses import dataclass
+
+import tyro
+from typing_extensions import Annotated
+
+import GPUtil
+
+from nerfstudio.configs.base_config import PrintableConfig
+from nerfstudio.utils.scripts import run_command
+
+# for the mipnerf360 experiments
+mipnerf360_capture_names = ["bicycle", "garden", "stump", "room", "counter", "kitchen", "bonsai"]  # 7 splits
+mipnerf360_table_rows = [
+    # nerfacto method
+    (
+        "nerfacto-w/o-pose-app",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras False --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False nerfstudio-data --downscale-factor 4 --train-split-percentage 0.875 mipnerf360-data",
+    ),
+]
+
+# for the ablation experiments
+ablations_capture_names = [
+    "Egypt",
+    "person",
+    "kitchen",
+    "plane",
+    "dozer",
+    "floating-tree",
+    "aspen",
+    "stump",
+    "sculpture",
+    "Giannini-Hall",
+]
+
+ablations_table_rows = [
+    ("nerfacto", "nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True"),
+    (
+        "w/o-pose",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.datamanager.camera-optimizer.mode off",
+    ),
+    (
+        "w/o-app",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.model.use-appearance-embedding False",
+    ),
+    (
+        "w/o-pose-app",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False",
+    ),
+    (
+        "1-prop-network",
+        "nerfacto",
+        '--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.num-proposal-samples-per-ray "256" --pipeline.model.num_proposal_iterations 1',
+    ),
+    (
+        "l2-contraction",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.scene-contraction-norm l2",
+    ),
+    (
+        "shared-prop-network",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-same-proposal-network True",
+    ),
+    (
+        "random-background-color",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.background-color random",
+    ),
+    (
+        "no-contraction",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-bounded True --pipeline.model.use-scene-contraction False nerfstudio-data --scale_factor 0.125",
+    ),
+    (
+        "synthetic-on-real",
+        "nerfacto",
+        "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False --pipeline.model.use-bounded True --pipeline.model.use-scene-contraction False nerfstudio-data --scale_factor 0.125",
+    ),
+]
+
+
+def launch_experiments(capture_names, table_rows, data_path: Path = Path("data/nerfstudio"), dry_run: bool = False):
+    """Launch the experiments."""
+
+    # make a list of all the jobs that need to be fun
+    jobs = []
+    for capture_name in capture_names:
+
+        for table_row_name, method, table_row_command in table_rows:
+            command = " ".join(
+                (
+                    f"ns-train {method}",
+                    "--vis wandb",
+                    f"--data { data_path / capture_name}",
+                    "--output-dir outputs/nerfacto-ablations",
+                    "--steps-per-eval-batch 0 --steps-per-eval-image 0",
+                    "--steps-per-eval-all-images 5000 --max-num-iterations 30001",
+                    f"--wandb-name {capture_name}_{table_row_name}",
+                    f"--experiment-name {capture_name}_{table_row_name}",
+                    # extra_string,
+                    table_row_command,
+                )
+            )
+            jobs.append(command)
+
+    while jobs:
+        # get GPUs that capacity to run these jobs
+        gpu_devices_available = GPUtil.getAvailable(order="first", limit=10, maxMemory=0.1)
+
+        print("Available GPUs: ", gpu_devices_available)
+
+        # thread list
+        threads = []
+        while gpu_devices_available and jobs:
+            gpu = gpu_devices_available.pop(0)
+            command = f"CUDA_VISIBLE_DEVICES={gpu} " + jobs.pop(0)
+
+            def task():
+                print("Starting command: ", command)
+                if not dry_run:
+                    _ = run_command(command, verbose=False)
+                print("Finished command: ", command)
+
+            threads.append(threading.Thread(target=task))
+            threads[-1].start()
+
+            # NOTE(ethan): here we need a delay, otherwise the wandb/tensorboard naming is messed up... not sure why
+            if not dry_run:
+                time.sleep(5)
+
+        # wait for all threads to finish
+        for t in threads:
+            t.join()
+
+        print("Finished all threads")
+
+
+@dataclass
+class Benchmark(PrintableConfig):
+    """Benchmark code."""
+
+    dry_run: bool = False
+
+    def main(self) -> None:
+        """Run the code."""
+        raise NotImplementedError
+
+
+@dataclass
+class BenchmarkMipNeRF360(Benchmark):
+    """Benchmark MipNeRF-360."""
+
+    def main(self, dry_run: bool = False):
+        launch_experiments(
+            mipnerf360_capture_names, mipnerf360_table_rows, data_path=Path("data/nerfstudio-data-mipnerf360", dry_run=dry_run)
+        )
+
+
+@dataclass
+class BenchmarkAblations(Benchmark):
+    """Benchmark ablations."""
+
+    def main(self, dry_run: bool = False):
+        launch_experiments(ablations_capture_names, ablations_table_rows, dry_run=dry_run)
+
+
+Commands = Union[
+    Annotated[BenchmarkMipNeRF360, tyro.conf.subcommand(name="nerfacto-on-mipnerf360")],
+    Annotated[BenchmarkAblations, tyro.conf.subcommand(name="nerfacto-ablations")],
+]
+
+
+def main(
+    benchmark: Benchmark,
+):
+    """Script to run the benchmark experiments for the Nerfstudio paper.
+    - nerfacto-on-mipnerf360: The MipNeRF-360 experiments on the MipNeRF-360 Dataset.
+    - nerfacto-ablations: The Nerfacto ablations on the Nerfstudio Dataset.
+
+    Args:
+        benchmark: The benchmark to run.
+    """
+
+    benchmark.main(dry_run=benchmark.dry_run)
+
+
+def entrypoint():
+    """Entrypoint for use with pyproject scripts."""
+    tyro.extras.set_accent_color("bright_yellow")
+    main(tyro.cli(Commands))
+
+
+if __name__ == "__main__":
+    entrypoint()
+
+# For sphinx docs
+get_parser_fn = lambda: tyro.extras.get_parser(Commands)  # noqa

pyproject.toml

@@ -105,7 +105,7 @@ dependency_links = [
 ]
 
 [tool.setuptools.packages.find]
-include = ["nerfstudio*","scripts*"]
+include = ["nerfstudio*","scripts*","projects*"]
 
 [tool.setuptools.package-data]
 "*" = ["*.json", "py.typed", "setup.bash", "setup.zsh"]