diff --git a/nerfstudio/scripts/datasets/process_project_aria.py b/nerfstudio/scripts/datasets/process_project_aria.py index f2fdcc94c..10f26653f 100644 --- a/nerfstudio/scripts/datasets/process_project_aria.py +++ b/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] diff --git a/nerfstudio/viewer_beta/control_panel.py b/nerfstudio/viewer_beta/control_panel.py index c94b4745f..a937edce1 100644 --- a/nerfstudio/viewer_beta/control_panel.py +++ b/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) diff --git a/nerfstudio/viewer_beta/render_panel.py b/nerfstudio/viewer_beta/render_panel.py index 0e666038e..6c31447bf 100644 --- a/nerfstudio/viewer_beta/render_panel.py +++ b/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__": diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 849cf3d4e..54f8e69ae 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/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, diff --git a/nerfstudio/viewer_beta/utils.py b/nerfstudio/viewer_beta/utils.py index 3614fce41..8b9e23793 100644 --- a/nerfstudio/viewer_beta/utils.py +++ b/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), ) diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 2b6f4e140..0c560f9d8 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/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 diff --git a/nerfstudio/viewer_beta/viewer_elements.py b/nerfstudio/viewer_beta/viewer_elements.py index e8de855e0..b53da0535 100644 --- a/nerfstudio/viewer_beta/viewer_elements.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index fd058f9e8..cd971dd6f 100644 --- a/pyproject.toml +++ b/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",