Render oriented crops in viewer beta (#2447)

* add obb to the rendering pipeline

* useless line

* lint

* remove todo

* pyright

* more pyright

nerfstudio/engine/trainer.py

@@ -170,10 +170,13 @@ class Trainer:
             )
             banner_messages = [f"Viewer at: {self.viewer_state.viewer_url}"]
         if self.config.is_viewer_beta_enabled() and self.local_rank == 0:
+            datapath = self.config.data
+            if datapath is None:
+                datapath = self.base_dir
             self.viewer_state = ViewerBetaState(
                 self.config.viewer,
                 log_filename=viewer_log_path,
-                datapath=self.base_dir,
+                datapath=datapath,
                 pipeline=self.pipeline,
                 trainer=self,
                 train_lock=self.train_lock,

nerfstudio/scripts/render.py

@@ -54,7 +54,7 @@ from nerfstudio.cameras.camera_paths import (
 )
 from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
-from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.model_components import renderers
 from nerfstudio.pipelines.base_pipeline import Pipeline
 from nerfstudio.utils import colormaps, install_checks
@@ -127,12 +127,10 @@ def _render_trajectory_video(
 
         with progress:
             for camera_idx in progress.track(range(cameras.size), description=""):
-                aabb_box = None
+                obb_box = None
                 if crop_data is not None:
-                    bounding_box_min = crop_data.center - crop_data.scale / 2.0
-                    bounding_box_max = crop_data.center + crop_data.scale / 2.0
-                    aabb_box = SceneBox(torch.stack([bounding_box_min, bounding_box_max]).to(pipeline.device))
-                camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx, aabb_box=aabb_box)
+                    obb_box = crop_data.obb
+                camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx, obb_box=obb_box)
 
                 if crop_data is not None:
                     with renderers.background_color_override_context(
@@ -289,10 +287,17 @@ class CropData:
 
     background_color: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0])
     """background color"""
-    center: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0])
-    """center of the crop"""
-    scale: Float[Tensor, "3"] = torch.Tensor([2.0, 2.0, 2.0])
-    """scale of the crop"""
+    obb: OrientedBox = OrientedBox(R=torch.eye(3), T=torch.zeros(3), S=torch.ones(3) * 2)
+    """Oriented box representing the crop region"""
+
+    # properties for backwards-compatibility interface
+    @property
+    def center(self):
+        return self.obb.T
+
+    @property
+    def scale(self):
+        return self.obb.S
 
 
 def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]:
@@ -305,13 +310,13 @@ def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]:
     """
     if "crop" not in camera_json or camera_json["crop"] is None:
         return None
-
     bg_color = camera_json["crop"]["crop_bg_color"]
-
+    center = camera_json["crop"]["crop_center"]
+    scale = camera_json["crop"]["crop_scale"]
+    rot = (0.0, 0.0, 0.0) if "crop_rot" not in camera_json["crop"] else tuple(camera_json["crop"]["crop_rot"])
     return CropData(
         background_color=torch.Tensor([bg_color["r"] / 255.0, bg_color["g"] / 255.0, bg_color["b"] / 255.0]),
-        center=torch.Tensor(camera_json["crop"]["crop_center"]),
-        scale=torch.Tensor(camera_json["crop"]["crop_scale"]),
+        obb=OrientedBox.from_params(center, rot, scale),
     )
 
 

nerfstudio/viewer_beta/render_panel.py

@@ -13,17 +13,19 @@
 # limitations under the License.
 
 from __future__ import annotations
-
+from pathlib import Path
 import colorsys
 import dataclasses
 import threading
 import time
 from typing import Dict, List, Optional, Tuple
-
+import datetime
+from nerfstudio.viewer_beta.control_panel import ControlPanel
 import numpy as onp
 import splines
 import splines.quaternion
 import viser
+import json
 import viser.transforms as tf
 
 
@@ -256,7 +258,11 @@ class CameraPath:
         )
 
 
-def populate_render_tab(server: viser.ViserServer) -> None:
+def populate_render_tab(
+    server: viser.ViserServer, config_path: Path, datapath: Path, control_panel: ControlPanel
+) -> None:
+    from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO
+
     fov_degrees = server.add_gui_slider(
         "FOV",
         initial_value=90.0,
@@ -528,6 +534,11 @@ def populate_render_tab(server: viser.ViserServer) -> None:
         play_button.visible = True
         pause_button.visible = False
 
+    # set the initial value to the current date-time string
+    now = datetime.datetime.now()
+    render_name_text = server.add_gui_text(
+        "Render Name", initial_value=now.strftime("%Y-%m-%d-%H-%M-%S"), hint="Name of the render"
+    )
     render_button = server.add_gui_button(
         "Generate Command",
         color="green",
@@ -537,10 +548,107 @@ def populate_render_tab(server: viser.ViserServer) -> None:
 
     @render_button.on_click
     def _(event: viser.GuiEvent) -> None:
-        """TODO: write the render JSON and show the render command."""
         assert event.client is not None
-        with event.client.add_gui_modal("TODO") as modal:
-            event.client.add_gui_markdown("TODO")
+        num_frames = int(framerate_slider.value * duration_number.value)
+        json_data = {}
+        # json data has the properties:
+        # keyframes: list of keyframes with
+        # matrix : flattened 4x4 matrix
+        # fov: float in degrees
+        # aspect: float
+        # camera_type: string of camera type
+        # render_height: int
+        # render_width: int
+        # fps: int
+        # seconds: float
+        # is_cycle: bool
+        # smoothness_value: float
+        # camera_path: list of frames with properties
+        # camera_to_world: flattened 4x4 matrix
+        # fov: float in degrees
+        # aspect: float
+        # first populate the keyframes:
+        keyframes = []
+        for keyframe, dummy in camera_path._keyframes.values():
+            pose = tf.SE3.from_rotation_and_translation(
+                tf.SO3(keyframe.wxyz) @ tf.SO3.from_x_radians(onp.pi),
+                keyframe.position / VISER_NERFSTUDIO_SCALE_RATIO,
+            )
+            keyframes.append(
+                {
+                    "matrix": pose.as_matrix().flatten().tolist(),
+                    "fov": onp.rad2deg(keyframe.override_fov_value)
+                    if keyframe.override_fov_enabled
+                    else fov_degrees.value,
+                    "aspect": keyframe.aspect,
+                }
+            )
+        json_data["keyframes"] = keyframes
+        json_data["camera_type"] = "perspective"
+        json_data["render_height"] = resolution.value[1]
+        json_data["render_width"] = resolution.value[0]
+        json_data["fps"] = framerate_slider.value
+        json_data["seconds"] = duration_number.value
+        json_data["is_cycle"] = loop.value
+        json_data["smoothness_value"] = smoothness.value
+        # now populate the camera path:
+        camera_path_list = []
+        for i in range(num_frames):
+            maybe_pose_and_fov = camera_path.interpolate_pose_and_fov(i / num_frames)
+            if maybe_pose_and_fov is None:
+                return
+            pose, fov = maybe_pose_and_fov
+            # rotate the axis of the camera 180 about x axis
+            pose = tf.SE3.from_rotation_and_translation(
+                pose.rotation() @ tf.SO3.from_x_radians(onp.pi),
+                pose.translation() / VISER_NERFSTUDIO_SCALE_RATIO,
+            )
+            camera_path_list.append(
+                {
+                    "camera_to_world": pose.as_matrix().flatten().tolist(),
+                    "fov": onp.rad2deg(fov),
+                    "aspect": resolution.value[0] / resolution.value[1],
+                }
+            )
+        json_data["camera_path"] = camera_path_list
+        # finally add crop data if crop is enabled
+        if control_panel.crop_viewport:
+            obb = control_panel.crop_obb
+            rpy = tf.SO3.from_matrix(obb.R.numpy()).as_rpy_radians()
+            color = control_panel.background_color
+            json_data["crop"] = {
+                "crop_center": obb.T.tolist(),
+                "crop_scale": obb.S.tolist(),
+                "crop_rot": [rpy.roll, rpy.pitch, rpy.yaw],
+                "crop_bg_color": {"r": color[0], "g": color[1], "b": color[2]},
+            }
+
+        # now write the json file
+        json_outfile = datapath / "camera_paths" / f"{render_name_text.value}.json"
+        with open(json_outfile.absolute(), "w") as outfile:
+            json.dump(json_data, outfile)
+        # now show the command
+        with event.client.add_gui_modal("Render Command") as modal:
+            dataname = datapath.name
+            command = " ".join(
+                [
+                    "ns-render camera-path",
+                    f"--load-config {config_path}",
+                    f"--camera-path-filename {json_outfile.absolute()}",
+                    f"--output-path renders/{dataname}/{render_name_text.value}.mp4",
+                ]
+            )
+            event.client.add_gui_markdown(
+                "\n".join(
+                    [
+                        "To render the trajectory, run the following from the command line:",
+                        "",
+                        "```",
+                        command,
+                        "```",
+                    ]
+                )
+            )
             close_button = event.client.add_gui_button("Close")
 
             @close_button.on_click
@@ -551,18 +659,3 @@ def populate_render_tab(server: viser.ViserServer) -> None:
     camera_path.default_fov = fov_degrees.value / 180.0 * onp.pi
 
     transform_controls: List[viser.SceneNodeHandle] = []
-
-
-def main() -> None:
-    """Launch a GUI with just the render panel, for development purposes."""
-    server = viser.ViserServer()
-    server.configure_theme(dark_mode=True, control_layout="collapsible")
-    server.world_axes.visible = True
-
-    populate_render_tab(server)
-    while True:
-        time.sleep(10.0)
-
-
-if __name__ == "__main__":
-    main()

nerfstudio/viewer_beta/viewer.py

@@ -154,11 +154,11 @@ class Viewer:
                 self._toggle_training_state,
                 self.set_camera_visibility,
             )
+        config_path = self.log_filename.parents[0] / "config.yml"
         with tabs.add_tab("Render", viser.Icon.CAMERA):
-            populate_render_tab(self.viser_server)
+            populate_render_tab(self.viser_server, config_path, self.datapath, self.control_panel)
 
         with tabs.add_tab("Export", viser.Icon.PACKAGE_EXPORT):
-            config_path = self.log_filename.parents[0] / "config.yml"
             populate_export_tab(self.viser_server, self.control_panel, config_path)
 
         def nested_folder_install(folder_labels: List[str], element: ViewerElement):