diff --git a/.vscode/settings.json b/.vscode/settings.json
index a02b23dd9..4c3a7447a 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -24,7 +24,7 @@
"typescript.suggestionActions.enabled": false,
"javascript.suggestionActions.enabled": false,
"[python]": {
- "editor.defaultFormatter": "ms-python.python",
+ "editor.defaultFormatter": "ms-python.black-formatter",
"editor.codeActionsOnSave": {
"source.organizeImports": true,
"source.fixAll": true
@@ -32,7 +32,7 @@
},
"editor.formatOnSave": true,
"python.envFile": "${workspaceFolder}/.env",
- "python.formatting.provider": "black",
+ "python.formatting.provider": "none",
"python.formatting.blackArgs": ["--line-length=120"],
"python.linting.pylintEnabled": false,
"python.linting.flake8Enabled": false,
diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py
index 8ad46e223..c2ef887a6 100644
--- a/nerfstudio/cameras/cameras.py
+++ b/nerfstudio/cameras/cameras.py
@@ -32,7 +32,7 @@ import nerfstudio.utils.math
import nerfstudio.utils.poses as pose_utils
from nerfstudio.cameras import camera_utils
from nerfstudio.cameras.rays import RayBundle
-from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.data.scene_box import SceneBox, OrientedBox
from nerfstudio.utils.tensor_dataclass import TensorDataclass
TORCH_DEVICE = Union[torch.device, str]
@@ -323,6 +323,7 @@ class Cameras(TensorDataclass):
keep_shape: Optional[bool] = None,
disable_distortion: bool = False,
aabb_box: Optional[SceneBox] = None,
+ obb_box: Optional[OrientedBox] = None,
) -> RayBundle:
"""Generates rays for the given camera indices.
@@ -465,20 +466,24 @@ class Cameras(TensorDataclass):
if keep_shape is False:
raybundle = raybundle.flatten()
- if aabb_box:
+ if aabb_box is not None or obb_box is not None:
with torch.no_grad():
- tensor_aabb = Parameter(aabb_box.aabb.flatten(), requires_grad=False)
-
rays_o = raybundle.origins.contiguous()
rays_d = raybundle.directions.contiguous()
- tensor_aabb = tensor_aabb.to(rays_o.device)
shape = rays_o.shape
rays_o = rays_o.reshape((-1, 3))
rays_d = rays_d.reshape((-1, 3))
- t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb)
+ if aabb_box is not None:
+ tensor_aabb = Parameter(aabb_box.aabb.flatten(), requires_grad=False)
+ tensor_aabb = tensor_aabb.to(rays_o.device)
+ t_min, t_max = nerfstudio.utils.math.intersect_aabb(rays_o, rays_d, tensor_aabb)
+ elif obb_box is not None:
+ t_min, t_max = nerfstudio.utils.math.intersect_obb(rays_o, rays_d, obb_box)
+ else:
+ assert False
t_min = t_min.reshape([shape[0], shape[1], 1])
t_max = t_max.reshape([shape[0], shape[1], 1])
@@ -747,8 +752,6 @@ class Cameras(TensorDataclass):
rotation = c2w[..., :3, :3]
- -torch.pi * ((x - cx) / fx)[0]
-
# interocular axis of the VR180 camera
vr180_x_axis = torch.tensor([1, 0, 0], device=c2w.device)
diff --git a/nerfstudio/data/scene_box.py b/nerfstudio/data/scene_box.py
index aaa055645..835c2b1cf 100644
--- a/nerfstudio/data/scene_box.py
+++ b/nerfstudio/data/scene_box.py
@@ -17,7 +17,8 @@ Dataset input structures.
"""
from dataclasses import dataclass
-from typing import Union
+from typing import Union, Tuple
+import viser.transforms as vtf
import torch
from jaxtyping import Float
@@ -76,3 +77,38 @@ class SceneBox:
xyzs = poses[..., :3, -1]
aabb = torch.stack([torch.min(xyzs, dim=0)[0], torch.max(xyzs, dim=0)[0]])
return SceneBox(aabb=aabb * scale_factor)
+
+
+@dataclass
+class OrientedBox:
+ R: Float[Tensor, "3 3"]
+ """R: rotation matrix."""
+ T: Float[Tensor, "3"]
+ """T: translation vector."""
+ S: Float[Tensor, "3"]
+ """S: scale vector."""
+
+ def within(self, pts: Float[Tensor, "n 3"]):
+ """Returns a boolean mask indicating whether each point is within the box."""
+ R, T, S = self.R, self.T, self.S.to(pts)
+ H = torch.eye(4, device=pts.device, dtype=pts.dtype)
+ H[:3, :3] = R
+ H[:3, 3] = T
+ H_world2bbox = torch.inverse(H)
+ pts = torch.cat((pts, torch.ones_like(pts[..., :1])), dim=-1)
+ pts = torch.matmul(H_world2bbox, pts.T).T[..., :3]
+
+ comp_l = torch.tensor(-S / 2)
+ comp_m = torch.tensor(S / 2)
+ mask = torch.all(torch.concat([pts > comp_l, pts < comp_m], dim=-1), dim=-1)
+ return mask
+
+ @staticmethod
+ def from_params(
+ pos: Tuple[float, float, float], rpy: Tuple[float, float, float], scale: Tuple[float, float, float]
+ ):
+ """Construct a box from position, rotation, and scale parameters."""
+ R = torch.tensor(vtf.SO3.from_rpy_radians(rpy[0], rpy[1], rpy[2]).as_matrix())
+ T = torch.tensor(pos)
+ S = torch.tensor(scale)
+ return OrientedBox(R=R, T=T, S=S)
diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py
index 1baeeb2b5..66d658690 100644
--- a/nerfstudio/exporter/exporter_utils.py
+++ b/nerfstudio/exporter/exporter_utils.py
@@ -33,6 +33,7 @@ from torch import Tensor
from nerfstudio.cameras.cameras import Cameras
from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.data.scene_box import OrientedBox
from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline
from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
@@ -84,8 +85,9 @@ def generate_point_cloud(
depth_output_name: str = "depth",
normal_output_name: Optional[str] = None,
use_bounding_box: bool = True,
- bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0),
- bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+ bounding_box_min: Optional[Tuple[float, float, float]] = None,
+ bounding_box_max: Optional[Tuple[float, float, float]] = None,
+ crop_obb: Optional[OrientedBox] = None,
std_ratio: float = 10.0,
) -> o3d.geometry.PointCloud:
"""Generate a point cloud from a nerf.
@@ -119,7 +121,8 @@ def generate_point_cloud(
rgbs = []
normals = []
view_directions = []
-
+ if use_bounding_box and (crop_obb is not None and bounding_box_max is not None):
+ CONSOLE.print("Provided aabb and crop_obb at the same time, using only the obb", style="bold yellow")
with progress as progress_bar:
task = progress_bar.add_task("Generating Point Cloud", total=num_points)
while not progress_bar.finished:
@@ -163,12 +166,15 @@ def generate_point_cloud(
normal = normal[mask]
if use_bounding_box:
- comp_l = torch.tensor(bounding_box_min, device=point.device)
- comp_m = torch.tensor(bounding_box_max, device=point.device)
- assert torch.all(
- comp_l < comp_m
- ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
- mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
+ if crop_obb is None:
+ comp_l = torch.tensor(bounding_box_min, device=point.device)
+ comp_m = torch.tensor(bounding_box_max, device=point.device)
+ assert torch.all(
+ comp_l < comp_m
+ ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
+ mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
+ else:
+ mask = crop_obb.within(point)
point = point[mask]
rgb = rgb[mask]
view_direction = view_direction[mask]
diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py
index 8bc5b4d78..20a506b1e 100644
--- a/nerfstudio/model_components/renderers.py
+++ b/nerfstudio/model_components/renderers.py
@@ -219,6 +219,7 @@ class RGBRenderer(nn.Module):
Returns:
Outputs of rgb values.
"""
+
if background_color is None:
background_color = self.background_color
diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py
index d5ccd82bb..2e4fc68ca 100644
--- a/nerfstudio/scripts/exporter.py
+++ b/nerfstudio/scripts/exporter.py
@@ -34,6 +34,7 @@ from typing_extensions import Annotated, Literal
from nerfstudio.cameras.rays import RayBundle
from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
+from nerfstudio.data.scene_box import OrientedBox
from nerfstudio.exporter import texture_utils, tsdf_utils
from nerfstudio.exporter.exporter_utils import (
collect_camera_poses,
@@ -109,10 +110,17 @@ class ExportPointCloud(Exporter):
"""Name of the RGB output."""
use_bounding_box: bool = True
"""Only query points within the bounding box"""
- bounding_box_min: Tuple[float, float, float] = (-1, -1, -1)
+ bounding_box_min: Optional[Tuple[float, float, float]] = (-1, -1, -1)
"""Minimum of the bounding box, used if use_bounding_box is True."""
- bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
+ bounding_box_max: Optional[Tuple[float, float, float]] = (1, 1, 1)
"""Maximum of the bounding box, used if use_bounding_box is True."""
+
+ obb_center: Optional[Tuple[float, float, float]] = None
+ """Center of the oriented bounding box."""
+ obb_rotation: Optional[Tuple[float, float, float]] = None
+ """Rotation of the oriented bounding box. Expressed as RPY Euler angles in radians"""
+ obb_scale: Optional[Tuple[float, float, float]] = None
+ """Scale of the oriented bounding box along each axis."""
num_rays_per_batch: int = 32768
"""Number of rays to evaluate per batch. Decrease if you run out of memory."""
std_ratio: float = 10.0
@@ -135,7 +143,9 @@ class ExportPointCloud(Exporter):
# Whether the normals should be estimated based on the point cloud.
estimate_normals = self.normal_method == "open3d"
-
+ crop_obb = None
+ if self.obb_center is not None and self.obb_rotation is not None and self.obb_scale is not None:
+ crop_obb = OrientedBox.from_params(self.obb_center, self.obb_rotation, self.obb_scale)
pcd = generate_point_cloud(
pipeline=pipeline,
num_points=self.num_points,
@@ -148,6 +158,7 @@ class ExportPointCloud(Exporter):
use_bounding_box=self.use_bounding_box,
bounding_box_min=self.bounding_box_min,
bounding_box_max=self.bounding_box_max,
+ crop_obb=crop_obb,
std_ratio=self.std_ratio,
)
torch.cuda.empty_cache()
@@ -264,6 +275,12 @@ class ExportPoissonMesh(Exporter):
"""Minimum of the bounding box, used if use_bounding_box is True."""
bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
"""Minimum of the bounding box, used if use_bounding_box is True."""
+ obb_center: Optional[Tuple[float, float, float]] = None
+ """Center of the oriented bounding box."""
+ obb_rotation: Optional[Tuple[float, float, float]] = None
+ """Rotation of the oriented bounding box. Expressed as RPY Euler angles in radians"""
+ obb_scale: Optional[Tuple[float, float, float]] = None
+ """Scale of the oriented bounding box along each axis."""
num_rays_per_batch: int = 32768
"""Number of rays to evaluate per batch. Decrease if you run out of memory."""
texture_method: Literal["point_cloud", "nerf"] = "nerf"
@@ -296,6 +313,10 @@ class ExportPoissonMesh(Exporter):
# Whether the normals should be estimated based on the point cloud.
estimate_normals = self.normal_method == "open3d"
+ if self.obb_center is not None and self.obb_rotation is not None and self.obb_scale is not None:
+ crop_obb = OrientedBox.from_params(self.obb_center, self.obb_rotation, self.obb_scale)
+ else:
+ crop_obb = None
pcd = generate_point_cloud(
pipeline=pipeline,
@@ -309,6 +330,7 @@ class ExportPoissonMesh(Exporter):
use_bounding_box=self.use_bounding_box,
bounding_box_min=self.bounding_box_min,
bounding_box_max=self.bounding_box_max,
+ crop_obb=crop_obb,
std_ratio=self.std_ratio,
)
torch.cuda.empty_cache()
diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py
index 788a0520d..8567a8b28 100644
--- a/nerfstudio/utils/math.py
+++ b/nerfstudio/utils/math.py
@@ -21,7 +21,7 @@ import torch
from jaxtyping import Bool, Float
from torch import Tensor
-from nerfstudio.utils.misc import torch_compile
+from nerfstudio.data.scene_box import OrientedBox
def components_from_spherical_harmonics(
@@ -199,7 +199,7 @@ def expected_sin(x_means: torch.Tensor, x_vars: torch.Tensor) -> torch.Tensor:
return torch.exp(-0.5 * x_vars) * torch.sin(x_means)
-@torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
+# @torch_compile(dynamic=True, mode="reduce-overhead", backend="eager")
def intersect_aabb(
origins: torch.Tensor,
directions: torch.Tensor,
@@ -240,6 +240,42 @@ def intersect_aabb(
return t_min, t_max
+def intersect_obb(
+ origins: torch.Tensor,
+ directions: torch.Tensor,
+ obb: OrientedBox,
+ max_bound: float = 1e10,
+ invalid_value: float = 1e10,
+):
+ """
+ Ray intersection with an oriented bounding box (OBB)
+
+ Args:
+ origins: [N,3] tensor of 3d positions
+ directions: [N,3] tensor of normalized directions
+ R: [3,3] rotation matrix
+ T: [3] translation vector
+ S: [3] extents of the bounding box
+ max_bound: Maximum value of t_max
+ invalid_value: Value to return in case of no intersection
+ """
+ # Transform ray to OBB space
+ R, T, S = obb.R, obb.T, obb.S.to(origins.device)
+ H = torch.eye(4, device=origins.device, dtype=origins.dtype)
+ H[:3, :3] = R
+ H[:3, 3] = T
+ H_world2bbox = torch.inverse(H)
+ origins = torch.cat((origins, torch.ones_like(origins[..., :1])), dim=-1)
+ origins = torch.matmul(H_world2bbox, origins.T).T[..., :3]
+ directions = torch.matmul(H_world2bbox[:3, :3], directions.T).T
+
+ # Compute intersection with axis-aligned bounding box with min as -S and max as +S
+ aabb = torch.concat((-S / 2, S / 2))
+ t_min, t_max = intersect_aabb(origins, directions, aabb, max_bound=max_bound, invalid_value=invalid_value)
+
+ return t_min, t_max
+
+
def safe_normalize(
vectors: Float[Tensor, "*batch_dim N"],
eps: float = 1e-10,
diff --git a/nerfstudio/viewer_beta/control_panel.py b/nerfstudio/viewer_beta/control_panel.py
index a56d4bb74..07f556dd4 100644
--- a/nerfstudio/viewer_beta/control_panel.py
+++ b/nerfstudio/viewer_beta/control_panel.py
@@ -16,11 +16,13 @@
from collections import defaultdict
from typing import Callable, DefaultDict, List, Tuple, get_args
+import numpy as np
import torch
-from viser import ViserServer
-
+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,
+ ViewerButton,
ViewerButtonGroup,
ViewerCheckbox,
ViewerDropdown,
@@ -30,6 +32,7 @@ from nerfstudio.viewer_beta.viewer_elements import ( # ViewerButtonGroup,
ViewerSlider,
ViewerVec3,
)
+from viser import ViserServer
class ControlPanel:
@@ -47,21 +50,29 @@ class ControlPanel:
self,
viser_server: ViserServer,
time_enabled: bool,
+ scale_ratio: float,
rerender_cb: Callable,
crop_update_cb: Callable,
update_output_cb: Callable,
update_split_output_cb: Callable,
+ toggle_training_state_cb: Callable,
+ camera_vis: Callable,
):
+ self.viser_scale_ratio = scale_ratio
# elements holds a mapping from tag: [elements]
self.viser_server = viser_server
self._elements_by_tag: DefaultDict[str, List[ViewerElement]] = defaultdict(lambda: [])
self._train_speed = ViewerButtonGroup(
- name="Train Speed ",
- default_value="Balanced",
- options=["Slow", "Balanced", "Fast"],
+ name="Train Speed",
+ default_value="Mid",
+ options=["Slow", "Mid", "Fast"],
cb_hook=lambda han: self._train_speed_cb(),
)
+ self._reset_camera = ViewerButton(
+ name="Reset Up Direction",
+ cb_hook=lambda han: self._reset_camera_cb(),
+ )
self._output_render = ViewerDropdown(
"Output Render",
"not set",
@@ -118,6 +129,9 @@ class ControlPanel:
step=0.05,
hint="Target training utilization, 0.0 is slow, 1.0 is fast. Doesn't affect final render quality",
)
+ self._layer_depth = ViewerCheckbox(
+ "Composite Depth", True, cb_hook=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"
)
@@ -130,27 +144,76 @@ class ControlPanel:
self._background_color = ViewerRGB(
"Background color", (38, 42, 55), cb_hook=crop_update_cb, hint="Color of the background"
)
- self._crop_min = ViewerVec3(
- "Crop Min", (-1, -1, -1), 0.05, cb_hook=crop_update_cb, hint="Minimum value of the crop"
+ self._crop_handle = self.viser_server.add_transform_controls("Crop", depth_test=False, line_width=4.0)
+
+ def update_center(han):
+ self._crop_handle.position = tuple(p * self.viser_scale_ratio for p in han.value)
+
+ self._crop_center = ViewerVec3(
+ "Crop Center",
+ (0.0, 0.0, 0.0),
+ step=0.01,
+ cb_hook=lambda e: [crop_update_cb(e), update_center(e)],
+ hint="Center of the crop box",
)
- self._crop_max = ViewerVec3(
- "Crop Max", (1, 1, 1), 0.05, cb_hook=crop_update_cb, hint="Maximum value of the crop"
+
+ def update_rot(han):
+ self._crop_handle.wxyz = vtf.SO3.from_rpy_radians(*han.value).wxyz
+
+ self._crop_rot = ViewerVec3(
+ "Crop Rotation",
+ (0.0, 0.0, 0.0),
+ step=0.01,
+ cb_hook=lambda e: [crop_update_cb(e), 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"
+ )
+
+ @self._crop_handle.on_update
+ def _update_crop_handle(han):
+ pos = self._crop_handle.position
+ self._crop_center.value = tuple(p / self.viser_scale_ratio for p in pos)
+ 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_enabled = time_enabled
+ self.stat_folder = self.viser_server.add_gui_folder("Stats")
+ with self.stat_folder:
+ self.markdown = self.viser_server.add_gui_markdown("Step: 0")
+ self.pause_train = viser_server.add_gui_button(label="Pause Training", disabled=False)
+ self.pause_train.on_click(lambda _: self.toggle_pause_button())
+ self.pause_train.on_click(lambda han: toggle_training_state_cb(han))
+ self.resume_train = viser_server.add_gui_button(label="Resume Training", disabled=False)
+ self.resume_train.on_click(lambda _: self.toggle_pause_button())
+ self.resume_train.on_click(lambda han: toggle_training_state_cb(han))
+ self.resume_train.visible = False
+ # Add buttons to toggle training image visibility
+ self.hide_images = viser_server.add_gui_button(label="Hide Train Cams", disabled=False)
+ self.hide_images.on_click(lambda _: camera_vis(False))
+ self.hide_images.on_click(lambda _: self.toggle_cameravis_button())
+ self.show_images = viser_server.add_gui_button(label="Show Train Cams", disabled=False)
+ self.show_images.on_click(lambda _: camera_vis(True))
+ self.show_images.on_click(lambda _: self.toggle_cameravis_button())
+ self.show_images.visible = False
+
self.add_element(self._train_speed)
self.add_element(self._train_util)
+ self.add_element(self._reset_camera)
with self.viser_server.add_gui_folder("Render Options"):
self.add_element(self._max_res)
self.add_element(self._output_render)
self.add_element(self._colormap)
+ self.add_element(self._layer_depth)
# colormap options
- with self.viser_server.add_gui_folder(" "):
- self.add_element(self._invert, additional_tags=("colormap",))
- self.add_element(self._normalize, additional_tags=("colormap",))
- self.add_element(self._min, additional_tags=("colormap",))
- self.add_element(self._max, additional_tags=("colormap",))
+ self.add_element(self._invert, additional_tags=("colormap",))
+ self.add_element(self._normalize, additional_tags=("colormap",))
+ self.add_element(self._min, additional_tags=("colormap",))
+ self.add_element(self._max, additional_tags=("colormap",))
# split options
with self.viser_server.add_gui_folder("Split Screen"):
@@ -159,19 +222,19 @@ class ControlPanel:
self.add_element(self._split_percentage, additional_tags=("split",))
self.add_element(self._split_output_render, additional_tags=("split",))
self.add_element(self._split_colormap, additional_tags=("split",))
- with self.viser_server.add_gui_folder(" "):
- self.add_element(self._split_invert, additional_tags=("split_colormap",))
- self.add_element(self._split_normalize, additional_tags=("split_colormap",))
- self.add_element(self._split_min, additional_tags=("split_colormap",))
- self.add_element(self._split_max, additional_tags=("split_colormap",))
+
+ self.add_element(self._split_invert, additional_tags=("split_colormap",))
+ self.add_element(self._split_normalize, additional_tags=("split_colormap",))
+ self.add_element(self._split_min, additional_tags=("split_colormap",))
+ self.add_element(self._split_max, additional_tags=("split_colormap",))
with self.viser_server.add_gui_folder("Crop Viewport"):
self.add_element(self._crop_viewport)
-
# Crop options
self.add_element(self._background_color, additional_tags=("crop",))
- self.add_element(self._crop_min, additional_tags=("crop",))
- self.add_element(self._crop_max, additional_tags=("crop",))
+ self.add_element(self._crop_center, additional_tags=("crop",))
+ self.add_element(self._crop_scale, additional_tags=("crop",))
+ self.add_element(self._crop_rot, additional_tags=("crop",))
self.add_element(self._time, additional_tags=("time",))
@@ -182,13 +245,36 @@ class ControlPanel:
if self._train_speed.value == "Fast":
self._train_util.value = 0.95
self._max_res.value = 256
- elif self._train_speed.value == "Balanced":
+ elif self._train_speed.value == "Mid":
self._train_util.value = 0.85
self._max_res.value = 512
elif self._train_speed.value == "Slow":
self._train_util.value = 0.5
self._max_res.value = 1024
+ def _reset_camera_cb(self) -> None:
+ for client in self.viser_server.get_clients().values():
+ client.camera.up_direction = vtf.SO3(client.camera.wxyz) @ np.array([0.0, -1.0, 0.0])
+
+ def toggle_pause_button(self) -> None:
+ self.pause_train.visible = not self.pause_train.visible
+ self.resume_train.visible = not self.resume_train.visible
+
+ def toggle_cameravis_button(self) -> None:
+ self.hide_images.visible = not self.hide_images.visible
+ self.show_images.visible = not self.show_images.visible
+
+ def update_step(self, step):
+ """
+ Args:
+ step: the train step to set the model to
+ """
+ with self.viser_server.atomic(), self.stat_folder:
+ # TODO change to a .value call instead of remove() and add, this makes it jittery
+ with self.viser_server.atomic():
+ self.markdown.remove()
+ self.markdown = self.viser_server.add_gui_markdown(f"Step: {step}")
+
def update_output_options(self, new_options: List[str]):
"""
Args:
@@ -226,6 +312,7 @@ class ControlPanel:
self._split_output_render.set_hidden(not self._split.value)
self._split_colormap.set_hidden(not self._split.value)
self._split_colormap.set_disabled(self.split_output_render == "rgb")
+ self._crop_handle.visible = self.crop_viewport
def update_colormap_options(self, dimensions: int, dtype: type) -> None:
"""update the colormap options based on the current render
@@ -286,24 +373,12 @@ class ControlPanel:
self._crop_viewport.value = value
@property
- def crop_min(self) -> Tuple[float, float, float]:
- """Returns the current crop min setting"""
- return self._crop_min.value
-
- @crop_min.setter
- def crop_min(self, value: Tuple[float, float, float]):
- """Sets the crop min setting"""
- self._crop_min.value = value
-
- @property
- def crop_max(self) -> Tuple[float, float, float]:
- """Returns the current crop max setting"""
- return self._crop_max.value
-
- @crop_max.setter
- def crop_max(self, value: Tuple[float, float, float]):
- """Sets the crop max setting"""
- self._crop_max.value = value
+ def crop_obb(self):
+ """Returns the current crop obb setting"""
+ rxyz = self._crop_rot.value
+ R = torch.tensor(vtf.SO3.from_rpy_radians(rxyz[0], rxyz[1], rxyz[2]).as_matrix())
+ obb = OrientedBox(R, torch.tensor(self._crop_center.value), torch.tensor(self._crop_scale.value))
+ return obb
@property
def background_color(self) -> Tuple[int, int, int]:
@@ -347,6 +422,10 @@ class ControlPanel:
invert=self._split_invert.value,
)
+ @property
+ def layer_depth(self):
+ return self._layer_depth.value
+
def _get_colormap_options(dimensions: int, dtype: type) -> List[Colormaps]:
"""
diff --git a/nerfstudio/viewer_beta/export_panel.py b/nerfstudio/viewer_beta/export_panel.py
new file mode 100644
index 000000000..425e10a19
--- /dev/null
+++ b/nerfstudio/viewer_beta/export_panel.py
@@ -0,0 +1,153 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import viser
+import viser.transforms as vtf
+from nerfstudio.data.scene_box import OrientedBox
+from nerfstudio.viewer_beta.control_panel import ControlPanel
+from typing_extensions import Literal
+
+
+def populate_export_tab(server: viser.ViserServer, control_panel: ControlPanel, config_path: Path) -> None:
+ crop_output = server.add_gui_checkbox("Use Crop", False)
+
+ @crop_output.on_update
+ def _(_) -> None:
+ control_panel.crop_viewport = crop_output.value
+
+ with server.add_gui_folder("Point Cloud"):
+ populate_point_cloud_tab(server, control_panel, config_path)
+ with server.add_gui_folder("Mesh"):
+ populate_mesh_tab(server, control_panel, config_path)
+
+
+def show_command_modal(client: viser.ClientHandle, what: Literal["mesh", "point cloud"], command: str) -> None:
+ """Show a modal to each currently connected client.
+
+ In the future, we should only show the modal to the client that pushes the
+ generation button.
+ """
+ with client.add_gui_modal(what.title() + " Export") as modal:
+ client.add_gui_markdown(
+ "\n".join(
+ [
+ f"To export a {what}, run the following from the command line:",
+ "",
+ "```",
+ command,
+ "```",
+ ]
+ )
+ )
+ close_button = client.add_gui_button("Close")
+
+ @close_button.on_click
+ def _(_) -> None:
+ modal.close()
+
+
+def get_crop_string(obb: OrientedBox):
+ """Takes in an oriented bounding box and returns a string of the form "--obb_{center,rotation,scale}
+ and each arg formatted with spaces around it
+ """
+ rpy = vtf.SO3.from_matrix(obb.R.numpy(force=True)).as_rpy_radians()
+ pos = obb.T.squeeze().tolist()
+ scale = obb.S.squeeze().tolist()
+ rpystring = " ".join([str(x) for x in rpy])
+ posstring = " ".join([str(x) for x in pos])
+ scalestring = " ".join([str(x) for x in scale])
+ return f"--obb_center {posstring} --obb_rotation {rpystring} --obb_scale {scalestring}"
+
+
+def populate_point_cloud_tab(
+ server: viser.ViserServer,
+ control_panel: ControlPanel,
+ config_path: Path,
+) -> None:
+ server.add_gui_markdown("Render depth, project to an oriented point cloud, and filter. ")
+ num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1)
+ remove_outliers = server.add_gui_checkbox("Remove outliers", True)
+ normals = server.add_gui_dropdown(
+ "Normals",
+ # TODO: options here could depend on what's available to the model.
+ ("open3d", "model_output"),
+ initial_value="open3d",
+ hint="Normal map source.",
+ )
+ output_dir = server.add_gui_text("Output Directory", initial_value="exports/pcd/")
+ generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2)
+
+ @generate_command.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client is not None
+ command = " ".join(
+ [
+ "ns-export pointcloud",
+ f"--load-config {config_path}",
+ f"--output-dir {output_dir.value}",
+ f"--num-points {num_points.value}",
+ f"--remove-outliers {remove_outliers.value}",
+ f"--normal-method {normals.value}",
+ f"--use_bounding_box {control_panel.crop_viewport}",
+ get_crop_string(control_panel.crop_obb),
+ ]
+ )
+ show_command_modal(event.client, "point cloud", command)
+
+
+def populate_mesh_tab(
+ server: viser.ViserServer,
+ control_panel: ControlPanel,
+ config_path: Path,
+) -> None:
+ server.add_gui_markdown(
+ "Render depth, project to an oriented point cloud, and run Poisson surface reconstruction."
+ )
+
+ normals = server.add_gui_dropdown(
+ "Normals",
+ ("open3d", "model_output"),
+ initial_value="open3d",
+ hint="Source for normal maps.",
+ )
+ num_faces = server.add_gui_number("# Faces", initial_value=50_000, min=1)
+ texture_resolution = server.add_gui_number("Texture Resolution", min=8, initial_value=2048)
+ output_directory = server.add_gui_text("Output Directory", initial_value="exports/mesh/")
+ num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1)
+ remove_outliers = server.add_gui_checkbox("Remove outliers", True)
+
+ generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2)
+
+ @generate_command.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client is not None
+ command = " ".join(
+ [
+ "ns-export poisson",
+ f"--load-config {config_path}",
+ f"--output-dir {output_directory.value}",
+ f"--target-num-faces {num_faces.value}",
+ f"--num-pixels-per-side {texture_resolution.value}",
+ f"--num-points {num_points.value}",
+ f"--remove-outliers {remove_outliers.value}",
+ f"--normal-method {normals.value}",
+ f"--use_bounding_box {control_panel.crop_viewport}",
+ get_crop_string(control_panel.crop_obb),
+ ]
+ )
+ show_command_modal(event.client, "mesh", command)
diff --git a/nerfstudio/viewer_beta/render_panel.py b/nerfstudio/viewer_beta/render_panel.py
new file mode 100644
index 000000000..02d4d6ce0
--- /dev/null
+++ b/nerfstudio/viewer_beta/render_panel.py
@@ -0,0 +1,568 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import colorsys
+import dataclasses
+import threading
+import time
+from typing import Dict, List, Optional, Tuple
+
+import numpy as onp
+import splines
+import splines.quaternion
+import viser
+import viser.transforms as tf
+
+
+@dataclasses.dataclass
+class Keyframe:
+ position: onp.ndarray
+ wxyz: onp.ndarray
+ override_fov_enabled: bool
+ override_fov_value: float
+ aspect: float
+
+ @staticmethod
+ def from_camera(camera: viser.CameraHandle, aspect: float) -> Keyframe:
+ return Keyframe(
+ camera.position,
+ camera.wxyz,
+ override_fov_enabled=False,
+ override_fov_value=camera.fov,
+ aspect=aspect,
+ )
+
+
+class CameraPath:
+ def __init__(self, server: viser.ViserServer):
+ self._server = server
+ self._keyframes: Dict[int, Tuple[Keyframe, viser.CameraFrustumHandle]] = {}
+ self._keyframe_counter: int = 0
+ self._spline: Optional[viser.SceneNodeHandle] = None
+ self._camera_edit_panel: Optional[viser.Gui3dContainerHandle] = None
+
+ self._orientation_spline: Optional[splines.quaternion.KochanekBartels] = None
+ self._position_spline: Optional[splines.KochanekBartels] = None
+ self._fov_spline: Optional[splines.KochanekBartels] = None
+ self._keyframes_visible: bool = True
+
+ # These parameters should be overridden externally.
+ self.loop: bool = False
+ self.smoothness: float = 0.5 # Tension / alpha term.
+ self.default_fov: float = 0.0
+
+ def set_keyframes_visible(self, visible: bool) -> None:
+ self._keyframes_visible = visible
+ for keyframe in self._keyframes.values():
+ keyframe[1].visible = visible
+
+ def add_camera(self, keyframe: Keyframe, keyframe_index: Optional[int] = None) -> None:
+ """Add a new camera, or replace an old one if `keyframe_index` is passed in."""
+ server = self._server
+
+ # Add a keyframe if we aren't replacing an existing one.
+ if keyframe_index is None:
+ keyframe_index = self._keyframe_counter
+ self._keyframe_counter += 1
+
+ frustum_handle = server.add_camera_frustum(
+ f"/render_cameras/{keyframe_index}",
+ fov=keyframe.override_fov_value if keyframe.override_fov_enabled else self.default_fov,
+ aspect=keyframe.aspect,
+ scale=0.1,
+ color=(127, 127, 127),
+ wxyz=keyframe.wxyz,
+ position=keyframe.position,
+ visible=self._keyframes_visible,
+ )
+
+ @frustum_handle.on_click
+ def _(_) -> None:
+ with server.add_3d_gui_container(
+ "/camera_edit_panel",
+ wxyz=keyframe.wxyz,
+ position=keyframe.position,
+ ) as camera_edit_panel:
+ self._camera_edit_panel = camera_edit_panel
+ override_fov = server.add_gui_checkbox("Override FOV", initial_value=keyframe.override_fov_enabled)
+ override_fov_degrees = server.add_gui_slider(
+ "Override FOV (degrees)",
+ 5.0,
+ 175.0,
+ step=0.1,
+ initial_value=keyframe.override_fov_value * 180.0 / onp.pi,
+ disabled=not keyframe.override_fov_enabled,
+ )
+ 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")
+
+ @override_fov.on_update
+ def _(_) -> None:
+ keyframe.override_fov_enabled = override_fov.value
+ override_fov_degrees.disabled = not override_fov.value
+ self.add_camera(keyframe, keyframe_index)
+
+ @override_fov_degrees.on_update
+ def _(_) -> None:
+ keyframe.override_fov_value = override_fov_degrees.value / 180.0 * onp.pi
+ self.add_camera(keyframe, keyframe_index)
+
+ @delete_button.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client is not None
+ with event.client.add_gui_modal("Confirm") as modal:
+ event.client.add_gui_markdown("Delete keyframe?")
+ confirm_button = event.client.add_gui_button("Yes", color="red", icon=viser.Icon.TRASH)
+ exit_button = event.client.add_gui_button("Cancel")
+
+ @confirm_button.on_click
+ def _(_) -> None:
+ assert camera_edit_panel is not None
+
+ keyframe_id = None
+ for i, keyframe_tuple in self._keyframes.items():
+ if keyframe_tuple[1] is frustum_handle:
+ keyframe_id = i
+ break
+ assert keyframe_id is not None
+
+ self._keyframes.pop(keyframe_id)
+ frustum_handle.remove()
+ camera_edit_panel.remove()
+ modal.close()
+ self.update_spline()
+
+ @exit_button.on_click
+ def _(_) -> None:
+ modal.close()
+
+ @go_to_button.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client is not None
+ client = event.client
+ T_world_current = tf.SE3.from_rotation_and_translation(
+ tf.SO3(client.camera.wxyz), client.camera.position
+ )
+ T_world_target = tf.SE3.from_rotation_and_translation(
+ tf.SO3(keyframe.wxyz), keyframe.position
+ ) @ tf.SE3.from_translation(onp.array([0.0, 0.0, -0.5]))
+
+ T_current_target = T_world_current.inverse() @ T_world_target
+
+ for j in range(10):
+ T_world_set = T_world_current @ tf.SE3.exp(T_current_target.log() * j / 9.0)
+
+ # Important bit: we atomically set both the orientation and the position
+ # of the camera.
+ with client.atomic():
+ client.camera.wxyz = T_world_set.rotation().wxyz
+ client.camera.position = T_world_set.translation()
+ time.sleep(1.0 / 30.0)
+
+ @close_button.on_click
+ def _(_) -> None:
+ assert camera_edit_panel is not None
+ camera_edit_panel.remove()
+
+ self._keyframes[keyframe_index] = (keyframe, frustum_handle)
+
+ def update_aspect(self, aspect: float) -> None:
+ for keyframe_index, frame in self._keyframes.items():
+ frame = dataclasses.replace(frame[0], aspect=aspect)
+ self.add_camera(frame, keyframe_index=keyframe_index)
+
+ def reset(self) -> None:
+ for frame in self._keyframes.values():
+ frame[1].remove()
+ self._keyframes.clear()
+ self.update_spline()
+
+ def interpolate_pose_and_fov(self, normalized_t: float) -> Optional[Tuple[tf.SE3, float]]:
+ if len(self._keyframes) < 2:
+ return None
+ # TODO: this doesn't need to be constantly re-instantiated.
+ self._fov_spline = splines.KochanekBartels(
+ [
+ keyframe[0].override_fov_value if keyframe[0].override_fov_enabled else self.default_fov
+ for keyframe in self._keyframes.values()
+ ],
+ tcb=(self.smoothness, 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
+ max_t = len(self._keyframes) if self.loop else len(self._keyframes) - 1
+ t = max_t * normalized_t
+ quat = self._orientation_spline.evaluate(t)
+ assert isinstance(quat, splines.quaternion.UnitQuaternion)
+ return (
+ tf.SE3.from_rotation_and_translation(
+ tf.SO3(onp.array([quat.scalar, *quat.vector])),
+ self._position_spline.evaluate(t),
+ ),
+ float(self._fov_spline.evaluate(t)),
+ )
+
+ def update_spline(self) -> None:
+ keyframes = list(self._keyframes.values())
+ if len(keyframes) <= 1:
+ if self._spline is not None:
+ self._spline.remove()
+ self._spline = None
+ return
+
+ # Update internal splines.
+ self._orientation_spline = splines.quaternion.KochanekBartels(
+ [
+ splines.quaternion.UnitQuaternion.from_unit_xyzw(onp.roll(keyframe[0].wxyz, shift=-1))
+ for keyframe in keyframes
+ ],
+ tcb=(self.smoothness, 0.0, 0.0),
+ endconditions="closed" if self.loop else "natural",
+ )
+ self._position_spline = splines.KochanekBartels(
+ [keyframe[0].position for keyframe in keyframes],
+ tcb=(self.smoothness, 0.0, 0.0),
+ endconditions="closed" if self.loop else "natural",
+ )
+
+ # Update visualized spline.
+ num_keyframes = len(keyframes) + 1 if self.loop else len(keyframes)
+ points_array = onp.array(
+ [self._position_spline.evaluate(t) for t in onp.linspace(0, num_keyframes - 1, num_keyframes * 100)]
+ )
+ colors_array = onp.array([colorsys.hls_to_rgb(h, 0.5, 1.0) for h in onp.linspace(0.0, 1.0, len(points_array))])
+ self._spline = self._server.add_point_cloud(
+ "/render_camera_spline",
+ points=points_array,
+ colors=colors_array,
+ point_size=0.035,
+ )
+
+
+def populate_render_tab(server: viser.ViserServer) -> None:
+ fov_degrees = server.add_gui_slider(
+ "FOV",
+ initial_value=90.0,
+ min=0.1,
+ max=175.0,
+ step=0.01,
+ hint="Field-of-view for rendering, which can also be overridden on a per-keyframe basis.",
+ )
+
+ @fov_degrees.on_update
+ def _(_) -> None:
+ fov_radians = fov_degrees.value / 180.0 * onp.pi
+ for client in server.get_clients().values():
+ client.camera.fov = fov_radians
+ camera_path.default_fov = fov_radians
+
+ # Updating the aspect ratio will also re-render the camera frustums.
+ # Could rethink this.
+ camera_path.update_aspect(resolution.value[0] / resolution.value[1])
+
+ resolution = server.add_gui_vector2(
+ "Resolution",
+ initial_value=(1920, 1080),
+ min=(50, 50),
+ max=(10_000, 10_000),
+ step=1,
+ hint="Tension parameter for adjusting smoothness of spline interpolation.",
+ )
+
+ @resolution.on_update
+ def _(_) -> None:
+ """Update the aspect ratio for all cameras when the resolution changes."""
+ camera_path.update_aspect(resolution.value[0] / resolution.value[1])
+
+ add_button = server.add_gui_button(
+ "Add keyframe",
+ icon=viser.Icon.PLUS,
+ hint="Add a new keyframe at the current pose.",
+ )
+
+ @add_button.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client_id is not None
+ camera = server.get_clients()[event.client_id].camera
+
+ # Add this camera to the path.
+ camera_path.add_camera(
+ Keyframe.from_camera(camera, aspect=resolution.value[0] / resolution.value[1]),
+ )
+ camera_path.update_spline()
+
+ reset_up_button = server.add_gui_button(
+ "Reset up direction",
+ icon=viser.Icon.ARROW_AUTOFIT_UP,
+ 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) @ onp.array([0.0, -1.0, 0.0])
+
+ clear_keyframes_button = server.add_gui_button(
+ "Clear keyframes",
+ icon=viser.Icon.TRASH,
+ hint="Remove all keyframes from the render path.",
+ )
+
+ @clear_keyframes_button.on_click
+ def _(event: viser.GuiEvent) -> None:
+ assert event.client_id is not None
+ client = server.get_clients()[event.client_id]
+ with client.add_gui_modal("Confirm") as modal:
+ client.add_gui_markdown("Clear all keyframes?")
+ confirm_button = client.add_gui_button("Yes", color="red", icon=viser.Icon.TRASH)
+ exit_button = client.add_gui_button("Cancel")
+
+ @confirm_button.on_click
+ def _(_) -> None:
+ camera_path.reset()
+ modal.close()
+
+ # Clear move handles.
+ if len(transform_controls) > 0:
+ for t in transform_controls:
+ t.remove()
+ transform_controls.clear()
+ return
+
+ @exit_button.on_click
+ def _(_) -> None:
+ modal.close()
+
+ loop = server.add_gui_checkbox("Loop", False)
+
+ @loop.on_update
+ def _(_) -> None:
+ camera_path.loop = loop.value
+ camera_path.update_spline()
+
+ smoothness = server.add_gui_slider(
+ "Spline Tension",
+ min=0.0,
+ max=1.0,
+ initial_value=0.0,
+ step=0.01,
+ hint="Tension parameter for adjusting smoothness of spline interpolation.",
+ )
+
+ @smoothness.on_update
+ def _(_) -> None:
+ camera_path.smoothness = smoothness.value
+ camera_path.update_spline()
+
+ move_checkbox = server.add_gui_checkbox(
+ "Move keyframes",
+ initial_value=False,
+ hint="Toggle move handles for keyframes in the scene.",
+ )
+
+ @move_checkbox.on_update
+ def _(event: viser.GuiEvent) -> None:
+ # Clear move handles when toggled off.
+ if move_checkbox.value is False:
+ for t in transform_controls:
+ t.remove()
+ transform_controls.clear()
+ return
+
+ def _make_transform_controls_callback(
+ keyframe: Tuple[Keyframe, viser.SceneNodeHandle],
+ controls: viser.TransformControlsHandle,
+ ) -> None:
+ @controls.on_update
+ def _(_) -> None:
+ keyframe[0].wxyz = controls.wxyz
+ keyframe[0].position = controls.position
+
+ keyframe[1].wxyz = controls.wxyz
+ keyframe[1].position = controls.position
+
+ camera_path.update_spline()
+
+ # Show move handles.
+ assert event.client is not None
+ for keyframe_index, keyframe in camera_path._keyframes.items():
+ controls = event.client.add_transform_controls(
+ f"/keyframe_move/{keyframe_index}",
+ scale=0.4,
+ wxyz=keyframe[0].wxyz,
+ position=keyframe[0].position,
+ )
+ transform_controls.append(controls)
+ _make_transform_controls_callback(keyframe, controls)
+
+ playback_folder = server.add_gui_folder("Playback")
+ with playback_folder:
+ duration_number = server.add_gui_number("Duration (sec)", min=0.0, max=1e8, step=0.0001, initial_value=4.0)
+ framerate_slider = server.add_gui_slider("FPS", min=1.0, max=240.0, step=1e-8, initial_value=30.0)
+ framerate_buttons = server.add_gui_button_group("", ("24", "30", "60"))
+
+ @framerate_buttons.on_click
+ def _(_) -> None:
+ framerate_slider.value = float(framerate_buttons.value)
+
+ 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)
+ show_checkbox = server.add_gui_checkbox(
+ "Show keyframes",
+ initial_value=True,
+ hint="Show keyframes in the scene.",
+ )
+
+ @show_checkbox.on_update
+ def _(_: viser.GuiEvent) -> None:
+ camera_path.set_keyframes_visible(show_checkbox.value)
+
+ 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."""
+ max_frame_index = int(framerate_slider.value * duration_number.value) - 1
+
+ if max_frame_index <= 0:
+ return None
+ with playback_folder:
+ preview_frame_slider = server.add_gui_slider(
+ "Preview frame",
+ min=0,
+ max=max_frame_index,
+ step=1,
+ initial_value=0,
+ # Place right after the pause button.
+ order=pause_button.order + 0.01,
+ )
+
+ @preview_frame_slider.on_update
+ def _(_) -> None:
+ max_frame_index = int(framerate_slider.value * duration_number.value) - 1
+ maybe_pose_and_fov = camera_path.interpolate_pose_and_fov(
+ preview_frame_slider.value / max_frame_index if max_frame_index > 0 else 0
+ )
+ if maybe_pose_and_fov is None:
+ return
+ pose, fov = maybe_pose_and_fov
+ 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 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
+
+ return preview_frame_slider
+
+ @attach_viewport_checkbox.on_update
+ def _(_) -> None:
+ if not attach_viewport_checkbox.value:
+ for client in server.get_clients().values():
+ client.camera.fov = fov_degrees.value
+
+ preview_frame_slider = add_preview_frame_slider()
+
+ @duration_number.on_update
+ @framerate_slider.on_update
+ def _(_) -> None:
+ nonlocal preview_frame_slider
+ old = preview_frame_slider
+ assert old is not None
+
+ preview_frame_slider = add_preview_frame_slider()
+ if preview_frame_slider is not None:
+ old.remove()
+ else:
+ preview_frame_slider = old
+
+ # Play the camera trajectory when the play button is pressed.
+ @play_button.on_click
+ def _(_) -> None:
+ play_button.visible = False
+ pause_button.visible = True
+
+ def play() -> None:
+ while not play_button.visible:
+ max_frame = int(framerate_slider.value * duration_number.value)
+ if max_frame > 0:
+ assert preview_frame_slider is not None
+ preview_frame_slider.value = (preview_frame_slider.value + 1) % max_frame
+ time.sleep(1.0 / framerate_slider.value)
+
+ threading.Thread(target=play).start()
+
+ # Play the camera trajectory when the play button is pressed.
+ @pause_button.on_click
+ def _(_) -> None:
+ play_button.visible = True
+ pause_button.visible = False
+
+ render_button = server.add_gui_button(
+ "Generate Command",
+ color="green",
+ icon=viser.Icon.FILE_EXPORT,
+ hint="Generate the ns-render command for rendering the camera path.",
+ )
+
+ @render_button.on_click
+ def _(event: viser.GuiEvent) -> None:
+ """TODO: write the render JSON and show the render command."""
+ assert event.client is not None
+ with event.client.add_gui_modal("TODO") as modal:
+ event.client.add_gui_markdown("TODO")
+ close_button = event.client.add_gui_button("Close")
+
+ @close_button.on_click
+ def _(_) -> None:
+ modal.close()
+
+ camera_path = CameraPath(server)
+ camera_path.default_fov = fov_degrees.value / 180.0 * onp.pi
+
+ transform_controls: List[viser.SceneNodeHandle] = []
+
+
+def main() -> None:
+ """Launch a GUI with just the render panel, for development purposes."""
+ server = viser.ViserServer()
+ server.configure_theme(dark_mode=True, control_layout="collapsible")
+ server.world_axes.visible = True
+
+ populate_render_tab(server)
+ while True:
+ time.sleep(10.0)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py
index 6be9b3403..6355a5ba6 100644
--- a/nerfstudio/viewer_beta/render_state_machine.py
+++ b/nerfstudio/viewer_beta/render_state_machine.py
@@ -17,17 +17,14 @@ from __future__ import annotations
import contextlib
import threading
-import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args
import torch
-
from nerfstudio.model_components.renderers import background_color_override_context
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 import utils
from nerfstudio.viewer_beta.utils import CameraState, get_camera
if TYPE_CHECKING:
@@ -55,7 +52,7 @@ class RenderStateMachine(threading.Thread):
viewer: the viewer state
"""
- def __init__(self, viewer: Viewer):
+ def __init__(self, viewer: Viewer, viser_scale_ratio: float):
threading.Thread.__init__(self)
self.transitions: Dict[RenderStates, Dict[RenderActions, RenderStates]] = {
s: {} for s in get_args(RenderStates)
@@ -74,11 +71,12 @@ class RenderStateMachine(threading.Thread):
self.next_action: Optional[RenderAction] = None
self.state: RenderStates = "low_static"
self.render_trigger = threading.Event()
- self.target_fps = 24
+ self.target_fps = 30
self.viewer = viewer
self.interrupt_render_flag = False
self.daemon = True
self.output_keys = {}
+ self.viser_scale_ratio = viser_scale_ratio
def action(self, action: RenderAction):
"""Takes an action and updates the state machine
@@ -88,9 +86,7 @@ class RenderStateMachine(threading.Thread):
"""
if self.next_action is None:
self.next_action = action
- elif action.action == "step" and (
- self.state == "low_move" or self.next_action.action in ("move", "static", "rerender")
- ):
+ elif action.action == "step" and (self.state == "low_move" or self.next_action.action in ("move", "rerender")):
# ignore steps if:
# 1. we are in low_moving state
# 2. the current next_action is move, static, or rerender
@@ -107,7 +103,6 @@ class RenderStateMachine(threading.Thread):
# handle interrupt logic
if self.state == "high" and self.next_action.action in ("move", "rerender"):
- print("interrupting render", self.next_action.action)
self.interrupt_render_flag = True
self.render_trigger.set()
@@ -117,14 +112,11 @@ class RenderStateMachine(threading.Thread):
Args:
camera_state: the current camera state
"""
-
# initialize the camera ray bundle
- utils.update_render_aabb(
- crop_viewport=self.viewer.control_panel.crop_viewport,
- crop_min=self.viewer.control_panel.crop_min,
- crop_max=self.viewer.control_panel.crop_max,
- model=self.viewer.get_model(),
- )
+ if self.viewer.control_panel.crop_viewport:
+ obb = self.viewer.control_panel.crop_obb
+ else:
+ obb = None
image_height, image_width = self._calculate_image_res(camera_state.aspect)
@@ -132,28 +124,39 @@ class RenderStateMachine(threading.Thread):
camera = camera.to(self.viewer.get_model().device)
assert camera is not None, "render called before viewer connected"
- with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext():
- camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=self.viewer.get_model().render_aabb)
-
- with TimeWriter(None, None, write=False) as vis_t:
+ with TimeWriter(None, None, write=False) as vis_t:
+ with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext():
+ camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb)
self.viewer.get_model().eval()
step = self.viewer.step
- if self.viewer.control_panel.crop_viewport:
- color = self.viewer.control_panel.background_color
- if color is None:
- background_color = torch.tensor([0.0, 0.0, 0.0], device=self.viewer.pipeline.model.device)
+ try:
+ if self.viewer.control_panel.crop_viewport:
+ color = self.viewer.control_panel.background_color
+ if color is None:
+ background_color = torch.tensor([0.0, 0.0, 0.0], device=self.viewer.pipeline.model.device)
+ else:
+ background_color = torch.tensor(
+ [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0],
+ device=self.viewer.get_model().device,
+ )
+ with background_color_override_context(
+ background_color
+ ), torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt):
+ outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)
else:
- background_color = torch.tensor(
- [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0],
- device=self.viewer.get_model().device,
- )
- with background_color_override_context(background_color), torch.no_grad():
- outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)
- else:
- with torch.no_grad():
- outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+ with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt):
+ outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)
+ except viewer_utils.IOChangeException:
+ self.viewer.get_model().train()
+ raise
self.viewer.get_model().train()
- num_rays = len(camera_ray_bundle)
+ num_rays = len(camera_ray_bundle)
+ if self.viewer.control_panel.layer_depth:
+ # convert to z_depth if depth compositing is enabled
+ R = camera.camera_to_worlds[0:3, 0:3].T
+ pts = camera_ray_bundle.directions * outputs["depth"]
+ pts = (R @ (pts.view(-1, 3).T)).T.view(*camera_ray_bundle.directions.shape)
+ outputs["gl_z_buf_depth"] = -pts[..., 2:3] # negative z axis is the coordinate convention
render_time = vis_t.duration
if writer.is_initialized():
writer.put_time(
@@ -164,28 +167,25 @@ class RenderStateMachine(threading.Thread):
def run(self):
"""Main loop for the render thread"""
while True:
- self.render_trigger.wait()
- self.render_trigger.clear()
+ if not self.render_trigger.wait(0.2):
+ # if we haven't received a trigger in a while, send a static action
+ if self.viewer.camera_state is not None:
+ self.action(RenderAction(action="static", camera_state=self.viewer.camera_state))
action = self.next_action
- assert action is not None, "Action should never be None at this point"
+ self.render_trigger.clear()
+ if action is None:
+ continue
self.next_action = None
if self.state == "high" and action.action == "static":
# if we are in high res and we get a static action, we don't need to do anything
continue
self.state = self.transitions[self.state][action.action]
try:
- with viewer_utils.SetTrace(self.check_interrupt):
- outputs = self._render_img(action.camera_state)
+ outputs = self._render_img(action.camera_state)
except viewer_utils.IOChangeException:
# if we got interrupted, don't send the output to the viewer
continue
self._send_output_to_viewer(outputs)
- # if we rendered a static low res, we need to self-trigger a static high-res
- if self.next_action is None:
- # if there hasn't been an action, wait for 1/target_fps seconds in case we get another move command
- time.sleep(1 / self.target_fps)
- if self.state in ["low_static", "low_move"]:
- self.action(RenderAction("static", action.camera_state))
def check_interrupt(self, frame, event, arg):
"""Raises interrupt when flag has been set and not already on lowest resolution.
@@ -234,11 +234,14 @@ class RenderStateMachine(threading.Thread):
selected_output[:, split_index] = torch.tensor([0.133, 0.157, 0.192], device=selected_output.device)
selected_output = (selected_output * 255).type(torch.uint8)
-
+ depth = (
+ outputs["gl_z_buf_depth"].cpu().numpy() * self.viser_scale_ratio if "gl_z_buf_depth" in outputs else None
+ )
self.viewer.viser_server.set_background_image(
selected_output.cpu().numpy(),
format=self.viewer.config.image_format,
jpeg_quality=self.viewer.config.jpeg_quality,
+ depth=depth,
)
def _calculate_image_res(self, aspect_ratio: float) -> Tuple[int, int]:
diff --git a/nerfstudio/viewer_beta/utils.py b/nerfstudio/viewer_beta/utils.py
index 0dd435b88..a3dd7b1c2 100644
--- a/nerfstudio/viewer_beta/utils.py
+++ b/nerfstudio/viewer_beta/utils.py
@@ -12,16 +12,18 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from __future__ import annotations
+
from dataclasses import dataclass
-from typing import Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
import numpy as np
import torch
from jaxtyping import Float
-
from nerfstudio.cameras.cameras import Cameras, CameraType
from nerfstudio.data.scene_box import SceneBox
from nerfstudio.models.base_model import Model
+from torch import nn
@dataclass
@@ -53,7 +55,7 @@ def get_camera(
pp_w = image_width / 2.0
pp_h = image_height / 2.0
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]]).float()
+ intrinsics_matrix = torch.tensor([[focal_length, 0, pp_w], [0, focal_length, pp_h], [0, 0, 1]], dtype=torch.float32)
camera_type = CameraType.PERSPECTIVE
@@ -63,7 +65,7 @@ def get_camera(
cx=pp_w,
cy=pp_h,
camera_type=camera_type,
- camera_to_worlds=camera_state.c2w,
+ camera_to_worlds=camera_state.c2w.to(torch.float32),
times=torch.tensor([0.0], dtype=torch.float32),
)
return camera
@@ -93,3 +95,53 @@ def update_render_aabb(
model.render_aabb = SceneBox(aabb=torch.stack([crop_min_tensor, crop_max_tensor], dim=0))
else:
model.render_aabb = None
+
+
+def parse_object(
+ obj: Any,
+ type_check,
+ tree_stub: str,
+) -> List[Tuple[str, Any]]:
+ """
+ obj: the object to parse
+ type_check: recursively adds instances of this type to the output
+ tree_stub: the path down the object tree to this object
+
+ Returns:
+ a list of (path/to/object, obj), which represents the path down the object tree
+ along with the object itself
+ """
+
+ def add(ret: List[Tuple[str, Any]], ts: str, v: Any):
+ """
+ helper that adds to ret, and if v exists already keeps the tree stub with
+ the shortest path
+ """
+ for i, (t, o) in enumerate(ret):
+ if o == v:
+ if len(t.split("/")) > len(ts.split("/")):
+ ret[i] = (ts, v)
+ return
+ ret.append((ts, v))
+
+ if not hasattr(obj, "__dict__"):
+ return []
+ ret = []
+ # get a list of the properties of the object, sorted by whether things are instances of type_check
+ obj_props = [(k, getattr(obj, k)) for k in dir(obj)]
+ for k, v in obj_props:
+ if k[0] == "_":
+ continue
+ new_tree_stub = f"{tree_stub}/{k}"
+ if isinstance(v, type_check):
+ add(ret, new_tree_stub, v)
+ elif isinstance(v, nn.Module):
+ if v is obj:
+ # some nn.Modules might contain infinite references, e.g. consider foo = nn.Module(), foo.bar = foo
+ # to stop infinite recursion, we skip such attributes
+ continue
+ lower_rets = parse_object(v, type_check, new_tree_stub)
+ # check that the values aren't already in the tree
+ for ts, o in lower_rets:
+ add(ret, ts, o)
+ return ret
diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py
index a5271571d..9e2314ad4 100644
--- a/nerfstudio/viewer_beta/viewer.py
+++ b/nerfstudio/viewer_beta/viewer.py
@@ -16,15 +16,16 @@
from __future__ import annotations
import threading
+import time
from pathlib import Path
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
import numpy as np
import torch
import torchvision
import viser
+import viser.theme
import viser.transforms as vtf
-
from nerfstudio.configs import base_config as cfg
from nerfstudio.data.datasets.base_dataset import InputDataset
from nerfstudio.models.base_model import Model
@@ -33,13 +34,19 @@ from nerfstudio.utils.decorators import check_main_thread, decorate_all
from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName
from nerfstudio.viewer.server import viewer_utils
from nerfstudio.viewer_beta.control_panel import ControlPanel
+from nerfstudio.viewer_beta.export_panel import populate_export_tab
+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
+from nerfstudio.viewer_beta.utils import CameraState, parse_object
+from nerfstudio.viewer_beta.viewer_elements import ViewerControl, ViewerElement
+
+from nerfstudio.cameras.camera_optimizers import CameraOptimizer
if TYPE_CHECKING:
from nerfstudio.engine.trainer import Trainer
-VISER_NERFSTUDIO_SCALE_RATIO: int = 10
+
+VISER_NERFSTUDIO_SCALE_RATIO: float = 10.0
@decorate_all([check_main_thread])
@@ -95,57 +102,148 @@ class Viewer:
self.train_btn_state: Literal["training", "paused", "completed"] = "training"
self._prev_train_state: Literal["training", "paused", "completed"] = "training"
- self.camera_message = None
-
+ self.client: Optional[viser.ClientHandle] = None
self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port)
-
- self.viser_server.on_client_connect(self.handle_new_client)
-
- self.control_panel = ControlPanel(
- self.viser_server,
- self.include_time,
- self._interrupt_render,
- self._crop_params_update,
- self._output_type_change,
- self._output_split_type_change,
+ buttons = (
+ viser.theme.TitlebarButton(
+ text="Getting Started",
+ icon=None,
+ href="https://nerf.studio",
+ ),
+ viser.theme.TitlebarButton(
+ text="Github",
+ icon="GitHub",
+ href="https://github.com/nerfstudio-project/nerfstudio",
+ ),
+ viser.theme.TitlebarButton(
+ text="Documentation",
+ icon="Description",
+ href="https://docs.nerf.studio",
+ ),
+ )
+ image = viser.theme.TitlebarImage(
+ image_url_light="https://docs.nerf.studio/en/latest/_static/imgs/logo.png",
+ image_url_dark="https://docs.nerf.studio/en/latest/_static/imgs/logo-dark.png",
+ image_alt="NerfStudio Logo",
+ href="https://docs.nerf.studio/",
+ )
+ titlebar_theme = viser.theme.TitlebarConfig(buttons=buttons, image=image)
+ self.viser_server.configure_theme(
+ titlebar_content=titlebar_theme,
+ control_layout="collapsible",
+ dark_mode=True,
+ brand_color=(255, 211, 105),
)
- self.render_statemachine = RenderStateMachine(self)
+ self.render_statemachine = RenderStateMachine(self, VISER_NERFSTUDIO_SCALE_RATIO)
+ self.viser_server.on_client_connect(self.handle_new_client)
+
+ tabs = self.viser_server.add_gui_tab_group()
+ control_tab = tabs.add_tab("Control", viser.Icon.SETTINGS)
+ with control_tab:
+ self.control_panel = ControlPanel(
+ self.viser_server,
+ self.include_time,
+ VISER_NERFSTUDIO_SCALE_RATIO,
+ self._interrupt_render,
+ self._crop_params_update,
+ self._output_type_change,
+ self._output_split_type_change,
+ self._toggle_training_state,
+ self.set_camera_visibility,
+ )
+ with tabs.add_tab("Render", viser.Icon.CAMERA):
+ populate_render_tab(self.viser_server)
+
+ with tabs.add_tab("Export", viser.Icon.PACKAGE_EXPORT):
+ config_path = self.log_filename.parents[0] / "config.yml"
+ populate_export_tab(self.viser_server, self.control_panel, config_path)
+
+ def nested_folder_install(folder_labels: List[str], element: ViewerElement):
+ if len(folder_labels) == 0:
+ 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)]
+ else:
+ with self.viser_server.add_gui_folder(folder_labels[0]):
+ nested_folder_install(folder_labels[1:], element)
+
+ with control_tab:
+ self.viewer_elements = []
+ self.viewer_elements.extend(parse_object(pipeline, ViewerElement, "Custom Elements"))
+ for param_path, element in self.viewer_elements:
+ folder_labels = param_path.split("/")[:-1]
+ nested_folder_install(folder_labels, element)
+
+ # scrape the trainer/pipeline for any ViewerControl objects to initialize them
+ self.viewer_controls: List[ViewerControl] = [
+ e for (_, e) in parse_object(self.trainer, ViewerControl, "Custom Elements")
+ ]
+ for c in self.viewer_controls:
+ c._setup(self)
self.render_statemachine.start()
def handle_new_client(self, client: viser.ClientHandle) -> None:
+ self.client = client
+ self.last_move_time = 0
+
@client.camera.on_update
- def _(_: viser.CameraHandle) -> None:
- R = vtf.SO3(wxyz=client.camera.wxyz)
+ def _(cam: viser.CameraHandle) -> None:
+ assert self.client is not None
+ self.last_move_time = time.time()
+ R = vtf.SO3(wxyz=self.client.camera.wxyz)
R = R @ vtf.SO3.from_x_radians(np.pi)
- R = torch.tensor(R.as_matrix(), dtype=torch.float32)
- pos = torch.tensor(client.camera.position, dtype=torch.float32) / VISER_NERFSTUDIO_SCALE_RATIO
+ R = torch.tensor(R.as_matrix())
+ pos = torch.tensor(self.client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
c2w = torch.concatenate([R, pos[:, None]], dim=1)
- self.camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w)
+ self.camera_state = CameraState(fov=self.client.camera.fov, aspect=self.client.camera.aspect, c2w=c2w)
self.render_statemachine.action(RenderAction("move", self.camera_state))
+ def set_camera_visibility(self, visible: bool) -> None:
+ """Toggle the visibility of the training cameras."""
+ with self.viser_server.atomic():
+ for idx in self.camera_handles:
+ self.camera_handles[idx].visible = visible
+
+ def update_camera_poses(self):
+ # Update the train camera locations based on optimization
+ assert self.camera_handles is not None
+ idxs = list(self.camera_handles.keys())
+ if hasattr(self.pipeline.datamanager, "train_camera_optimizer"):
+ camera_optimizer = self.pipeline.datamanager.train_camera_optimizer
+ else:
+ camera_optimizer = self.pipeline.model.camera_optimizer
+ with torch.no_grad():
+ assert isinstance(camera_optimizer, CameraOptimizer)
+ c2ws_delta = camera_optimizer(torch.tensor(idxs, device=camera_optimizer.device)).cpu().numpy()
+ for idx in idxs:
+ # both are numpy arrays
+ c2w_orig = self.original_c2w[idx]
+ c2w_delta = c2ws_delta[idx, ...]
+ c2w = c2w_orig @ np.concatenate((c2w_delta, np.array([[0, 0, 0, 1]])), axis=0)
+ R = vtf.SO3.from_matrix(c2w[:3, :3]) # type: ignore
+ R = R @ vtf.SO3.from_x_radians(np.pi)
+ self.camera_handles[idx].position = c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO
+ self.camera_handles[idx].wxyz = R.wxyz
+
def _interrupt_render(self, _) -> None:
"""Interrupt current render."""
if self.camera_state is not None:
self.render_statemachine.action(RenderAction("rerender", self.camera_state))
+ def _toggle_training_state(self, _) -> None:
+ """Toggle the trainer's training state."""
+ if self.trainer is not None:
+ if self.trainer.training_state == "training":
+ self.trainer.training_state = "paused"
+ elif self.trainer.training_state == "paused":
+ self.trainer.training_state = "training"
+
def _crop_params_update(self, _) -> None:
"""Update crop parameters"""
- print("Crop params not set up")
- # crop_min = torch.tensor(self.control_panel.crop_min, dtype=torch.float32)
- # crop_max = torch.tensor(self.control_panel.crop_max, dtype=torch.float32)
- # scene_box = SceneBox(aabb=torch.stack([crop_min, crop_max], dim=0))
- # self.viser_server.update_scene_box(scene_box)
- # crop_scale = crop_max - crop_min
- # crop_center = (crop_max + crop_min) / 2.0
- # self.viser_server.send_crop_params(
- # crop_enabled=self.control_panel.crop_viewport,
- # crop_bg_color=self.control_panel.background_color,
- # crop_scale=tuple(crop_scale.tolist()),
- # crop_center=tuple(crop_center.tolist()),
- # )
- # if self.camera_message is not None:
- # self.render_statemachine.action(RenderAction("rerender", self.camera_message))
+ if self.camera_state is not None:
+ self.render_statemachine.action(RenderAction("move", self.camera_state))
def _output_type_change(self, _):
self.output_type_changed = True
@@ -181,8 +279,9 @@ class Viewer:
dataset: dataset to render in the scene
train_state: Current status of training
"""
-
# draw the training cameras and images
+ self.camera_handles: Dict[int, viser.CameraFrustumHandle] = {}
+ self.original_c2w: Dict[int, np.ndarray] = {}
image_indices = self._pick_drawn_image_idxs(len(train_dataset))
for idx in image_indices:
image = train_dataset[idx]["image"]
@@ -195,16 +294,26 @@ class Viewer:
c2w = camera.camera_to_worlds.cpu().numpy()
R = vtf.SO3.from_matrix(c2w[:3, :3])
R = R @ vtf.SO3.from_x_radians(np.pi)
- self.viser_server.add_camera_frustum(
- name=f"Camera {idx}",
- fov=float(camera.fx[0]),
- scale=0.2,
+ camera_handle = self.viser_server.add_camera_frustum(
+ name=f"/cameras/camera_{idx:05d}",
+ fov=float(2 * np.arctan(camera.cx / camera.fx[0])),
+ scale=0.1,
aspect=float(camera.cx[0] / camera.cy[0]),
image=image_uint8,
wxyz=R.wxyz,
position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,
)
+ @camera_handle.on_click
+ def _(event: viser.ClickEvent[viser.CameraFrustumHandle]) -> None:
+ assert self.client is not None
+ with self.client.atomic():
+ self.client.camera.position = event.target.position
+ self.client.camera.wxyz = event.target.wxyz
+
+ self.camera_handles[idx] = camera_handle
+ self.original_c2w[idx] = c2w
+
self.train_state = train_state
self.train_util = 0.9
@@ -219,7 +328,10 @@ class Viewer:
if self.camera_state is None:
return
-
+ # this stops training while moving to make the response smoother
+ while time.time() - self.last_move_time < 0.1:
+ time.sleep(0.05)
+ # self.render_statemachine.action(RenderAction("static", self.camera_state))
if self.trainer is not None and self.trainer.training_state == "training" and self.train_util != 1:
if (
EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]
@@ -239,6 +351,8 @@ class Viewer:
if step > self.last_step + render_freq:
self.last_step = step
self.render_statemachine.action(RenderAction("step", self.camera_state))
+ self.update_camera_poses()
+ self.control_panel.update_step(step)
def update_colormap_options(self, dimensions: int, dtype: type) -> None:
"""update the colormap options based on the current render
diff --git a/nerfstudio/viewer_beta/viewer_elements.py b/nerfstudio/viewer_beta/viewer_elements.py
index f9fe0e071..916566d5d 100644
--- a/nerfstudio/viewer_beta/viewer_elements.py
+++ b/nerfstudio/viewer_beta/viewer_elements.py
@@ -19,8 +19,12 @@
from __future__ import annotations
from abc import abstractmethod
-from typing import Any, Callable, Generic, List, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Generic, List, Optional, Tuple, Union
+import numpy as np
+import torch
+import viser.transforms as vtf
from typing_extensions import LiteralString, TypeVar
from viser import (
GuiButtonGroupHandle,
@@ -30,10 +34,131 @@ from viser import (
ViserServer,
)
+from nerfstudio.cameras.cameras import Cameras
+from nerfstudio.viewer_beta.utils import CameraState, get_camera
+
+if TYPE_CHECKING:
+ from nerfstudio.viewer_beta.viewer import Viewer
+
TValue = TypeVar("TValue")
TString = TypeVar("TString", default=str, bound=str)
+@dataclass
+class ViewerClick:
+ """
+ Class representing a click in the viewer as a ray.
+ """
+
+ # the information here matches the information in the ClickMessage,
+ # but we implement a wrapper as an abstraction layer
+ origin: Tuple[float, float, float]
+ """The origin of the click in world coordinates (center of camera)"""
+ direction: Tuple[float, float, float]
+ """
+ The direction of the click if projected from the camera through the clicked pixel,
+ in world coordinates
+ """
+
+
+class ViewerControl:
+ """
+ class for exposing non-gui controls of the viewer to the user
+ """
+
+ def __init__(self):
+ # this should be a user-facing constructor, since it will be used inside the model/pipeline class
+ self.click_cbs = []
+
+ def _setup(self, viewer: Viewer):
+ """
+ Internal use only, setup the viewer control with the viewer state object
+
+ Args:
+ viewer: The viewer object (viewer.py)
+ """
+ self.viewer: Viewer = viewer
+ self.viser_server: ViserServer = viewer.viser_server
+
+ def set_pose(
+ self,
+ position: Optional[Tuple[float, float, float]] = None,
+ look_at: Optional[Tuple[float, float, float]] = None,
+ instant: bool = False,
+ ):
+ """
+ Set the camera position of the viewer camera.
+
+ Args:
+ position: The new position of the camera in world coordinates
+ look_at: The new look_at point of the camera in world coordinates
+ instant: If the camera should move instantly or animate to the new position
+ """
+ raise NotImplementedError()
+
+ def set_fov(self, fov):
+ """
+ Set the FOV of the viewer camera
+
+ Args:
+ fov: The new FOV of the camera in degrees
+
+ """
+ raise NotImplementedError()
+
+ def set_crop(self, min_point: Tuple[float, float, float], max_point: Tuple[float, float, float]):
+ """
+ Set the scene crop box of the viewer to the specified min,max point
+
+ Args:
+ min_point: The minimum point of the crop box
+ max_point: The maximum point of the crop box
+
+ """
+ raise NotImplementedError()
+
+ def get_camera(self, img_height: int, img_width: int) -> Optional[Cameras]:
+ """
+ Returns the Cameras object representing the current camera for the viewer, or None if the viewer
+ is not connected yet
+
+ Args:
+ img_height: The height of the image to get camera intrinsics for
+ img_width: The width of the image to get camera intrinsics for
+ """
+ assert self.viewer.client is not None
+
+ from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO
+
+ R = vtf.SO3(wxyz=self.viewer.client.camera.wxyz)
+ R = R @ vtf.SO3.from_x_radians(np.pi)
+ R = torch.tensor(R.as_matrix())
+ pos = torch.tensor(self.viewer.client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO
+ c2w = torch.concatenate([R, pos[:, None]], dim=1)
+ camera_state = CameraState(fov=self.viewer.client.camera.fov, aspect=self.viewer.client.camera.aspect, c2w=c2w)
+ return get_camera(camera_state, img_height, img_width)
+
+ def register_click_cb(self, cb: Callable):
+ """
+ Add a callback which will be called when a click is detected in the viewer.
+
+ Args:
+ cb: The callback to call when a click is detected.
+ The callback should take a ViewerClick object as an argument
+ """
+ self.click_cbs.append(cb)
+
+ def on_click(self, msg):
+ """
+ Internal use only, register a click in the viewer which propagates to all self.click_cbs
+ """
+ raise NotImplementedError()
+
+ @property
+ def server(self):
+ return self.viser_server
+
+
class ViewerElement(Generic[TValue]):
"""Base class for all viewer elements
@@ -100,7 +225,7 @@ class ViewerButton(ViewerElement[bool]):
super().__init__(name, disabled=disabled, cb_hook=cb_hook)
def _create_gui_handle(self, viser_server: ViserServer) -> None:
- self.gui_handle = viser_server.add_gui_button(self.name, disabled=self.disabled)
+ self.gui_handle = viser_server.add_gui_button(label=self.name, disabled=self.disabled)
def install(self, viser_server: ViserServer) -> None:
self._create_gui_handle(viser_server)
diff --git a/pyproject.toml b/pyproject.toml
index 6d8fb7804..1347f954b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,12 +46,13 @@ dependencies = [
"requests",
"rich>=12.5.1",
"scikit-image>=0.19.3",
+ "splines==0.3.0",
"tensorboard>=2.13.0",
"torch>=1.13.1",
"torchvision>=0.14.1",
"torchmetrics[image]>=1.0.1",
"typing_extensions>=4.4.0",
- "viser==0.1.0",
+ "viser==0.1.2",
"nuscenes-devkit>=1.1.1",
"wandb>=0.13.3",
"xatlas",
diff --git a/tests/utils/test_visualization.py b/tests/utils/test_visualization.py
index 8fad19579..b3d9f53d1 100644
--- a/tests/utils/test_visualization.py
+++ b/tests/utils/test_visualization.py
@@ -42,6 +42,9 @@ def test_apply_boolean_colormap():
def test_cube_center():
"""Test adding a cube"""
cube = plotly_utils.get_cube(1.0, torch.Tensor([2.0, 3.0, 4.0]))
+ assert cube.x is not None
+ assert cube.y is not None
+ assert cube.z is not None
assert cube.x[0] == 1.5
assert cube.y[0] == 2.5
assert cube.z[-1] == 4.5
@@ -50,6 +53,9 @@ def test_cube_center():
def test_aabb_center():
"""Test adding a cube with different dsid length"""
cube = plotly_utils.get_cube(torch.Tensor([2.0, 1.0, 3.0]))
+ assert cube.x is not None
+ assert cube.y is not None
+ assert cube.z is not None
assert cube.x[0] == -1
assert cube.y[0] == -0.5
assert cube.z[-1] == 1.5