Split up render commands (#1923)

* Split up render commands

* bump viewer

nerfstudio/cameras/camera_paths.py

@@ -28,7 +28,7 @@ from nerfstudio.viewer.server.utils import three_js_perspective_camera_focal_len
 
 
 def get_interpolated_camera_path(cameras: Cameras, steps: int, order_poses: bool) -> Cameras:
-    """Generate a camera path between two cameras.
+    """Generate a camera path between two cameras. Uses the camera type of the first camera
 
     Args:
         cameras: Cameras object containing intrinsics of all cameras.
@@ -41,7 +41,14 @@ def get_interpolated_camera_path(cameras: Cameras, steps: int, order_poses: bool
     poses = cameras.camera_to_worlds.cpu().numpy()
     poses, Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=steps, order_poses=order_poses)
 
-    cameras = Cameras(fx=Ks[:, 0, 0], fy=Ks[:, 1, 1], cx=Ks[0, 0, 2], cy=Ks[0, 1, 2], camera_to_worlds=poses)
+    cameras = Cameras(
+        fx=Ks[:, 0, 0],
+        fy=Ks[:, 1, 1],
+        cx=Ks[0, 0, 2],
+        cy=Ks[0, 1, 2],
+        camera_type=cameras.camera_type[0],
+        camera_to_worlds=poses,
+    )
     return cameras
 
 

nerfstudio/scripts/render.py

@@ -25,7 +25,7 @@ import sys
 from contextlib import ExitStack
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import mediapy as media
 import numpy as np
@@ -42,7 +42,7 @@ from rich.progress import (
 )
 from rich.table import Table
 from torchtyping import TensorType
-from typing_extensions import assert_never
+from typing_extensions import Annotated
 
 from nerfstudio.cameras.camera_paths import (
     get_interpolated_camera_path,
@@ -67,7 +67,6 @@ def _render_trajectory_video(
     rendered_resolution_scaling_factor: float = 1.0,
     seconds: float = 5.0,
     output_format: Literal["images", "video"] = "video",
-    camera_type: CameraType = CameraType.PERSPECTIVE,
     colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions(),
 ) -> None:
     """Helper function to create a video of the spiral trajectory.
@@ -81,7 +80,6 @@ def _render_trajectory_video(
         rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.
         seconds: Length of output video.
         output_format: How to save output data.
-        camera_type: Camera projection format type.
         colormap_options: Options for colormap.
     """
     CONSOLE.print("[bold green]Creating trajectory " + output_format)
@@ -172,7 +170,8 @@ def _render_trajectory_video(
         title_style=style.Style(bold=True),
     )
     if output_format == "video":
-        if camera_type == CameraType.EQUIRECTANGULAR:
+        if cameras.camera_type[0] == CameraType.EQUIRECTANGULAR.value:
+            CONSOLE.print("Adding spherical camera data")
             insert_spherical_metadata_into_file(output_filename)
         table.add_row("Video", str(output_filename))
     else:
@@ -283,91 +282,47 @@ def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]:
 
 
 @dataclass
-class RenderTrajectory:
-    """Load a checkpoint, render a trajectory, and save to a video file.
-    The following trajectory options are available,
-    filename: Load from trajectory created using viewer or blender vfx plugin.
-    interpolate: Create trajectory by interpolating between eval dataset images.
-    spiral: Create a spiral trajectory (can be hit or miss).
-    """
+class BaseRender:
+    """Base class for rendering."""
 
     load_config: Path
     """Path to config YAML file."""
-    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
-    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
-    traj: Literal["spiral", "filename", "interpolate"] = "spiral"
-    """Trajectory type to render. Select between spiral-shaped trajectory, trajectory loaded from
-    a viewer-generated file and interpolated camera paths from the eval dataset."""
-    downscale_factor: int = 1
-    """Scaling factor to apply to the camera image resolution."""
-    camera_path_filename: Path = Path("camera_path.json")
-    """Filename of the camera path to render."""
     output_path: Path = Path("renders/output.mp4")
-    """Name of the output file."""
-    seconds: float = 5.0
-    """How long the video should be."""
-    output_format: Literal["images", "video"] = "video"
-    """How to save output data."""
-    interpolation_steps: int = 10
-    """Number of interpolation steps between eval dataset cameras."""
-    order_poses: bool = False
-    """Whether to order camera poses by proximity (only for interpolate trajectory)."""
-    adjust_frame_rate: bool = False
-    """Whether to adjust the frame rate of the output video (only for interpolate trajectory)."""
-    frame_rate: int = 24
-    """Frame rate of the output video (only for interpolate trajectory)."""
-    seconds_between_poses: int = 3
-    """Seconds between poses for interpolated camera paths."""
+    """Path to output video file."""
+    downscale_factor: float = 1.0
+    """Scaling factor to apply to the camera image resolution."""
     eval_num_rays_per_chunk: Optional[int] = None
     """Specifies number of rays per chunk during eval."""
     colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions()
     """Colormap options."""
 
+
+@dataclass
+class RenderCameraPath(BaseRender):
+    """Render a camera path generated by the viewer or blender add-on."""
+
+    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
+    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
+    camera_path_filename: Path = Path("camera_path.json")
+    """Filename of the camera path to render."""
+    output_format: Literal["images", "video"] = "video"
+    """How to save output data."""
+
     def main(self) -> None:
         """Main function."""
         _, pipeline, _, _ = eval_setup(
             self.load_config,
             eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
-            test_mode="test" if self.traj in ["spiral", "interpolate"] else "inference",
+            test_mode="inference",
         )
 
         install_checks.check_ffmpeg_installed()
 
-        seconds = self.seconds
-        crop_data = None
-
-        # TODO(ethan): use camera information from parsing args
-        if self.traj == "spiral":
-            camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten()
-            # TODO(ethan): pass in the up direction of the camera
-            camera_type = CameraType.PERSPECTIVE
-            camera_path = get_spiral_path(camera_start, steps=30, radius=0.1)
-        elif self.traj == "filename":
-            with open(self.camera_path_filename, "r", encoding="utf-8") as f:
-                camera_path = json.load(f)
-            seconds = camera_path["seconds"]
-            if "camera_type" not in camera_path:
-                camera_type = CameraType.PERSPECTIVE
-            elif camera_path["camera_type"] == "fisheye":
-                camera_type = CameraType.FISHEYE
-            elif camera_path["camera_type"] == "equirectangular":
-                camera_type = CameraType.EQUIRECTANGULAR
-            else:
-                camera_type = CameraType.PERSPECTIVE
-            crop_data = get_crop_from_json(camera_path)
-            camera_path = get_path_from_json(camera_path)
-        elif self.traj == "interpolate":
-            camera_type = CameraType.PERSPECTIVE
-            if self.adjust_frame_rate:
-                self.interpolation_steps = self.seconds_between_poses * self.frame_rate
-                seconds = self.seconds_between_poses * len(pipeline.datamanager.eval_dataloader.cameras)
-            camera_path = get_interpolated_camera_path(
-                cameras=pipeline.datamanager.eval_dataloader.cameras,
-                steps=self.interpolation_steps,
-                order_poses=self.order_poses,
-            )
-        else:
-            assert_never(self.traj)
+        with open(self.camera_path_filename, "r", encoding="utf-8") as f:
+            camera_path = json.load(f)
+        seconds = camera_path["seconds"]
+        crop_data = get_crop_from_json(camera_path)
+        camera_path = get_path_from_json(camera_path)
 
         _render_trajectory_video(
             pipeline,
@@ -378,19 +333,119 @@ class RenderTrajectory:
             crop_data=crop_data,
             seconds=seconds,
             output_format=self.output_format,
-            camera_type=camera_type,
             colormap_options=self.colormap_options,
         )
 
 
+@dataclass
+class RenderInterpolated(BaseRender):
+    """Render a trajectory that interpolates between training or eval dataset images."""
+
+    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
+    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
+    pose_source: Literal["eval", "train"] = "eval"
+    """Pose source to render."""
+    interpolation_steps: int = 10
+    """Number of interpolation steps between eval dataset cameras."""
+    order_poses: bool = False
+    """Whether to order camera poses by proximity."""
+    frame_rate: int = 24
+    """Frame rate of the output video."""
+    output_format: Literal["images", "video"] = "video"
+    """How to save output data."""
+
+    def main(self) -> None:
+        """Main function."""
+        _, pipeline, _, _ = eval_setup(
+            self.load_config,
+            eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
+            test_mode="test",
+        )
+
+        install_checks.check_ffmpeg_installed()
+
+        if self.pose_source == "eval":
+            cameras = pipeline.datamanager.eval_dataset.cameras
+        else:
+            cameras = pipeline.datamanager.train_dataset.cameras
+
+        seconds = self.interpolation_steps * len(cameras) / self.frame_rate
+        camera_path = get_interpolated_camera_path(
+            cameras=cameras,
+            steps=self.interpolation_steps,
+            order_poses=self.order_poses,
+        )
+
+        _render_trajectory_video(
+            pipeline,
+            camera_path,
+            output_filename=self.output_path,
+            rendered_output_names=self.rendered_output_names,
+            rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
+            seconds=seconds,
+            output_format=self.output_format,
+            colormap_options=self.colormap_options,
+        )
+
+
+@dataclass
+class SpiralRender(BaseRender):
+    """Render a spiral trajectory (often not great)."""
+
+    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
+    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
+    seconds: float = 3.0
+    """How long the video should be."""
+    output_format: Literal["images", "video"] = "video"
+    """How to save output data."""
+    frame_rate: int = 24
+    """Frame rate of the output video (only for interpolate trajectory)."""
+    radius: float = 0.1
+    """Radius of the spiral."""
+
+    def main(self) -> None:
+        """Main function."""
+        _, pipeline, _, _ = eval_setup(
+            self.load_config,
+            eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
+            test_mode="test",
+        )
+
+        install_checks.check_ffmpeg_installed()
+
+        steps = int(self.frame_rate * self.seconds)
+        camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten()
+        camera_path = get_spiral_path(camera_start, steps=steps, radius=self.radius)
+
+        _render_trajectory_video(
+            pipeline,
+            camera_path,
+            output_filename=self.output_path,
+            rendered_output_names=self.rendered_output_names,
+            rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
+            seconds=self.seconds,
+            output_format=self.output_format,
+            colormap_options=self.colormap_options,
+        )
+
+
+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")],
+    ]
+]
+
+
 def entrypoint():
     """Entrypoint for use with pyproject scripts."""
     tyro.extras.set_accent_color("bright_yellow")
-    tyro.cli(RenderTrajectory).main()
+    tyro.cli(Commands).main()
 
 
 if __name__ == "__main__":
     entrypoint()
 
 # For sphinx docs
-get_parser_fn = lambda: tyro.extras.get_parser(RenderTrajectory)  # noqa
+get_parser_fn = lambda: tyro.extras.get_parser(Commands)  # noqa

nerfstudio/utils/colormaps.py

@@ -26,7 +26,7 @@ from nerfstudio.utils import colors
 Colormaps = Literal["default", "turbo", "viridis", "magma", "inferno", "cividis"]
 
 
-@dataclass
+@dataclass(frozen=True)
 class ColormapOptions:
     """Options for colormap"""
 

nerfstudio/viewer/app/package.json

@@ -1,7 +1,7 @@
 {
   "name": "viewer",
   "homepage": ".",
-  "version": "23-05-01-0",
+  "version": "23-05-15-0",
   "private": true,
   "dependencies": {
     "@emotion/react": "^11.10.4",
@@ -93,4 +93,4 @@
     "resolve-url-loader": "^5.0.0",
     "yaml": "^2.2.2"
   }
-}
+}

nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx

@@ -35,7 +35,7 @@ export default function RenderModal(props: RenderModalProps) {
   const config_filename = `${config_base_dir}/config.yml`;
   const camera_path_filename = `${export_path}.json`;
   const data_base_dir_leaf = data_base_dir.split('/').pop();
-  const cmd = `ns-render --load-config ${config_filename} --traj filename --camera-path-filename ${data_base_dir}/camera_paths/${camera_path_filename} --output-path renders/${data_base_dir_leaf}/${export_path}.mp4`;
+  const cmd = `ns-render camera-path --load-config ${config_filename} --camera-path-filename ${data_base_dir}/camera_paths/${camera_path_filename} --output-path renders/${data_base_dir_leaf}/${export_path}.mp4`;
 
   const text_intro = `To render a full resolution video, run the following command in a terminal.`;