From 9db97773b93a42b8bcff74d7f2331d19fe0f590f Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Jun 2023 05:03:06 +0000 Subject: [PATCH 1/4] fix render statemachine logic for viewer beta --- nerfstudio/viewer_beta/render_state_machine.py | 10 +++++++++- nerfstudio/viewer_beta/viewer.py | 4 ++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index d18200761..6be9b3403 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -17,6 +17,7 @@ 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 @@ -97,12 +98,16 @@ class RenderStateMachine(threading.Thread): elif self.next_action == "rerender": # never overwrite rerenders pass + elif action.action == "static" and self.next_action.action == "move": + # don't overwrite a move action with a static: static is always self-fired + return else: # monimal use case, just set the next action self.next_action = action # 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() @@ -176,7 +181,10 @@ class RenderStateMachine(threading.Thread): 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.state == "low_static": + 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): diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index f4e8d3429..51ca6ab14 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -126,8 +126,8 @@ class Viewer: def _interrupt_render(self, _) -> None: """Interrupt current render.""" - 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("rerender", self.camera_state)) def _crop_params_update(self, _) -> None: """Update crop parameters""" From efd18266f724141a5cabe53aecd8cb94d1a220e7 Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 16 Jun 2023 06:05:38 +0000 Subject: [PATCH 2/4] clicking on cameras in the viewer --- nerfstudio/viewer_beta/viewer.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 51ca6ab14..413405368 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -95,8 +95,7 @@ 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 = None self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port) self.viser_server.on_client_connect(self.handle_new_client) @@ -114,14 +113,16 @@ class Viewer: self.render_statemachine.start() def handle_new_client(self, client: viser.ClientHandle) -> None: + self.client = client + @client.camera.on_update def _(_: viser.CameraHandle) -> None: - R = vtf.SO3(wxyz=client.camera.wxyz) + 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(client.camera.position, dtype=torch.float64) / VISER_NERFSTUDIO_SCALE_RATIO + 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 _interrupt_render(self, _) -> None: @@ -190,8 +191,8 @@ 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}", + camera_handle = self.viser_server.add_camera_frustum( + name=f"/camera_{idx:05d}", fov=float(camera.fx[0]), scale=0.2, aspect=float(camera.cx[0] / camera.cy[0]), @@ -200,6 +201,12 @@ class Viewer: position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO, ) + @camera_handle.on_click + def _(cam: viser.CameraHandle) -> None: + with self.client.atomic(): + self.client.camera.position = cam.position + self.client.camera.wxyz = cam.wxyz + self.train_state = train_state self.train_util = 0.9 From cb66e3149fb923b30d3545351b28cd936bf615c7 Mon Sep 17 00:00:00 2001 From: Justin Date: Sun, 18 Jun 2023 04:58:49 +0000 Subject: [PATCH 3/4] fix fov calculation, update camera positions while training --- nerfstudio/configs/method_configs.py | 2 +- nerfstudio/engine/trainer.py | 1 + .../viewer_beta/render_state_machine.py | 4 ++- nerfstudio/viewer_beta/viewer.py | 27 +++++++++++++++++-- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index aa5eb0f67..7c2253f0c 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -89,7 +89,7 @@ method_configs["nerfacto"] = TrainerConfig( eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", - optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=0), scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), ), ), diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index afc0d94f3..a9bb876c2 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -351,6 +351,7 @@ class Trainer: except RuntimeError: time.sleep(0.03) # sleep to allow buffer to reset CONSOLE.log("Viewer failed. Continuing training.") + raise @check_viewer_enabled def _train_complete_viewer(self) -> None: diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 6be9b3403..c56994a41 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -167,7 +167,8 @@ class RenderStateMachine(threading.Thread): self.render_trigger.wait() self.render_trigger.clear() action = self.next_action - assert action is not None, "Action should never be None at this point" + 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 @@ -180,6 +181,7 @@ class RenderStateMachine(threading.Thread): # if we got interrupted, don't send the output to the viewer continue self._send_output_to_viewer(outputs) + self.viewer.update_camera_poses() # 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 diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 413405368..af08bbcc3 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -125,6 +125,23 @@ class Viewer: 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 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()) + camera_optimizer = self.trainer.pipeline.datamanager.train_camera_optimizer + with torch.no_grad(): + 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]) + 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: @@ -180,6 +197,8 @@ class Viewer: # draw the training cameras and images image_indices = self._pick_drawn_image_idxs(len(dataset)) + self.camera_handles = {} + self.original_c2w = {} for idx in image_indices: image = dataset[idx]["image"] camera = dataset.cameras[idx] @@ -193,12 +212,13 @@ class Viewer: R = R @ vtf.SO3.from_x_radians(np.pi) camera_handle = self.viser_server.add_camera_frustum( name=f"/camera_{idx:05d}", - fov=float(camera.fx[0]), - scale=0.2, + fov=2 * np.arctan(camera.cx / camera.fx[0]), + scale=1, aspect=float(camera.cx[0] / camera.cy[0]), image=image_uint8, wxyz=R.wxyz, position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO, + jpeg_quality=50, ) @camera_handle.on_click @@ -207,6 +227,9 @@ class Viewer: self.client.camera.position = cam.position self.client.camera.wxyz = cam.wxyz + self.camera_handles[idx] = camera_handle + self.original_c2w[idx] = c2w + self.train_state = train_state self.train_util = 0.9 From 0a9eb8a7c9ab080452cb7491aeae73e9dabf711f Mon Sep 17 00:00:00 2001 From: Justin Date: Sun, 18 Jun 2023 05:01:07 +0000 Subject: [PATCH 4/4] cleanup --- nerfstudio/configs/method_configs.py | 2 +- nerfstudio/engine/trainer.py | 1 - nerfstudio/viewer_beta/render_state_machine.py | 1 - nerfstudio/viewer_beta/viewer.py | 1 - 4 files changed, 1 insertion(+), 4 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 7c2253f0c..aa5eb0f67 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -89,7 +89,7 @@ method_configs["nerfacto"] = TrainerConfig( eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", - optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=0), + optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), ), ), diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index a9bb876c2..afc0d94f3 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -351,7 +351,6 @@ class Trainer: except RuntimeError: time.sleep(0.03) # sleep to allow buffer to reset CONSOLE.log("Viewer failed. Continuing training.") - raise @check_viewer_enabled def _train_complete_viewer(self) -> None: diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index c56994a41..0b00f4dac 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -107,7 +107,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() diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index af08bbcc3..fdc0f21cc 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -218,7 +218,6 @@ class Viewer: image=image_uint8, wxyz=R.wxyz, position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO, - jpeg_quality=50, ) @camera_handle.on_click