diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 6355a5ba6..ce4cbb132 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -21,6 +21,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args import torch +from viser import ClientHandle 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 @@ -52,7 +53,7 @@ class RenderStateMachine(threading.Thread): viewer: the viewer state """ - def __init__(self, viewer: Viewer, viser_scale_ratio: float): + def __init__(self, viewer: Viewer, viser_scale_ratio: float, client: ClientHandle): threading.Thread.__init__(self) self.transitions: Dict[RenderStates, Dict[RenderActions, RenderStates]] = { s: {} for s in get_args(RenderStates) @@ -77,6 +78,8 @@ class RenderStateMachine(threading.Thread): self.daemon = True self.output_keys = {} self.viser_scale_ratio = viser_scale_ratio + self.client = client + self.running = True def action(self, action: RenderAction): """Takes an action and updates the state machine @@ -107,7 +110,7 @@ class RenderStateMachine(threading.Thread): self.render_trigger.set() def _render_img(self, camera_state: CameraState): - """Takes the current camera, generates rays, and renders the iamge + """Takes the current camera, generates rays, and renders the image Args: camera_state: the current camera state @@ -166,7 +169,7 @@ class RenderStateMachine(threading.Thread): def run(self): """Main loop for the render thread""" - while True: + while self.running: 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: @@ -237,7 +240,8 @@ class RenderStateMachine(threading.Thread): depth = ( outputs["gl_z_buf_depth"].cpu().numpy() * self.viser_scale_ratio if "gl_z_buf_depth" in outputs else None ) - self.viewer.viser_server.set_background_image( + + self.client.set_background_image( selected_output.cpu().numpy(), format=self.viewer.config.image_format, jpeg_quality=self.viewer.config.jpeg_quality, diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index cc352d553..41b142aea 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -103,8 +103,8 @@ class Viewer: self.step = 0 self.train_btn_state: Literal["training", "paused", "completed"] = "training" self._prev_train_state: Literal["training", "paused", "completed"] = "training" + self.last_move_time = 0 - self.client: Optional[viser.ClientHandle] = None self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port, share=share) buttons = ( viser.theme.TitlebarButton( @@ -137,7 +137,8 @@ class Viewer: brand_color=(255, 211, 105), ) - self.render_statemachine = RenderStateMachine(self, VISER_NERFSTUDIO_SCALE_RATIO) + self.render_statemachines: Dict[int, RenderStateMachine] = {} + self.viser_server.on_client_disconnect(self.handle_disconnect) self.viser_server.on_client_connect(self.handle_new_client) tabs = self.viser_server.add_gui_tab_group() @@ -202,24 +203,30 @@ class Viewer: ] for c in self.viewer_controls: c._setup(self) - self.render_statemachine.start() + + def get_camera_state(self, client: viser.ClientHandle) -> CameraState: + R = vtf.SO3(wxyz=client.camera.wxyz) + R = R @ vtf.SO3.from_x_radians(np.pi) + R = torch.tensor(R.as_matrix()) + pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO + c2w = torch.concatenate([R, pos[:, None]], dim=1) + camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w) + return camera_state + + def handle_disconnect(self, client: viser.ClientHandle) -> None: + self.render_statemachines[client.client_id].running = False + self.render_statemachines.pop(client.client_id) def handle_new_client(self, client: viser.ClientHandle) -> None: - self.client = client - self.last_move_time = 0 + self.render_statemachines[client.client_id] = RenderStateMachine(self, VISER_NERFSTUDIO_SCALE_RATIO, client) + self.render_statemachines[client.client_id].start() @client.camera.on_update def _(cam: viser.CameraHandle) -> None: - assert self.client is not None - with client.atomic(): - 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()) - 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=self.client.camera.fov, aspect=self.client.camera.aspect, c2w=c2w) - self.render_statemachine.action(RenderAction("move", self.camera_state)) + self.last_move_time = time.time() + with self.viser_server.atomic(): + camera_state = self.get_camera_state(client) + self.render_statemachines[client.client_id].action(RenderAction("move", camera_state)) def set_camera_visibility(self, visible: bool) -> None: """Toggle the visibility of the training cameras.""" @@ -250,8 +257,11 @@ class Viewer: def _interrupt_render(self, _) -> None: """Interrupt current render.""" - if self.camera_state is not None: - self.render_statemachine.action(RenderAction("rerender", self.camera_state)) + clients = self.viser_server.get_clients() + for id in clients: + camera_state = self.get_camera_state(clients[id]) + if camera_state is not None: + self.render_statemachines[id].action(RenderAction("rerender", camera_state)) def _toggle_training_state(self, _) -> None: """Toggle the trainer's training state.""" @@ -263,8 +273,11 @@ class Viewer: def _crop_params_update(self, _) -> None: """Update crop parameters""" - if self.camera_state is not None: - self.render_statemachine.action(RenderAction("move", self.camera_state)) + clients = self.viser_server.get_clients() + for id in clients: + camera_state = self.get_camera_state(clients[id]) + if camera_state is not None: + self.render_statemachines[id].action(RenderAction("move", camera_state)) def _output_type_change(self, _): self.output_type_changed = True @@ -327,10 +340,9 @@ class Viewer: @camera_handle.on_click def _(event: viser.SceneNodePointerEvent[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 + with event.client.atomic(): + event.client.camera.position = event.target.position + event.client.camera.wxyz = event.target.wxyz self.camera_handles[idx] = camera_handle self.original_c2w[idx] = c2w @@ -347,7 +359,7 @@ class Viewer: """ self.step = step - if self.camera_state is None: + if len(self.render_statemachines) == 0: return # this stops training while moving to make the response smoother while time.time() - self.last_move_time < 0.1: @@ -371,7 +383,11 @@ class Viewer: render_freq = 30 if step > self.last_step + render_freq: self.last_step = step - self.render_statemachine.action(RenderAction("step", self.camera_state)) + clients = self.viser_server.get_clients() + for id in clients: + camera_state = self.get_camera_state(clients[id]) + if camera_state is not None: + self.render_statemachines[id].action(RenderAction("step", camera_state)) self.update_camera_poses() self.control_panel.update_step(step) diff --git a/nerfstudio/viewer_beta/viewer_elements.py b/nerfstudio/viewer_beta/viewer_elements.py index 916566d5d..82a8b3d44 100644 --- a/nerfstudio/viewer_beta/viewer_elements.py +++ b/nerfstudio/viewer_beta/viewer_elements.py @@ -117,7 +117,7 @@ class ViewerControl: """ raise NotImplementedError() - def get_camera(self, img_height: int, img_width: int) -> Optional[Cameras]: + def get_camera(self, img_height: int, img_width: int, client_id: Optional[int] = None) -> Optional[Cameras]: """ Returns the Cameras object representing the current camera for the viewer, or None if the viewer is not connected yet @@ -126,16 +126,21 @@ class ViewerControl: 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 + clients = self.viser_server.get_clients() + if len(clients) == 0: + return None + if not client_id: + client_id = list(clients.keys())[0] from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO - R = vtf.SO3(wxyz=self.viewer.client.camera.wxyz) + client = clients[client_id] + R = vtf.SO3(wxyz=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 + pos = torch.tensor(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO c2w = torch.concatenate([R, pos[:, None]], dim=1) - camera_state = CameraState(fov=self.viewer.client.camera.fov, aspect=self.viewer.client.camera.aspect, c2w=c2w) + camera_state = CameraState(fov=client.camera.fov, aspect=client.camera.aspect, c2w=c2w) return get_camera(camera_state, img_height, img_width) def register_click_cb(self, cb: Callable):