Merge branch 'main' into ra/nd-gauss

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
 

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
 

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"

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
 

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

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"