Merge remote-tracking branch 'origin' into generative

scripts/downloads/download_data.py

@@ -287,6 +287,75 @@ class PhototourismDownload(DatasetDownload):
         os.remove(download_path)
 
 
+# credit to https://autonomousvision.github.io/sdfstudio/
+# pylint: disable=line-too-long
+sdfstudio_downloads = {
+    "sdfstudio-demo-data": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/sdfstudio-demo-data.tar",
+    "dtu": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/DTU.tar",
+    "replica": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/Replica.tar",
+    "scannet": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/scannet.tar",
+    "tanks-and-temple": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/tnt_advanced.tar",
+    "tanks-and-temple-highres": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/highresTNT.tar",
+    "heritage": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/Heritage-Recon.tar",
+    "neural-rgbd-data": "http://kaldir.vc.in.tum.de/neural_rgbd/neural_rgbd_data.zip",
+    "all": None,
+}
+
+SDFstudioCaptureName = tyro.extras.literal_type_from_choices(sdfstudio_downloads.keys())
+
+
+@dataclass
+class SDFstudioDemoDownload(DatasetDownload):
+    """Download the sdfstudio dataset."""
+
+    dataset_name: SDFstudioCaptureName = "sdfstudio-demo-data"
+
+    def download(self, save_dir: Path):
+        """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF)."""
+        # TODO: give this code the same structure as download_nerfstudio
+
+        if self.dataset_name == "all":
+            for dataset_name in sdfstudio_downloads:
+                if dataset_name != "all":
+                    SDFstudioDemoDownload(dataset_name=dataset_name).download(save_dir)
+            return
+
+        assert (
+            self.dataset_name in sdfstudio_downloads
+        ), f"Capture name {self.dataset_name} not found in {sdfstudio_downloads.keys()}"
+
+        url = sdfstudio_downloads[self.dataset_name]
+
+        target_path = str(save_dir / self.dataset_name)
+        os.makedirs(target_path, exist_ok=True)
+
+        file_format = url[-4:]
+
+        download_path = Path(f"{target_path}{file_format}")
+        tmp_path = str(save_dir / ".temp")
+        shutil.rmtree(tmp_path, ignore_errors=True)
+        os.makedirs(tmp_path, exist_ok=True)
+
+        os.system(f"curl -L {url} > {download_path}")
+        if file_format == ".tar":
+            with tarfile.open(download_path, "r") as tar_ref:
+                tar_ref.extractall(str(tmp_path))
+        elif file_format == ".zip":
+            with zipfile.ZipFile(download_path, "r") as zip_ref:
+                zip_ref.extractall(str(target_path))
+            return
+        else:
+            raise NotImplementedError
+
+        inner_folders = os.listdir(tmp_path)
+        assert len(inner_folders) == 1, "There is more than one folder inside this zip file."
+        folder = os.path.join(tmp_path, inner_folders[0])
+        shutil.rmtree(target_path)
+        shutil.move(folder, target_path)
+        shutil.rmtree(tmp_path)
+        os.remove(download_path)
+
+
 Commands = Union[
     Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")],
     Annotated[FriendsDownload, tyro.conf.subcommand(name="friends")],
@@ -294,6 +363,7 @@ Commands = Union[
     Annotated[Record3dDownload, tyro.conf.subcommand(name="record3d")],
     Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")],
     Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")],
+    Annotated[SDFstudioDemoDownload, tyro.conf.subcommand(name="sdfstudio")],
 ]
 
 

scripts/process_data.py

@@ -737,7 +737,7 @@ class ProcessRealityCapture:
     """Process RealityCapture data into a nerfstudio dataset.
 
     This script assumes that cameras have been aligned using RealityCapture. After alignment, it is necessary to
-    export the camera poses as a `.csv` file.
+    export the camera poses as a `.csv` file using the `Internal/External camera parameters` option.
 
     This script does the following:
 

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)