Add ns-viewer command to run viewer only with model in eval mode (#1379)

* Add ns-viewer command to run only viewer with model in eval mode

* Remove profiler and logger

* Fix ns-viewer

* ns-viewer remove writer.setup_local_writer

* ns-viewer add writer.setup_local_writer

* ns-viewer: update docs

* Delete outputs

docs/quickstart/first_nerf.md

@@ -30,7 +30,7 @@ Navigating to the link at the end of the terminal will load the webviewer. If yo
 - All data configurations must go at the end. In this case, `nerfstudio-data` and all of its corresponding configurations come at the end after the model and viewer specification.
   :::
 
-## Resume from checkpoint / visualize existing run
+## Resume from checkpoint
 
 It is possible to load a pretrained model by running
 
@@ -40,6 +40,14 @@ ns-train nerfacto --data data/nerfstudio/poster --load-dir {outputs/.../nerfstud
 
 This will automatically start training. If you do not want it to train, add `--viewer.start-train False` to your training command.
 
+## Visualize existing run
+
+Given a pretrained model checkpoint, you can start the viewer by running
+
+```bash
+ns-viewer --load-config {outputs/.../config.yml}
+```
+
 ## Exporting Results
 
 Once you have a NeRF model you can either render out a video or export a point cloud.

pyproject.toml

@@ -100,6 +100,7 @@ ns-install-cli = "scripts.completions.install:entrypoint"
 ns-process-data = "scripts.process_data:entrypoint"
 ns-download-data = "scripts.downloads.download_data:entrypoint"
 ns-train = "scripts.train:entrypoint"
+ns-viewer = "scripts.viewer.run_viewer:entrypoint"
 ns-eval = "scripts.eval:entrypoint"
 ns-render = "scripts.render:entrypoint"
 ns-export = "scripts.exporter:entrypoint"

scripts/viewer/run_viewer.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+"""
+Starts viewer in eval mode.
+"""
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+
+import tyro
+from rich.console import Console
+
+from nerfstudio.configs.base_config import ViewerConfig
+from nerfstudio.engine.trainer import TrainerConfig
+from nerfstudio.pipelines.base_pipeline import Pipeline
+from nerfstudio.utils import writer
+from nerfstudio.utils.eval_utils import eval_setup
+from nerfstudio.utils.writer import EventName, TimeWriter
+from nerfstudio.viewer.server import viewer_utils
+
+CONSOLE = Console(width=120)
+
+
+@dataclass
+class ViewerConfigWithoutNumRays(ViewerConfig):
+    """Configuration for viewer instantiation"""
+
+    num_rays_per_chunk: tyro.conf.Suppress[int] = -1
+    start_train: tyro.conf.Suppress[bool] = False
+
+    def as_viewer_config(self):
+        """Converts the instance to ViewerConfig"""
+        return ViewerConfig(**{x.name: getattr(self, x.name) for x in fields(self)})
+
+
+@dataclass
+class RunViewer:
+    """Load a checkpoint and start the viewer."""
+
+    load_config: Path
+    """Path to config YAML file."""
+    viewer: ViewerConfigWithoutNumRays = field(default_factory=ViewerConfigWithoutNumRays)
+    """Viewer configuration"""
+
+    def main(self) -> None:
+        """Main function."""
+        config, pipeline, _ = eval_setup(
+            self.load_config,
+            eval_num_rays_per_chunk=None,
+            test_mode="test",
+        )
+        num_rays_per_chunk = config.viewer.num_rays_per_chunk
+        assert self.viewer.num_rays_per_chunk == -1
+        config.vis = "viewer"
+        config.viewer = self.viewer.as_viewer_config()
+        config.viewer.num_rays_per_chunk = num_rays_per_chunk
+
+        self._start_viewer(config, pipeline)
+
+    def _start_viewer(self, config: TrainerConfig, pipeline):
+        base_dir = config.get_base_dir()
+        viewer_log_path = base_dir / config.viewer.relative_log_filename
+        viewer_state, banner_messages = viewer_utils.setup_viewer(
+            config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data
+        )
+
+        # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated
+        config.logging.local_writer.enable = False
+        writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages)
+
+        assert viewer_state and pipeline.datamanager.train_dataset
+        viewer_state.init_scene(
+            dataset=pipeline.datamanager.train_dataset,
+            start_train=False,
+        )
+        while True:
+            viewer_state.vis["renderingState/isTraining"].write(False)
+            self._update_viewer_state(viewer_state, config, pipeline)
+
+    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, config: TrainerConfig, pipeline: Pipeline):
+        """Updates the viewer state by rendering out scene with current pipeline
+        Returns the time taken to render scene.
+
+        """
+        # NOTE: step must be > 0 otherwise the rendering would not happen
+        step = 1
+        num_rays_per_batch = config.pipeline.datamanager.train_num_rays_per_batch
+        with TimeWriter(writer, EventName.ITER_VIS_TIME) as _:
+            try:
+                viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch)
+            except RuntimeError:
+                time.sleep(0.03)  # sleep to allow buffer to reset
+                assert viewer_state.vis is not None
+                viewer_state.vis["renderingState/log_errors"].write(
+                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
+                )
+
+    def save_checkpoint(self, *args, **kwargs):
+        """
+        Mock method because we pass this instance to viewer_state.update_scene
+        """
+
+
+def entrypoint():
+    """Entrypoint for use with pyproject scripts."""
+    tyro.extras.set_accent_color("bright_yellow")
+    tyro.cli(RunViewer).main()
+
+
+if __name__ == "__main__":
+    entrypoint()
+
+# For sphinx docs
+get_parser_fn = lambda: tyro.extras.get_parser(RunViewer)  # noqa

scripts/viewer/view_dataset.py

@@ -3,22 +3,21 @@
 view_dataset.py
 """
 
-from rich.console import Console
-
-CONSOLE = Console(width=120)
 import time
 from datetime import timedelta
 from pathlib import Path
 
 import torch
 import tyro
+from rich.console import Console
 
 from nerfstudio.configs.base_config import ViewerConfig
-from nerfstudio.data.datamanagers import AnnotatedDataParserUnion
+from nerfstudio.data.datamanagers.base_datamanager import AnnotatedDataParserUnion
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.viewer.server import viewer_utils
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)
+CONSOLE = Console(width=120)
 
 # speedup for when input size to model doesn't change (much)
 torch.backends.cudnn.benchmark = True  # type: ignore
@@ -30,9 +29,10 @@ def main(
     log_base_dir: Path = Path("/tmp/nerfstudio_viewer_logs"),
 ) -> None:
     """Main function."""
-    viewer_state = viewer_utils.ViewerState(
+    viewer_state, _ = viewer_utils.setup_viewer(
         viewer,
         log_filename=log_base_dir / viewer.relative_log_filename,
+        datapath=dataparser.data,
     )
     dataset = InputDataset(dataparser.setup().get_dataparser_outputs(split="train"))
     viewer_state.init_scene(dataset=dataset, start_train=False)