diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index e7071c329..73e40a9d8 100644 --- a/docs/quickstart/first_nerf.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index ba57c01d4..4ee2042f1 100644 --- a/pyproject.toml +++ b/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" diff --git a/scripts/viewer/__init__.py b/scripts/viewer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/viewer/run_viewer.py b/scripts/viewer/run_viewer.py new file mode 100644 index 000000000..9a96c6703 --- /dev/null +++ b/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 diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 1f5aa2201..4bd526552 100644 --- a/scripts/viewer/view_dataset.py +++ b/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)