diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index 4c63f0f86..4c60a040d 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -692,6 +692,17 @@ class SplatfactoModel(Model): else: return 1 + def _downscale_if_required(self, image): + d = self._get_downscale_factor() + if d > 1: + newsize = [image.shape[0] // d, image.shape[1] // d] + + # torchvision can be slow to import, so we do it lazily. + import torchvision.transforms.functional as TF + + return TF.resize(image.permute(2, 0, 1), newsize, antialias=None).permute(1, 2, 0) + return image + def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: """Takes in a Ray Bundle and returns a dictionary of outputs. @@ -767,12 +778,6 @@ class SplatfactoModel(Model): W, H = int(camera.width.item()), int(camera.height.item()) self.last_size = (H, W) projmat = projection_matrix(0.001, 1000, fovx, fovy, device=self.device) - BLOCK_X, BLOCK_Y = 16, 16 - tile_bounds = ( - int((W + BLOCK_X - 1) // BLOCK_X), - int((H + BLOCK_Y - 1) // BLOCK_Y), - 1, - ) if crop_ids is not None: opacities_crop = self.opacities[crop_ids] @@ -790,7 +795,7 @@ class SplatfactoModel(Model): quats_crop = self.quats colors_crop = torch.cat((features_dc_crop[:, None, :], features_rest_crop), dim=1) - + BLOCK_WIDTH = 16 # this controls the tile size of rasterization, 16 is a good default self.xys, depths, self.radii, conics, comp, num_tiles_hit, cov3d = project_gaussians( # type: ignore means_crop, torch.exp(scales_crop), @@ -804,7 +809,7 @@ class SplatfactoModel(Model): cy, H, W, - tile_bounds, + BLOCK_WIDTH, ) # type: ignore # rescale the camera back to original dimensions before returning @@ -854,6 +859,7 @@ class SplatfactoModel(Model): opacities, H, W, + BLOCK_WIDTH, background=background, return_alpha=True, ) # type: ignore @@ -871,6 +877,7 @@ class SplatfactoModel(Model): torch.sigmoid(opacities_crop), H, W, + BLOCK_WIDTH, background=torch.zeros(3, device=self.device), )[..., 0:1] # type: ignore depth_im = torch.where(alpha > 0, depth_im / alpha, depth_im.detach().max()) @@ -901,16 +908,7 @@ class SplatfactoModel(Model): """ if image.dtype == torch.uint8: image = image.float() / 255.0 - d = self._get_downscale_factor() - if d > 1: - newsize = [image.shape[0] // d, image.shape[1] // d] - - # torchvision can be slow to import, so we do it lazily. - import torchvision.transforms.functional as TF - - gt_img = TF.resize(image.permute(2, 0, 1), newsize, antialias=None).permute(1, 2, 0) - else: - gt_img = image + gt_img = self._downscale_if_required(image) return gt_img.to(self.device) def composite_with_background(self, image, background) -> torch.Tensor: @@ -956,8 +954,9 @@ class SplatfactoModel(Model): # This is a little bit sketchy for the SSIM loss. if "mask" in batch: # batch["mask"] : [H, W, 1] - assert batch["mask"].shape[:2] == gt_img.shape[:2] == pred_img.shape[:2] - mask = batch["mask"].to(self.device) + mask = self._downscale_if_required(batch["mask"]) + mask = mask.to(self.device) + assert mask.shape[:2] == gt_img.shape[:2] == pred_img.shape[:2] gt_img = gt_img * mask pred_img = pred_img * mask diff --git a/nerfstudio/viewer/render_panel.py b/nerfstudio/viewer/render_panel.py index f5c06ad19..a6bbba3a3 100644 --- a/nerfstudio/viewer/render_panel.py +++ b/nerfstudio/viewer/render_panel.py @@ -21,7 +21,7 @@ import json import threading import time from pathlib import Path -from typing import Dict, List, Literal, Optional, Tuple +from typing import Dict, List, Literal, Optional, Tuple, Union import numpy as np import scipy @@ -39,6 +39,8 @@ class Keyframe: wxyz: np.ndarray override_fov_enabled: bool override_fov_rad: float + override_time_enabled: bool + override_time_val: float aspect: float override_transition_enabled: bool override_transition_sec: Optional[float] @@ -50,6 +52,8 @@ class Keyframe: camera.wxyz, override_fov_enabled=False, override_fov_rad=camera.fov, + override_time_enabled=False, + override_time_val=0.0, aspect=aspect, override_transition_enabled=False, override_transition_sec=None, @@ -57,7 +61,9 @@ class Keyframe: class CameraPath: - def __init__(self, server: viser.ViserServer, duration_element: viser.GuiInputHandle[float]): + def __init__( + self, server: viser.ViserServer, duration_element: viser.GuiInputHandle[float], time_enabled: bool = False + ): self._server = server self._keyframes: Dict[int, Tuple[Keyframe, viser.CameraFrustumHandle]] = {} self._keyframe_counter: int = 0 @@ -76,6 +82,8 @@ class CameraPath: self.framerate: float = 30.0 self.tension: float = 0.5 # Tension / alpha term. self.default_fov: float = 0.0 + self.time_enabled = time_enabled + self.default_render_time: float = 0.0 self.default_transition_sec: float = 0.0 self.show_spline: bool = True @@ -129,6 +137,30 @@ class CameraPath: initial_value=keyframe.override_fov_rad * 180.0 / np.pi, disabled=not keyframe.override_fov_enabled, ) + if self.time_enabled: + override_time = server.add_gui_checkbox( + "Override Time", initial_value=keyframe.override_time_enabled + ) + override_time_val = server.add_gui_slider( + "Override Time", + 0.0, + 1.0, + step=0.01, + initial_value=keyframe.override_time_val, + disabled=not keyframe.override_time_enabled, + ) + + @override_time.on_update + def _(_) -> None: + keyframe.override_time_enabled = override_time.value + override_time_val.disabled = not override_time.value + self.add_camera(keyframe, keyframe_index) + + @override_time_val.on_update + def _(_) -> None: + keyframe.override_time_val = override_time_val.value + self.add_camera(keyframe, keyframe_index) + delete_button = server.add_gui_button("Delete", color="red", icon=viser.Icon.TRASH) go_to_button = server.add_gui_button("Go to") close_button = server.add_gui_button("Close") @@ -251,7 +283,9 @@ class CameraPath: # 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]]: + def interpolate_pose_and_fov_rad( + self, normalized_t: float + ) -> Optional[Union[Tuple[tf.SE3, float], Tuple[tf.SE3, float, float]]]: if len(self._keyframes) < 2: return None @@ -264,22 +298,43 @@ class CameraPath: endconditions="closed" if self.loop else "natural", ) + self._time_spline = splines.KochanekBartels( + [ + keyframe[0].override_time_val if keyframe[0].override_time_enabled else self.default_render_time + for keyframe in self._keyframes.values() + ], + tcb=(self.tension, 0.0, 0.0), + endconditions="closed" if self.loop else "natural", + ) + assert self._orientation_spline is not None assert self._position_spline is not None assert self._fov_spline is not None + if self.time_enabled: + assert self._time_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) 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), - ), - float(self._fov_spline.evaluate(spline_t)), - ) + if self.time_enabled: + return ( + tf.SE3.from_rotation_and_translation( + tf.SO3(np.array([quat.scalar, *quat.vector])), + self._position_spline.evaluate(spline_t), + ), + float(self._fov_spline.evaluate(spline_t)), + float(self._time_spline.evaluate(spline_t)), + ) + else: + return ( + tf.SE3.from_rotation_and_translation( + tf.SO3(np.array([quat.scalar, *quat.vector])), + self._position_spline.evaluate(spline_t), + ), + float(self._fov_spline.evaluate(spline_t)), + ) def update_spline(self) -> None: num_frames = int(self.compute_duration() * self.framerate) @@ -456,6 +511,7 @@ class RenderTabState: preview_render: bool preview_fov: float + preview_time: float preview_aspect: float preview_camera_type: Literal["Perspective", "Fisheye", "Equirectangular"] @@ -471,6 +527,7 @@ def populate_render_tab( render_tab_state = RenderTabState( preview_render=False, preview_fov=0.0, + preview_time=0.0, preview_aspect=1.0, preview_camera_type="Perspective", ) @@ -484,6 +541,21 @@ def populate_render_tab( hint="Field-of-view for rendering, which can also be overridden on a per-keyframe basis.", ) + render_time = None + if control_panel is not None and control_panel._time_enabled: + render_time = server.add_gui_slider( + "Default Time", + initial_value=0.0, + min=0.0, + max=1.0, + step=0.01, + hint="Rendering time step, which can also be overridden on a per-keyframe basis.", + ) + + @render_time.on_update + def _(_) -> None: + camera_path.default_render_time = render_time.value + @fov_degrees.on_update def _(_) -> None: fov_radians = fov_degrees.value / 180.0 * np.pi @@ -703,7 +775,7 @@ def populate_render_tab( preview_camera_handle.remove() preview_camera_handle = None - def compute_and_update_preview_camera_state() -> Optional[Tuple[tf.SE3, float]]: + def compute_and_update_preview_camera_state() -> Optional[Union[Tuple[tf.SE3, float], Tuple[tf.SE3, float, float]]]: """Update the render tab state with the current preview camera pose. Returns current camera pose + FOV if available.""" @@ -715,11 +787,20 @@ def populate_render_tab( if maybe_pose_and_fov_rad is None: remove_preview_camera() return - pose, fov_rad = maybe_pose_and_fov_rad + time = None + if len(maybe_pose_and_fov_rad) == 3: # Time is enabled. + pose, fov_rad, time = maybe_pose_and_fov_rad + render_tab_state.preview_time = time + else: + 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 + + if time is not None: + return pose, fov_rad, time + else: + return pose, fov_rad def add_preview_frame_slider() -> Optional[viser.GuiInputHandle[int]]: """Helper for creating the current frame # slider. This is removed and @@ -745,7 +826,10 @@ def populate_render_tab( maybe_pose_and_fov_rad = compute_and_update_preview_camera_state() if maybe_pose_and_fov_rad is None: return - pose, fov_rad = maybe_pose_and_fov_rad + if len(maybe_pose_and_fov_rad) == 3: # Time is enabled. + pose, fov_rad, time = maybe_pose_and_fov_rad + else: + pose, fov_rad = maybe_pose_and_fov_rad preview_camera_handle = server.add_camera_frustum( "/preview_camera", @@ -776,7 +860,10 @@ def populate_render_tab( if maybe_pose_and_fov_rad is None: remove_preview_camera() return - pose, fov = maybe_pose_and_fov_rad + if len(maybe_pose_and_fov_rad) == 3: # Time is enabled. + pose, fov, time = maybe_pose_and_fov_rad + else: + pose, fov = maybe_pose_and_fov_rad del fov # Hide all scene nodes when we're previewing the render. @@ -902,8 +989,10 @@ def populate_render_tab( wxyz=pose.rotation().wxyz, # 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_enabled=abs(frame["fov"] - json_data.get("default_fov", 0.0)) > 1e-3, override_fov_rad=frame["fov"] / 180.0 * np.pi, + override_time_enabled=frame.get("override_time_enabled", False), + override_time_val=frame.get("render_time", None), aspect=frame["aspect"], override_transition_enabled=frame.get("override_transition_enabled", None), override_transition_sec=frame.get("override_transition_sec", None), @@ -977,18 +1066,22 @@ def populate_render_tab( tf.SO3(keyframe.wxyz) @ tf.SO3.from_x_radians(np.pi), keyframe.position / VISER_NERFSTUDIO_SCALE_RATIO, ) - keyframes.append( - { - "matrix": pose.as_matrix().flatten().tolist(), - "fov": np.rad2deg(keyframe.override_fov_rad) - if keyframe.override_fov_enabled - else fov_degrees.value, - "aspect": keyframe.aspect, - "override_transition_enabled": keyframe.override_transition_enabled, - "override_transition_sec": keyframe.override_transition_sec, - } - ) + keyframe_dict = { + "matrix": pose.as_matrix().flatten().tolist(), + "fov": np.rad2deg(keyframe.override_fov_rad) if keyframe.override_fov_enabled else fov_degrees.value, + "aspect": keyframe.aspect, + "override_transition_enabled": keyframe.override_transition_enabled, + "override_transition_sec": keyframe.override_transition_sec, + } + if render_time is not None: + keyframe_dict["render_time"] = ( + keyframe.override_time_val if keyframe.override_time_enabled else render_time.value + ) + keyframe_dict["override_time_enabled"] = keyframe.override_time_enabled + keyframes.append(keyframe_dict) json_data["default_fov"] = fov_degrees.value + if render_time is not None: + json_data["default_time"] = render_time.value if render_time is not None else None json_data["default_transition_sec"] = transition_sec_number.value json_data["keyframes"] = keyframes json_data["camera_type"] = camera_type.value.lower() @@ -1004,19 +1097,24 @@ def populate_render_tab( maybe_pose_and_fov = camera_path.interpolate_pose_and_fov_rad(i / num_frames) if maybe_pose_and_fov is None: return - pose, fov = maybe_pose_and_fov + time = None + if len(maybe_pose_and_fov) == 3: # Time is enabled. + pose, fov, time = maybe_pose_and_fov + else: + 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(np.pi), pose.translation() / VISER_NERFSTUDIO_SCALE_RATIO, ) - camera_path_list.append( - { - "camera_to_world": pose.as_matrix().flatten().tolist(), - "fov": np.rad2deg(fov), - "aspect": resolution.value[0] / resolution.value[1], - } - ) + camera_path_list_dict = { + "camera_to_world": pose.as_matrix().flatten().tolist(), + "fov": np.rad2deg(fov), + "aspect": resolution.value[0] / resolution.value[1], + } + if time is not None: + camera_path_list_dict["render_time"] = time + camera_path_list.append(camera_path_list_dict) json_data["camera_path"] = camera_path_list # finally add crop data if crop is enabled if control_panel is not None: @@ -1064,7 +1162,10 @@ def populate_render_tab( def _(_) -> None: modal.close() - camera_path = CameraPath(server, duration_number) + if control_panel is not None: + camera_path = CameraPath(server, duration_number, control_panel._time_enabled) + else: + camera_path = CameraPath(server, duration_number) camera_path.default_fov = fov_degrees.value / 180.0 * np.pi camera_path.default_transition_sec = transition_sec_number.value diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py index d264d0686..b0afb490c 100644 --- a/nerfstudio/viewer/render_state_machine.py +++ b/nerfstudio/viewer/render_state_machine.py @@ -129,10 +129,11 @@ class RenderStateMachine(threading.Thread): image_height, image_width = self._calculate_image_res(camera_state.aspect) - if self.viewer.include_time: - camera = get_camera(camera_state, image_height, image_width, self.viewer.control_panel.time) - else: - camera = get_camera(camera_state, image_height, image_width) + # These 2 lines make the control panel's time option independent from the render panel's. + # When outside of render preview, it will use the control panel's time. + if not self.viewer.render_tab_state.preview_render and self.viewer.include_time: + camera_state.time = self.viewer.control_panel.time + camera = get_camera(camera_state, image_height, image_width) camera = camera.to(self.viewer.get_model().device) assert isinstance(camera, Cameras) assert camera is not None, "render called before viewer connected" diff --git a/nerfstudio/viewer/utils.py b/nerfstudio/viewer/utils.py index 247d42a90..a8792d494 100644 --- a/nerfstudio/viewer/utils.py +++ b/nerfstudio/viewer/utils.py @@ -39,10 +39,12 @@ class CameraState: """The camera matrix.""" camera_type: Literal[CameraType.PERSPECTIVE, CameraType.EQUIRECTANGULAR, CameraType.FISHEYE] """Type of camera to render.""" + time: float = 0.0 + """The rendering time of the camera state.""" def get_camera( - camera_state: CameraState, image_height: int, image_width: Optional[Union[int, float]] = None, time: float = 0.0 + camera_state: CameraState, image_height: int, image_width: Optional[Union[int, float]] = None ) -> Cameras: """Returns the camera intrinsics matrix and the camera to world homogeneous matrix. @@ -74,7 +76,7 @@ def get_camera( cy=pp_h, camera_type=camera_state.camera_type, camera_to_worlds=camera_state.c2w.to(torch.float32)[None, ...], - times=torch.tensor([time], dtype=torch.float32), + times=torch.tensor([camera_state.time], dtype=torch.float32), ) return camera diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index 208b8da68..db84e9550 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -305,6 +305,7 @@ class Viewer: fov=self.render_tab_state.preview_fov, aspect=self.render_tab_state.preview_aspect, c2w=c2w, + time=self.render_tab_state.preview_time, camera_type=CameraType.PERSPECTIVE if camera_type == "Perspective" else CameraType.FISHEYE diff --git a/pyproject.toml b/pyproject.toml index e2a731f64..42ce41b70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ dependencies = [ "xatlas", "trimesh>=3.20.2", "timm==0.6.7", - "gsplat>=0.1.5", + "gsplat>=0.1.6", "pytorch-msssim", "pathos", "packaging"