tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-10 18:17:31 +08:00
Code

Full implementation of features for dynamic scenes into the Viser viewer. (#2936)

Browse Source 
* Add time buttons & export render_time

* Added preview of dynamic scenes, time interpolation for rendering, etc. (Fixes #2934)

* Fixes to comply with Ruff

* Fixed UnboundLocalError in static mode

* Formatting

* Fixed pyrights compliance

* Nit remove unused sleep

* Fixed JSON export/import for dynamic camera paths

* Fixed fov override check for loading json camera paths

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>
This commit is contained in:
Sacha Lewin 2024-02-21 01:02:05 +01:00 committed by GitHub
parent 36c48c7177
commit a8b3f71eeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 147 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

173
nerfstudio/viewer/render_panel.py
View File

@ -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

9
nerfstudio/viewer/render_state_machine.py
View File

@ -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"

6
nerfstudio/viewer/utils.py
View File

@ -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

1
nerfstudio/viewer/viewer.py
View File

@ -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