Improved render panel for beta viewer (#2667)

* Continuous velocities for render splines

* Configure ruff line length

* Improved render preview

* No longer need to update camera FOV in render preview

* Details

* Polishing

* Cleanup

* Details

* Address pyright errors

* type: ignore for broken (?) Aria stub

nerfstudio/scripts/datasets/process_project_aria.py

@@ -103,7 +103,7 @@ def get_camera_calibs(provider: VrsDataProvider) -> Dict[str, AriaCameraCalibrat
 
 
 def read_trajectory_csv_to_dict(file_iterable_csv: str) -> TimedPoses:
-    closed_loop_traj = mps.read_closed_loop_trajectory(file_iterable_csv)
+    closed_loop_traj = mps.read_closed_loop_trajectory(file_iterable_csv)  # type: ignore
 
     timestamps_secs, poses = zip(
         *[(it.tracking_timestamp.total_seconds(), it.transform_world_device) for it in closed_loop_traj]

nerfstudio/viewer_beta/control_panel.py

@@ -18,9 +18,8 @@ from typing import Callable, DefaultDict, List, Tuple, get_args
 
 import numpy as np
 import torch
-import viser.transforms as vtf
-from viser import ViserServer
 import viser
+import viser.transforms as vtf
 from nerfstudio.data.scene_box import OrientedBox
 from nerfstudio.utils.colormaps import ColormapOptions, Colormaps
 from nerfstudio.viewer_beta.viewer_elements import (  # ViewerButtonGroup,
@@ -33,6 +32,7 @@ from nerfstudio.viewer_beta.viewer_elements import (  # ViewerButtonGroup,
     ViewerSlider,
     ViewerVec3,
 )
+from viser import ViserServer
 
 
 class ControlPanel:
@@ -42,7 +42,6 @@ class ControlPanel:
         time_enabled: whether or not the time slider should be enabled
         rerender_cb: a callback that will be called when the user changes a parameter that requires a rerender
             (eg train speed, max res, etc)
-        crop_update_cb: a callback that will be called when the user changes the crop parameters
         update_output_cb: a callback that will be called when the user changes the output render
         default_composite_depth: whether to default to compositing depth or not
     """
@@ -52,8 +51,7 @@ class ControlPanel:
         viser_server: ViserServer,
         time_enabled: bool,
         scale_ratio: float,
-        rerender_cb: Callable,
-        crop_update_cb: Callable,
+        rerender_cb: Callable[[], None],
         update_output_cb: Callable,
         update_split_output_cb: Callable,
         default_composite_depth: bool = True,
@@ -71,51 +69,53 @@ class ControlPanel:
             cb_hook=lambda han: self._train_speed_cb(),
         )
         self._output_render = ViewerDropdown(
-            "Output Render",
+            "Output type",
             "not set",
             ["not set"],
-            cb_hook=lambda han: [self.update_control_panel(), update_output_cb(han), rerender_cb(han)],
+            cb_hook=lambda han: [self.update_control_panel(), update_output_cb(han), rerender_cb()],
             hint="The output to render",
         )
         self._colormap = ViewerDropdown[Colormaps](
-            "Colormap", "default", ["default"], cb_hook=rerender_cb, hint="The colormap to use"
+            "Colormap", "default", ["default"], cb_hook=lambda _: rerender_cb(), hint="The colormap to use"
         )
-        self._invert = ViewerCheckbox("Invert", False, cb_hook=rerender_cb, hint="Invert the colormap")
-        self._normalize = ViewerCheckbox("Normalize", True, cb_hook=rerender_cb, hint="Normalize the colormap")
-        self._min = ViewerNumber("Min", 0.0, cb_hook=rerender_cb, hint="Min value of the colormap")
-        self._max = ViewerNumber("Max", 1.0, cb_hook=rerender_cb, hint="Max value of the colormap")
+        self._invert = ViewerCheckbox("Invert", False, cb_hook=lambda _: rerender_cb(), hint="Invert the colormap")
+        self._normalize = ViewerCheckbox(
+            "Normalize", True, cb_hook=lambda _: rerender_cb(), hint="Normalize the colormap"
+        )
+        self._min = ViewerNumber("Min", 0.0, cb_hook=lambda _: rerender_cb(), hint="Min value of the colormap")
+        self._max = ViewerNumber("Max", 1.0, cb_hook=lambda _: rerender_cb(), hint="Max value of the colormap")
 
         self._split = ViewerCheckbox(
             "Enable",
             False,
-            cb_hook=lambda han: [self.update_control_panel(), rerender_cb(han)],
+            cb_hook=lambda han: [self.update_control_panel(), rerender_cb()],
             hint="Render two outputs",
         )
         self._split_percentage = ViewerSlider(
-            "Split Percentage", 0.5, 0.0, 1.0, 0.01, cb_hook=rerender_cb, hint="Where to split"
+            "Split percentage", 0.5, 0.0, 1.0, 0.01, cb_hook=lambda _: rerender_cb(), hint="Where to split"
         )
         self._split_output_render = ViewerDropdown(
-            "Output Render Split",
+            "Output render split",
             "not set",
             ["not set"],
-            cb_hook=lambda han: [self.update_control_panel(), update_split_output_cb(han), rerender_cb(han)],
+            cb_hook=lambda han: [self.update_control_panel(), update_split_output_cb(han), rerender_cb()],
             hint="The second output",
         )
         # Hack: spaces are after at the end of the names to make them unique
         self._split_colormap = ViewerDropdown[Colormaps](
-            "Colormap ", "default", ["default"], cb_hook=rerender_cb, hint="Colormap of the second output"
+            "Colormap ", "default", ["default"], cb_hook=lambda _: rerender_cb(), hint="Colormap of the second output"
         )
         self._split_invert = ViewerCheckbox(
-            "Invert ", False, cb_hook=rerender_cb, hint="Invert the colormap of the second output"
+            "Invert ", False, cb_hook=lambda _: rerender_cb(), hint="Invert the colormap of the second output"
         )
         self._split_normalize = ViewerCheckbox(
-            "Normalize ", True, cb_hook=rerender_cb, hint="Normalize the colormap of the second output"
+            "Normalize ", True, cb_hook=lambda _: rerender_cb(), hint="Normalize the colormap of the second output"
         )
         self._split_min = ViewerNumber(
-            "Min ", 0.0, cb_hook=rerender_cb, hint="Min value of the colormap of the second output"
+            "Min ", 0.0, cb_hook=lambda _: rerender_cb(), hint="Min value of the colormap of the second output"
         )
         self._split_max = ViewerNumber(
-            "Max ", 1.0, cb_hook=rerender_cb, hint="Max value of the colormap of the second output"
+            "Max ", 1.0, cb_hook=lambda _: rerender_cb(), hint="Max value of the colormap of the second output"
         )
 
         self._train_util = ViewerSlider(
@@ -127,22 +127,28 @@ class ControlPanel:
             hint="Target training utilization, 0.0 is slow, 1.0 is fast. Doesn't affect final render quality",
         )
         self._layer_depth = ViewerCheckbox(
-            "Composite Depth",
+            "Composite depth",
             self.default_composite_depth,
-            cb_hook=rerender_cb,
+            cb_hook=lambda _: rerender_cb(),
             hint="Allow NeRF to occlude 3D browser objects",
         )
         self._max_res = ViewerSlider(
-            "Max Res", 512, 64, 2048, 100, cb_hook=rerender_cb, hint="Maximum resolution to render in viewport"
+            "Max res",
+            512,
+            64,
+            2048,
+            100,
+            cb_hook=lambda _: rerender_cb(),
+            hint="Maximum resolution to render in viewport",
         )
         self._crop_viewport = ViewerCheckbox(
             "Enable ",
             False,
-            cb_hook=lambda han: [self.update_control_panel(), crop_update_cb(han), rerender_cb(han)],
+            cb_hook=lambda han: [self.update_control_panel(), rerender_cb()],
             hint="Crop the scene to a specified box",
         )
         self._background_color = ViewerRGB(
-            "Background color", (38, 42, 55), cb_hook=crop_update_cb, hint="Color of the background"
+            "Background color", (38, 42, 55), cb_hook=lambda _: rerender_cb(), hint="Color of the background"
         )
         self._crop_handle = self.viser_server.add_transform_controls("Crop", depth_test=False, line_width=4.0)
 
@@ -150,10 +156,10 @@ class ControlPanel:
             self._crop_handle.position = tuple(p * self.viser_scale_ratio for p in han.value)  # type: ignore
 
         self._crop_center = ViewerVec3(
-            "Crop Center",
+            "Crop center",
             (0.0, 0.0, 0.0),
             step=0.01,
-            cb_hook=lambda e: [crop_update_cb(e), update_center(e)],
+            cb_hook=lambda e: [rerender_cb(), update_center(e)],
             hint="Center of the crop box",
         )
 
@@ -161,15 +167,15 @@ class ControlPanel:
             self._crop_handle.wxyz = vtf.SO3.from_rpy_radians(*han.value).wxyz
 
         self._crop_rot = ViewerVec3(
-            "Crop Rotation",
+            "Crop rotation",
             (0.0, 0.0, 0.0),
             step=0.01,
-            cb_hook=lambda e: [crop_update_cb(e), update_rot(e)],
+            cb_hook=lambda e: [rerender_cb(), update_rot(e)],
             hint="Rotation of the crop box",
         )
 
         self._crop_scale = ViewerVec3(
-            "Crop Scale", (1.0, 1.0, 1.0), step=0.01, cb_hook=crop_update_cb, hint="Scale of the crop box"
+            "Crop scale", (1.0, 1.0, 1.0), step=0.01, cb_hook=lambda _: rerender_cb(), hint="Size of the crop box."
         )
 
         @self._crop_handle.on_update
@@ -179,7 +185,7 @@ class ControlPanel:
             rpy = vtf.SO3(self._crop_handle.wxyz).as_rpy_radians()
             self._crop_rot.value = (float(rpy.roll), float(rpy.pitch), float(rpy.yaw))
 
-        self._time = ViewerSlider("Time", 0.0, 0.0, 1.0, 0.01, cb_hook=rerender_cb, hint="Time to render")
+        self._time = ViewerSlider("Time", 0.0, 0.0, 1.0, 0.01, cb_hook=lambda _: rerender_cb(), hint="Time to render")
         self._time_enabled = time_enabled
 
         self.add_element(self._train_speed)
@@ -219,7 +225,10 @@ class ControlPanel:
 
         self.add_element(self._time, additional_tags=("time",))
         self._reset_camera = viser_server.add_gui_button(
-            label="Reset Up Dir", disabled=False, icon=viser.Icon.ARROW_BIG_UP_LINES, color="gray"
+            label="Reset Up Direction",
+            icon=viser.Icon.ARROW_BIG_UP_LINES,
+            color="gray",
+            hint="Set the up direction of the camera orbit controls to the camera's current up direction.",
         )
         self._reset_camera.on_click(self._reset_camera_cb)
 

nerfstudio/viewer_beta/render_panel.py

@@ -21,14 +21,14 @@ import json
 import threading
 import time
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Literal, Optional, Tuple
 
 import numpy as np
+import scipy.interpolate
 import splines
 import splines.quaternion
 import viser
 import viser.transforms as tf
-
 from nerfstudio.viewer_beta.control_panel import ControlPanel
 
 
@@ -209,21 +209,50 @@ class CameraPath:
             frame = dataclasses.replace(frame[0], aspect=aspect)
             self.add_camera(frame, keyframe_index=keyframe_index)
 
+    def get_aspect(self) -> float:
+        """Get W/H aspect ratio, which is shared across all keyframes."""
+        assert len(self._keyframes) > 0
+        return next(iter(self._keyframes.values()))[0].aspect
+
     def reset(self) -> None:
         for frame in self._keyframes.values():
             frame[1].remove()
         self._keyframes.clear()
         self.update_spline()
 
-    def interpolate_pose_and_fov_rad(self, normalized_t: float) -> Optional[Tuple[tf.SE3, float]]:
-        if len(self._keyframes) < 2:
-            return None
+    def spline_t_from_t_sec(self, time: np.ndarray) -> np.ndarray:
+        """From a time value in seconds, compute a t value for our geometric
+        spline interpolation. An increment of 1 for the latter will move the
+        camera forward by one keyframe.
 
+        We use a PCHIP spline here to guarantee monotonicity.
+        """
         transition_times_cumsum = self.compute_transition_times_cumsum()
         spline_indices = np.arange(transition_times_cumsum.shape[0])
 
-        def spline_t_from_t_sec(time: Union[float, np.ndarray]) -> np.ndarray:
-            return np.interp(time, transition_times_cumsum, spline_indices)
+        if self.loop:
+            # In the case of a loop, we pad the spline to match the start/end
+            # slopes.
+            interpolator = scipy.interpolate.PchipInterpolator(
+                x=np.concatenate(
+                    [
+                        [-(transition_times_cumsum[-1] - transition_times_cumsum[-2])],
+                        transition_times_cumsum,
+                        transition_times_cumsum[-1:] + transition_times_cumsum[1:2],
+                    ],
+                    axis=0,
+                ),
+                y=np.concatenate([[-1], spline_indices, [spline_indices[-1] + 1]], axis=0),
+            )
+        else:
+            interpolator = scipy.interpolate.PchipInterpolator(x=transition_times_cumsum, y=spline_indices)
+
+        # Clip to account for floating point error.
+        return np.clip(interpolator(time), 0, spline_indices[-1])
+
+    def interpolate_pose_and_fov_rad(self, normalized_t: float) -> Optional[Tuple[tf.SE3, float]]:
+        if len(self._keyframes) < 2:
+            return None
 
         self._fov_spline = splines.KochanekBartels(
             [
@@ -239,32 +268,29 @@ class CameraPath:
         assert self._fov_spline is not None
         max_t = self.compute_duration()
         t = max_t * normalized_t
+        spline_t = float(self.spline_t_from_t_sec(np.array(t)))
 
-        quat = self._orientation_spline.evaluate(spline_t_from_t_sec(t))
+        quat = self._orientation_spline.evaluate(spline_t)
         assert isinstance(quat, splines.quaternion.UnitQuaternion)
         return (
             tf.SE3.from_rotation_and_translation(
                 tf.SO3(np.array([quat.scalar, *quat.vector])),
-                self._position_spline.evaluate(spline_t_from_t_sec(t)),
+                self._position_spline.evaluate(spline_t),
             ),
-            float(self._fov_spline.evaluate(spline_t_from_t_sec(t))),
+            float(self._fov_spline.evaluate(spline_t)),
         )
 
     def update_spline(self) -> None:
         num_frames = int(self.compute_duration() * self.framerate)
-        if num_frames <= 0 or not self.show_spline:
+        keyframes = list(self._keyframes.values())
+
+        if num_frames <= 0 or not self.show_spline or len(keyframes) < 2:
             for node in self._spline_nodes:
                 node.remove()
             self._spline_nodes.clear()
             return
 
-        # Update internal splines.
-        keyframes = list(self._keyframes.values())
         transition_times_cumsum = self.compute_transition_times_cumsum()
-        spline_indices = np.arange(transition_times_cumsum.shape[0])
-
-        def spline_t_from_t_sec(time: Union[float, np.ndarray]) -> np.ndarray:
-            return np.interp(time, transition_times_cumsum, spline_indices)
 
         self._orientation_spline = splines.quaternion.KochanekBartels(
             [
@@ -282,7 +308,7 @@ class CameraPath:
 
         # Update visualized spline.
         points_array = self._position_spline.evaluate(
-            spline_t_from_t_sec(np.linspace(0, transition_times_cumsum[-1], num_frames))
+            self.spline_t_from_t_sec(np.linspace(0, transition_times_cumsum[-1], num_frames))
         )
         colors_array = np.array([colorsys.hls_to_rgb(h, 0.5, 1.0) for h in np.linspace(0.0, 1.0, len(points_array))])
 
@@ -313,7 +339,11 @@ class CameraPath:
         def make_transition_handle(i: int) -> None:
             assert self._position_spline is not None
             transition_pos = self._position_spline.evaluate(
-                spline_t_from_t_sec((transition_times_cumsum[i] + transition_times_cumsum[i + 1]) / 2.0)
+                float(
+                    self.spline_t_from_t_sec(
+                        (transition_times_cumsum[i] + transition_times_cumsum[i + 1]) / 2.0,
+                    )
+                )
             )
             transition_sphere = self._server.add_icosphere(
                 f"/render_camera_spline/transition_{i}",
@@ -340,7 +370,8 @@ class CameraPath:
                 ) as camera_edit_panel:
                     self._camera_edit_panel = camera_edit_panel
                     override_transition_enabled = server.add_gui_checkbox(
-                        "Override transition", initial_value=keyframe.override_transition_enabled
+                        "Override transition",
+                        initial_value=keyframe.override_transition_enabled,
                     )
                     override_transition_sec = server.add_gui_number(
                         "Override transition (sec)",
@@ -418,14 +449,31 @@ class CameraPath:
         return np.array(out)
 
 
+@dataclasses.dataclass
+class RenderTabState:
+    """Useful GUI handles exposed by the render tab."""
+
+    preview_render: bool
+    preview_fov: float
+    preview_aspect: float
+    preview_camera_type: Literal["Perspective", "Fisheye", "Equirectangular"]
+
+
 def populate_render_tab(
     server: viser.ViserServer,
     config_path: Path,
     datapath: Path,
     control_panel: Optional[ControlPanel] = None,
-) -> None:
+) -> RenderTabState:
     from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO
 
+    render_tab_state = RenderTabState(
+        preview_render=False,
+        preview_fov=0.0,
+        preview_aspect=1.0,
+        preview_camera_type="Perspective",
+    )
+
     fov_degrees = server.add_gui_slider(
         "Default FOV",
         initial_value=75.0,
@@ -445,6 +493,7 @@ def populate_render_tab(
         # Updating the aspect ratio will also re-render the camera frustums.
         # Could rethink this.
         camera_path.update_aspect(resolution.value[0] / resolution.value[1])
+        compute_and_update_preview_camera_state()
 
     resolution = server.add_gui_vector2(
         "Resolution",
@@ -454,16 +503,20 @@ def populate_render_tab(
         step=1,
         hint="Render output resolution in pixels.",
     )
-    resolution.on_update(lambda _: camera_path.update_aspect(resolution.value[0] / resolution.value[1]))
+
+    @resolution.on_update
+    def _(_) -> None:
+        camera_path.update_aspect(resolution.value[0] / resolution.value[1])
+        compute_and_update_preview_camera_state()
 
     camera_type = server.add_gui_dropdown(
-        "Camera Type",
+        "Camera type",
         ("Perspective", "Fisheye", "Equirectangular"),
         initial_value="Perspective",
         hint="Camera model to render with. This is applied to all keyframes.",
     )
     add_button = server.add_gui_button(
-        "Add keyframe",
+        "Add Keyframe",
         icon=viser.Icon.PLUS,
         hint="Add a new keyframe at the current pose.",
     )
@@ -483,19 +536,8 @@ def populate_render_tab(
         duration_number.value = camera_path.compute_duration()
         camera_path.update_spline()
 
-    reset_up_button = server.add_gui_button(
-        "Reset up direction",
-        icon=viser.Icon.ARROW_BIG_UP_LINES,
-        hint="Reset the orbit up direction.",
-    )
-
-    @reset_up_button.on_click
-    def _(event: viser.GuiEvent) -> None:
-        assert event.client is not None
-        event.client.camera.up_direction = tf.SO3(event.client.camera.wxyz) @ np.array([0.0, -1.0, 0.0])
-
     clear_keyframes_button = server.add_gui_button(
-        "Clear keyframes",
+        "Clear Keyframes",
         icon=viser.Icon.TRASH,
         hint="Remove all keyframes from the render path.",
     )
@@ -535,7 +577,7 @@ def populate_render_tab(
         duration_number.value = camera_path.compute_duration()
 
     tension_slider = server.add_gui_slider(
-        "Spline Tension",
+        "Spline tension",
         min=0.0,
         max=1.0,
         initial_value=0.0,
@@ -554,6 +596,8 @@ def populate_render_tab(
         hint="Toggle move handles for keyframes in the scene.",
     )
 
+    transform_controls: List[viser.SceneNodeHandle] = []
+
     @move_checkbox.on_update
     def _(event: viser.GuiEvent) -> None:
         # Clear move handles when toggled off.
@@ -614,7 +658,11 @@ def populate_render_tab(
     with playback_folder:
         play_button = server.add_gui_button("Play", icon=viser.Icon.PLAYER_PLAY)
         pause_button = server.add_gui_button("Pause", icon=viser.Icon.PLAYER_PAUSE, visible=False)
-        attach_viewport_checkbox = server.add_gui_checkbox("Attach viewport", initial_value=False)
+        preview_render_button = server.add_gui_button(
+            "Preview Render", hint="Show a preview of the render in the viewport."
+        )
+        preview_render_stop_button = server.add_gui_button("Exit Render Preview", color="red", visible=False)
+
         transition_sec_number = server.add_gui_number(
             "Transition (sec)",
             min=0.001,
@@ -654,6 +702,24 @@ def populate_render_tab(
             preview_camera_handle.remove()
             preview_camera_handle = None
 
+    def compute_and_update_preview_camera_state() -> Optional[Tuple[tf.SE3, float]]:
+        """Update the render tab state with the current preview camera pose.
+        Returns current camera pose + FOV if available."""
+
+        if preview_frame_slider is None:
+            return
+        maybe_pose_and_fov_rad = camera_path.interpolate_pose_and_fov_rad(
+            preview_frame_slider.value / get_max_frame_index()
+        )
+        if maybe_pose_and_fov_rad is None:
+            remove_preview_camera()
+            return
+        pose, fov_rad = maybe_pose_and_fov_rad
+        render_tab_state.preview_fov = fov_rad
+        render_tab_state.preview_aspect = camera_path.get_aspect()
+        render_tab_state.preview_camera_type = camera_type.value
+        return pose, fov_rad
+
     def add_preview_frame_slider() -> Optional[viser.GuiInputHandle[int]]:
         """Helper for creating the current frame # slider. This is removed and
         re-added anytime the `max` value changes."""
@@ -666,18 +732,17 @@ def populate_render_tab(
                 step=1,
                 initial_value=0,
                 # Place right after the pause button.
-                order=pause_button.order + 0.01,
+                order=preview_render_stop_button.order + 0.01,
+                disabled=get_max_frame_index() == 1,
             )
+            play_button.disabled = preview_frame_slider.disabled
+            preview_render_button.disabled = preview_frame_slider.disabled
 
         @preview_frame_slider.on_update
         def _(_) -> None:
             nonlocal preview_camera_handle
-
-            maybe_pose_and_fov_rad = camera_path.interpolate_pose_and_fov_rad(
-                preview_frame_slider.value / get_max_frame_index()
-            )
+            maybe_pose_and_fov_rad = compute_and_update_preview_camera_state()
             if maybe_pose_and_fov_rad is None:
-                remove_preview_camera()
                 return
             pose, fov_rad = maybe_pose_and_fov_rad
 
@@ -689,57 +754,61 @@ def populate_render_tab(
                 wxyz=pose.rotation().wxyz,
                 position=pose.translation(),
                 color=(10, 200, 30),
-                # Hack: hide green frustum if the viewport is attached.
-                # This is a waste of bandwidth, but will ensure that any old
-                # frustums are removed/aren't rendered.
-                #
-                # Easy to fix with a global variable.
-                visible=not attach_viewport_checkbox.value,
             )
-            if attach_viewport_checkbox.value:
+            if render_tab_state.preview_render:
                 for client in server.get_clients().values():
                     client.camera.wxyz = pose.rotation().wxyz
                     client.camera.position = pose.translation()
-                    client.camera.fov = fov_rad
 
         return preview_frame_slider
 
-    @attach_viewport_checkbox.on_update
+    # We back up the camera poses before and after we start previewing renders.
+    camera_pose_backup_from_id: Dict[int, tuple] = {}
+
+    @preview_render_button.on_click
     def _(_) -> None:
-        if preview_frame_slider is None:
-            remove_preview_camera()
-            return
-        maybe_pose_and_fov_rad = camera_path.interpolate_pose_and_fov_rad(
-            preview_frame_slider.value / get_max_frame_index()
-        )
+        render_tab_state.preview_render = True
+        preview_render_button.visible = False
+        preview_render_stop_button.visible = True
+
+        maybe_pose_and_fov_rad = compute_and_update_preview_camera_state()
         if maybe_pose_and_fov_rad is None:
             remove_preview_camera()
             return
         pose, fov = maybe_pose_and_fov_rad
-        server.add_camera_frustum(
-            "/preview_camera",
-            fov=fov,
-            aspect=resolution.value[0] / resolution.value[1],
-            scale=0.35,
-            wxyz=pose.rotation().wxyz,
-            position=pose.translation(),
-            color=(10, 200, 30),
-            # Hack: hide green frustum if the viewport is attached.
-            # This is a waste of bandwidth, but will ensure that any old
-            # frustums are removed/aren't rendered.
-            #
-            # Easy to fix with a global variable.
-            visible=not attach_viewport_checkbox.value,
-        )
-        if not attach_viewport_checkbox.value:
-            for client in server.get_clients().values():
-                client.camera.fov = fov_degrees.value / 180 * np.pi
-        else:
-            if attach_viewport_checkbox.value:
-                for client in server.get_clients().values():
-                    client.camera.wxyz = pose.rotation().wxyz
-                    client.camera.position = pose.translation()
-                    client.camera.fov = fov
+        del fov
+
+        # Hide all scene nodes when we're previewing the render.
+        server.set_global_scene_node_visibility(False)
+
+        # Back up and then set camera poses.
+        for client in server.get_clients().values():
+            camera_pose_backup_from_id[client.client_id] = (
+                client.camera.position,
+                client.camera.look_at,
+                client.camera.up_direction,
+            )
+            client.camera.wxyz = pose.rotation().wxyz
+            client.camera.position = pose.translation()
+
+    @preview_render_stop_button.on_click
+    def _(_) -> None:
+        render_tab_state.preview_render = False
+        preview_render_button.visible = True
+        preview_render_stop_button.visible = False
+
+        # Revert camera poses.
+        for client in server.get_clients().values():
+            if client.client_id not in camera_pose_backup_from_id:
+                continue
+            cam_position, cam_look_at, cam_up = camera_pose_backup_from_id.pop(client.client_id)
+            client.camera.position = cam_position
+            client.camera.look_at = cam_look_at
+            client.camera.up_direction = cam_up
+            client.flush()
+
+        # Un-hide scene nodes.
+        server.set_global_scene_node_visibility(True)
 
     preview_frame_slider = add_preview_frame_slider()
 
@@ -823,13 +892,16 @@ def populate_render_tab(
                         pose = tf.SE3.from_matrix(np.array(frame["matrix"]).reshape(4, 4))
                         # apply the x rotation by 180 deg
                         pose = tf.SE3.from_rotation_and_translation(
-                            pose.rotation() @ tf.SO3.from_x_radians(np.pi), pose.translation()
+                            pose.rotation() @ tf.SO3.from_x_radians(np.pi),
+                            pose.translation(),
                         )
                         camera_path.add_camera(
                             Keyframe(
                                 position=pose.translation() * VISER_NERFSTUDIO_SCALE_RATIO,
                                 wxyz=pose.rotation().wxyz,
-                                override_fov_enabled=True,
+                                # There are some floating point conversions between degrees and radians, so the fov and
+                                # default_Fov values will not be exactly matched.
+                                override_fov_enabled=abs(frame["fov"] - json_data.get("default_fov", 0.0)) < 1e-3,
                                 override_fov_rad=frame["fov"] / 180.0 * np.pi,
                                 aspect=frame["aspect"],
                                 override_transition_enabled=frame.get("override_transition_enabled", None),
@@ -853,7 +925,9 @@ def populate_render_tab(
     # 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 name",
+        initial_value=now.strftime("%Y-%m-%d-%H-%M-%S"),
+        hint="Name of the render",
     )
     render_button = server.add_gui_button(
         "Generate Command",
@@ -862,6 +936,18 @@ def populate_render_tab(
         hint="Generate the ns-render command for rendering the camera path.",
     )
 
+    reset_up_button = server.add_gui_button(
+        "Reset Up Direction",
+        icon=viser.Icon.ARROW_BIG_UP_LINES,
+        color="gray",
+        hint="Set the up direction of the camera orbit controls to the camera's current up direction.",
+    )
+
+    @reset_up_button.on_click
+    def _(event: viser.GuiEvent) -> None:
+        assert event.client is not None
+        event.client.camera.up_direction = tf.SO3(event.client.camera.wxyz) @ np.array([0.0, -1.0, 0.0])
+
     @render_button.on_click
     def _(event: viser.GuiEvent) -> None:
         assert event.client is not None
@@ -901,6 +987,7 @@ def populate_render_tab(
                     "override_transition_sec": keyframe.override_transition_sec,
                 }
             )
+        json_data["default_fov"] = fov_degrees.value
         json_data["default_transition_sec"] = transition_sec_number.value
         json_data["keyframes"] = keyframes
         json_data["camera_type"] = camera_type.value.lower()
@@ -980,7 +1067,7 @@ def populate_render_tab(
     camera_path.default_fov = fov_degrees.value / 180.0 * np.pi
     camera_path.default_transition_sec = transition_sec_number.value
 
-    transform_controls: List[viser.SceneNodeHandle] = []
+    return render_tab_state
 
 
 if __name__ == "__main__":

nerfstudio/viewer_beta/render_state_machine.py

@@ -20,15 +20,16 @@ import threading
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args
 
+import numpy as np
 import torch
-from viser import ClientHandle
+from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.model_components.renderers import background_color_override_context
+from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
 from nerfstudio.utils import colormaps, writer
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
 from nerfstudio.viewer.server import viewer_utils
 from nerfstudio.viewer_beta.utils import CameraState, get_camera
-from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
-from nerfstudio.cameras.cameras import Cameras
+from viser import ClientHandle
 
 if TYPE_CHECKING:
     from nerfstudio.viewer_beta.viewer import Viewer
@@ -96,7 +97,7 @@ class RenderStateMachine(threading.Thread):
             #  1. we are in low_moving state
             #  2. the current next_action is move, static, or rerender
             return
-        elif self.next_action == "rerender":
+        elif self.next_action.action == "rerender":
             # never overwrite rerenders
             pass
         elif action.action == "static" and self.next_action.action == "move":
@@ -254,9 +255,33 @@ class RenderStateMachine(threading.Thread):
         depth = (
             outputs["gl_z_buf_depth"].cpu().numpy() * self.viser_scale_ratio if "gl_z_buf_depth" in outputs else None
         )
-        jpg_quality = self.viewer.config.jpeg_quality if static_render else 40
+
+        # Convert to numpy.
+        selected_output = selected_output.cpu().numpy()
+        assert selected_output.shape[-1] == 3
+
+        # Pad image if the aspect ratio (W/H) doesn't match the client!
+        current_h, current_w = selected_output.shape[:2]
+        desired_aspect = self.client.camera.aspect
+        pad_width = int(max(0, (desired_aspect * current_h - current_w) // 2))
+        pad_height = int(max(0, (current_w / desired_aspect - current_h) // 2))
+        if pad_width > 5 or pad_height > 5:
+            selected_output = np.pad(
+                selected_output,
+                ((pad_height, pad_height), (pad_width, pad_width), (0, 0)),
+                mode="constant",
+                constant_values=0,
+            )
+
+        jpg_quality = (
+            self.viewer.config.jpeg_quality
+            if static_render
+            else 75
+            if self.viewer.render_tab_state.preview_render
+            else 40
+        )
         self.client.set_background_image(
-            selected_output.cpu().numpy(),
+            selected_output,
             format=self.viewer.config.image_format,
             jpeg_quality=jpg_quality,
             depth=depth,

nerfstudio/viewer_beta/utils.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -31,11 +31,13 @@ class CameraState:
     """A dataclass for storing the camera state."""
 
     fov: float
-    """ The field of view of the camera. """
+    """The field of view of the camera."""
     aspect: float
-    """ The aspect ratio of the image. """
+    """The aspect ratio of the image. """
     c2w: Float[torch.Tensor, "3 4"]
-    """ The camera matrix. """
+    """The camera matrix."""
+    camera_type: Literal[CameraType.PERSPECTIVE, CameraType.EQUIRECTANGULAR, CameraType.FISHEYE]
+    """Type of camera to render."""
 
 
 def get_camera(
@@ -57,14 +59,19 @@ def get_camera(
     focal_length = pp_h / np.tan(fov / 2.0)
     intrinsics_matrix = torch.tensor([[focal_length, 0, pp_w], [0, focal_length, pp_h], [0, 0, 1]], dtype=torch.float32)
 
-    camera_type = CameraType.PERSPECTIVE
+    if camera_state.camera_type is CameraType.EQUIRECTANGULAR:
+        fx = float(image_width / 2)
+        fy = float(image_height)
+    else:
+        fx = intrinsics_matrix[0, 0]
+        fy = intrinsics_matrix[1, 1]
 
     camera = Cameras(
-        fx=intrinsics_matrix[0, 0],
-        fy=intrinsics_matrix[1, 1],
+        fx=fx,
+        fy=fy,
         cx=pp_w,
         cy=pp_h,
-        camera_type=camera_type,
+        camera_type=camera_state.camera_type,
         camera_to_worlds=camera_state.c2w.to(torch.float32)[None, ...],
         times=torch.tensor([0.0], dtype=torch.float32),
     )

nerfstudio/viewer_beta/viewer.py

@@ -26,8 +26,8 @@ import torchvision
 import viser
 import viser.theme
 import viser.transforms as vtf
-
 from nerfstudio.cameras.camera_optimizers import CameraOptimizer
+from nerfstudio.cameras.cameras import CameraType
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.models.base_model import Model
@@ -41,6 +41,7 @@ from nerfstudio.viewer_beta.render_panel import populate_render_tab
 from nerfstudio.viewer_beta.render_state_machine import RenderAction, RenderStateMachine
 from nerfstudio.viewer_beta.utils import CameraState, parse_object
 from nerfstudio.viewer_beta.viewer_elements import ViewerControl, ViewerElement
+from typing_extensions import assert_never
 
 if TYPE_CHECKING:
     from nerfstudio.engine.trainer import Trainer
@@ -80,6 +81,7 @@ class Viewer:
         train_lock: Optional[threading.Lock] = None,
         share: bool = False,
     ):
+        self.ready = False  # Set to True at end of constructor.
         self.config = config
         self.trainer = trainer
         self.last_step = 0
@@ -182,15 +184,16 @@ class Viewer:
                 self.viser_server,
                 self.include_time,
                 VISER_NERFSTUDIO_SCALE_RATIO,
-                self._interrupt_render,
-                self._crop_params_update,
+                self._trigger_rerender,
                 self._output_type_change,
                 self._output_split_type_change,
                 default_composite_depth=self.config.default_composite_depth,
             )
         config_path = self.log_filename.parents[0] / "config.yml"
         with tabs.add_tab("Render", viser.Icon.CAMERA):
-            populate_render_tab(self.viser_server, config_path, self.datapath, self.control_panel)
+            self.render_tab_state = populate_render_tab(
+                self.viser_server, config_path, self.datapath, self.control_panel
+            )
 
         with tabs.add_tab("Export", viser.Icon.PACKAGE_EXPORT):
             populate_export_tab(self.viser_server, self.control_panel, config_path)
@@ -203,7 +206,7 @@ class Viewer:
                 element.install(self.viser_server)
                 # also rewire the hook to rerender
                 prev_cb = element.cb_hook
-                element.cb_hook = lambda element: [prev_cb(element), self._interrupt_render(element)]
+                element.cb_hook = lambda element: [prev_cb(element), self._trigger_rerender()]
             else:
                 # recursively create folders
                 # If the folder name is "Custom Elements/a/b", then:
@@ -237,6 +240,8 @@ class Viewer:
         for c in self.viewer_controls:
             c._setup(self)
 
+        self.ready = True
+
     def toggle_pause_button(self) -> None:
         self.pause_train.visible = not self.pause_train.visible
         self.resume_train.visible = not self.resume_train.visible
@@ -266,7 +271,27 @@ class Viewer:
         R = torch.tensor(R.as_matrix())
         pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
         c2w = torch.concatenate([R, pos[:, None]], dim=1)
-        camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w)
+        if self.ready and self.render_tab_state.preview_render:
+            camera_type = self.render_tab_state.preview_camera_type
+            camera_state = CameraState(
+                fov=self.render_tab_state.preview_fov,
+                aspect=self.render_tab_state.preview_aspect,
+                c2w=c2w,
+                camera_type=CameraType.PERSPECTIVE
+                if camera_type == "Perspective"
+                else CameraType.FISHEYE
+                if camera_type == "Fisheye"
+                else CameraType.EQUIRECTANGULAR
+                if camera_type == "Equirectangular"
+                else assert_never(camera_type),
+            )
+        else:
+            camera_state = CameraState(
+                fov=client.camera.fov,
+                aspect=client.camera.aspect,
+                c2w=c2w,
+                camera_type=CameraType.PERSPECTIVE,
+            )
         return camera_state
 
     def handle_disconnect(self, client: viser.ClientHandle) -> None:
@@ -278,7 +303,9 @@ class Viewer:
         self.render_statemachines[client.client_id].start()
 
         @client.camera.on_update
-        def _(cam: viser.CameraHandle) -> None:
+        def _(_: viser.CameraHandle) -> None:
+            if not self.ready:
+                return
             self.last_move_time = time.time()
             with self.viser_server.atomic():
                 camera_state = self.get_camera_state(client)
@@ -314,13 +341,14 @@ class Viewer:
             self.camera_handles[key].position = c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO
             self.camera_handles[key].wxyz = R.wxyz
 
-    def _interrupt_render(self, _) -> None:
+    def _trigger_rerender(self) -> None:
         """Interrupt current render."""
+        if not self.ready:
+            return
         clients = self.viser_server.get_clients()
         for id in clients:
             camera_state = self.get_camera_state(clients[id])
-            if camera_state is not None:
-                self.render_statemachines[id].action(RenderAction("rerender", camera_state))
+            self.render_statemachines[id].action(RenderAction("move", camera_state))
 
     def _toggle_training_state(self, _) -> None:
         """Toggle the trainer's training state."""
@@ -330,14 +358,6 @@ class Viewer:
             elif self.trainer.training_state == "paused":
                 self.trainer.training_state = "training"
 
-    def _crop_params_update(self, _) -> None:
-        """Update crop parameters"""
-        clients = self.viser_server.get_clients()
-        for id in clients:
-            camera_state = self.get_camera_state(clients[id])
-            if camera_state is not None:
-                self.render_statemachines[id].action(RenderAction("move", camera_state))
-
     def _output_type_change(self, _):
         self.output_type_changed = True
 

nerfstudio/viewer_beta/viewer_elements.py

@@ -35,7 +35,7 @@ from viser import (
     ViserServer,
 )
 
-from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.viewer_beta.utils import CameraState, get_camera
 
 if TYPE_CHECKING:
@@ -141,7 +141,9 @@ class ViewerControl:
         R = torch.tensor(R.as_matrix())
         pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
         c2w = torch.concatenate([R, pos[:, None]], dim=1)
-        camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w)
+        camera_state = CameraState(
+            fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w, camera_type=CameraType.PERSPECTIVE
+        )
         return get_camera(camera_state, img_height, img_width)
 
     def register_click_cb(self, cb: Callable):
@@ -483,7 +485,11 @@ class ViewerDropdown(ViewerParameter[TString], Generic[TString]):
     def _create_gui_handle(self, viser_server: ViserServer) -> None:
         assert self.gui_handle is None, "gui_handle should be initialized once"
         self.gui_handle = viser_server.add_gui_dropdown(
-            self.name, self.options, self.default_value, disabled=self.disabled, hint=self.hint  # type: ignore
+            self.name,
+            self.options,
+            self.default_value,
+            disabled=self.disabled,
+            hint=self.hint,  # type: ignore
         )
 
     def set_options(self, new_options: List[TString]) -> None:

pyproject.toml

@@ -56,7 +56,7 @@ dependencies = [
     "torchvision>=0.14.1",
     "torchmetrics[image]>=1.0.1",
     "typing_extensions>=4.4.0",
-    "viser==0.1.14",
+    "viser==0.1.17",
     "nuscenes-devkit>=1.1.1",
     "wandb>=0.13.3",
     "xatlas",