ns-render - render all images (#2459)

* ns-render - render all images

* Fix commonprefix->commonpath

* Fix raw rendering

nerfstudio/scripts/render.py

@@ -23,7 +23,8 @@ import os
 import struct
 import shutil
 import sys
-from contextlib import ExitStack
+import gzip
+from contextlib import ExitStack, contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Union
@@ -52,8 +53,11 @@ from nerfstudio.cameras.camera_paths import (
     get_path_from_json,
     get_spiral_path,
 )
-from nerfstudio.cameras.cameras import Cameras, CameraType
-from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
+from nerfstudio.cameras.cameras import Cameras, CameraType, RayBundle
+from nerfstudio.data.datasets.base_dataset import Dataset
+from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig
+from nerfstudio.data.utils.dataloaders import FixedIndicesEvalDataloader
+from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.model_components import renderers
 from nerfstudio.pipelines.base_pipeline import Pipeline
@@ -585,11 +589,209 @@ class SpiralRender(BaseRender):
         )
 
 
+@contextmanager
+def _disable_datamanager_setup(cls):
+    """
+    Disables setup_train or setup_eval for faster initialization.
+    """
+    old_setup_train = getattr(cls, "setup_train")
+    old_setup_eval = getattr(cls, "setup_eval")
+    setattr(cls, "setup_train", lambda *args, **kwargs: None)
+    setattr(cls, "setup_eval", lambda *args, **kwargs: None)
+    yield cls
+    setattr(cls, "setup_train", old_setup_train)
+    setattr(cls, "setup_eval", old_setup_eval)
+
+
+@dataclass
+class DatasetRender(BaseRender):
+    """Render all images in the dataset."""
+
+    output_path: Path = Path("renders")
+    """Path to output video file."""
+    data: Optional[Path] = None
+    """Override path to the dataset."""
+    downscale_factor: Optional[float] = None
+    """Scaling factor to apply to the camera image resolution."""
+    split: Literal["train", "val", "test", "train+test"] = "test"
+    """Split to render."""
+    rendered_output_names: Optional[List[str]] = field(default_factory=lambda: None)
+    """Name of the renderer outputs to use. rgb, depth, raw-depth, gt-rgb etc. By default all outputs are rendered."""
+
+    def main(self):
+        config: TrainerConfig
+
+        def update_config(config: TrainerConfig) -> TrainerConfig:
+            data_manager_config = config.pipeline.datamanager
+            assert isinstance(data_manager_config, VanillaDataManagerConfig)
+            data_manager_config.eval_image_indices = None
+            data_manager_config.eval_num_images_to_sample_from = -1
+            data_manager_config.eval_num_times_to_repeat_images = -1
+            data_manager_config.train_num_images_to_sample_from = -1
+            data_manager_config.train_num_times_to_repeat_images = -1
+            data_manager_config.data = self.data
+            if self.downscale_factor is not None:
+                assert hasattr(data_manager_config.dataparser, "downscale_factor")
+                setattr(data_manager_config.dataparser, "downscale_factor", self.downscale_factor)
+            return config
+
+        config, pipeline, _, _ = eval_setup(
+            self.load_config,
+            eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
+            test_mode="inference",
+            update_config_callback=update_config,
+        )
+        data_manager_config = config.pipeline.datamanager
+        assert isinstance(data_manager_config, VanillaDataManagerConfig)
+
+        for split in self.split.split("+"):
+            datamanager: VanillaDataManager
+            dataset: Dataset
+            if split == "train":
+                with _disable_datamanager_setup(data_manager_config._target):  # pylint: disable=protected-access
+                    datamanager = data_manager_config.setup(test_mode="test", device=pipeline.device)
+
+                dataset = datamanager.train_dataset
+                dataparser_outputs = getattr(dataset, "_dataparser_outputs", datamanager.train_dataparser_outputs)
+            else:
+                with _disable_datamanager_setup(data_manager_config._target):  # pylint: disable=protected-access
+                    datamanager = data_manager_config.setup(test_mode=split, device=pipeline.device)
+
+                dataset = datamanager.eval_dataset
+                dataparser_outputs = getattr(dataset, "_dataparser_outputs", None)
+                if dataparser_outputs is None:
+                    dataparser_outputs = datamanager.dataparser.get_dataparser_outputs(split=datamanager.test_split)
+            dataloader = FixedIndicesEvalDataloader(
+                input_dataset=dataset,
+                device=datamanager.device,
+                num_workers=datamanager.world_size * 4,
+            )
+            images_root = Path(os.path.commonpath(dataparser_outputs.image_filenames))
+            with Progress(
+                TextColumn(f":movie_camera: Rendering split {split} :movie_camera:"),
+                BarColumn(),
+                TaskProgressColumn(
+                    text_format="[progress.percentage]{task.completed}/{task.total:>.0f}({task.percentage:>3.1f}%)",
+                    show_speed=True,
+                ),
+                ItersPerSecColumn(suffix="fps"),
+                TimeRemainingColumn(elapsed_when_finished=False, compact=False),
+                TimeElapsedColumn(),
+            ) as progress:
+                for camera_idx, (ray_bundle, batch) in enumerate(progress.track(dataloader, total=len(dataset))):
+                    ray_bundle: RayBundle
+                    with torch.no_grad():
+                        outputs = pipeline.model.get_outputs_for_camera_ray_bundle(ray_bundle)
+
+                    gt_batch = batch.copy()
+                    gt_batch["rgb"] = gt_batch.pop("image")
+                    all_outputs = (
+                        list(outputs.keys())
+                        + [f"raw-{x}" for x in outputs.keys()]
+                        + [f"gt-{x}" for x in gt_batch.keys()]
+                        + [f"raw-gt-{x}" for x in gt_batch.keys()]
+                    )
+                    rendered_output_names = self.rendered_output_names
+                    if rendered_output_names is None:
+                        rendered_output_names = ["gt-rgb"] + list(outputs.keys())
+                    for rendered_output_name in rendered_output_names:
+                        if rendered_output_name not in all_outputs:
+                            CONSOLE.rule("Error", style="red")
+                            CONSOLE.print(
+                                f"Could not find {rendered_output_name} in the model outputs", justify="center"
+                            )
+                            CONSOLE.print(
+                                f"Please set --rendered-output-name to one of: {all_outputs}", justify="center"
+                            )
+                            sys.exit(1)
+
+                        is_raw = False
+                        is_depth = rendered_output_name.find("depth") != -1
+                        image_name = f"{camera_idx:05d}"
+
+                        # Try to get the original filename
+                        image_name = (
+                            dataparser_outputs.image_filenames[camera_idx].with_suffix("").relative_to(images_root)
+                        )
+
+                        output_path = self.output_path / split / rendered_output_name / image_name
+                        output_path.parent.mkdir(exist_ok=True, parents=True)
+
+                        output_name = rendered_output_name
+                        if output_name.startswith("raw-"):
+                            output_name = output_name[4:]
+                            is_raw = True
+                            if output_name.startswith("gt-"):
+                                output_name = output_name[3:]
+                                output_image = gt_batch[output_name]
+                            else:
+                                output_image = outputs[output_name]
+                                if is_depth:
+                                    # Divide by the dataparser scale factor
+                                    output_image.div_(dataparser_outputs.dataparser_scale)
+                        else:
+                            if output_name.startswith("gt-"):
+                                output_name = output_name[3:]
+                                output_image = gt_batch[output_name]
+                            else:
+                                output_image = outputs[output_name]
+                        del output_name
+
+                        # Map to color spaces / numpy
+                        if is_raw:
+                            output_image = output_image.cpu().numpy()
+                        elif is_depth:
+                            output_image = (
+                                colormaps.apply_depth_colormap(
+                                    output_image,
+                                    accumulation=outputs["accumulation"],
+                                    near_plane=self.depth_near_plane,
+                                    far_plane=self.depth_far_plane,
+                                    colormap_options=self.colormap_options,
+                                )
+                                .cpu()
+                                .numpy()
+                            )
+                        else:
+                            output_image = (
+                                colormaps.apply_colormap(
+                                    image=output_image,
+                                    colormap_options=self.colormap_options,
+                                )
+                                .cpu()
+                                .numpy()
+                            )
+
+                        # Save to file
+                        if is_raw:
+                            with gzip.open(output_path.with_suffix(".npy.gz"), "wb") as f:
+                                np.save(f, output_image)
+                        elif self.image_format == "png":
+                            media.write_image(output_path.with_suffix(".png"), output_image, fmt="png")
+                        elif self.image_format == "jpeg":
+                            media.write_image(
+                                output_path.with_suffix(".jpg"), output_image, fmt="jpeg", quality=self.jpeg_quality
+                            )
+                        else:
+                            raise ValueError(f"Unknown image format {self.image_format}")
+
+        table = Table(
+            title=None,
+            show_header=False,
+            box=box.MINIMAL,
+            title_style=style.Style(bold=True),
+        )
+        for split in self.split.split("+"):
+            table.add_row(f"Outputs {split}", str(self.output_path / split))
+        CONSOLE.print(Panel(table, title="[bold][green]:tada: Render on split {} Complete :tada:[/bold]", expand=False))
+
+
 Commands = tyro.conf.FlagConversionOff[
     Union[
         Annotated[RenderCameraPath, tyro.conf.subcommand(name="camera-path")],
         Annotated[RenderInterpolated, tyro.conf.subcommand(name="interpolate")],
         Annotated[SpiralRender, tyro.conf.subcommand(name="spiral")],
+        Annotated[DatasetRender, tyro.conf.subcommand(name="dataset")],
     ]
 ]
 

nerfstudio/utils/eval_utils.py

@@ -20,7 +20,7 @@ from __future__ import annotations
 import os
 import sys
 from pathlib import Path
-from typing import Literal, Optional, Tuple
+from typing import Literal, Optional, Tuple, Callable
 
 import torch
 import yaml
@@ -69,6 +69,7 @@ def eval_setup(
     config_path: Path,
     eval_num_rays_per_chunk: Optional[int] = None,
     test_mode: Literal["test", "val", "inference"] = "test",
+    update_config_callback: Optional[Callable[[TrainerConfig], TrainerConfig]] = None,
 ) -> Tuple[TrainerConfig, Pipeline, Path, int]:
     """Shared setup for loading a saved pipeline for evaluation.
 
@@ -79,6 +80,7 @@ def eval_setup(
             'val': loads train/val datasets into memory
             'test': loads train/test dataset into memory
             'inference': does not load any dataset into memory
+        update_config_callback: Callback to update the config before loading the pipeline
 
 
     Returns:
@@ -92,6 +94,9 @@ def eval_setup(
     if eval_num_rays_per_chunk:
         config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk
 
+    if update_config_callback is not None:
+        config = update_config_callback(config)
+
     # load checkpoints from wherever they were saved
     # TODO: expose the ability to choose an arbitrary checkpoint
     config.load_dir = config.get_checkpoint_dir()