From 3e1d4af8217134fc34b42690345c5847147bc732 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Tue, 19 Sep 2023 16:22:55 -0700 Subject: [PATCH] Add share link support for beta viewer (#2445) * Add share link support for beta viewer * More consistent CLI experience * fix viewer jitter by adding atomic update on camera pose * change to make_share_url * fix * remove incorrect assert --------- Co-authored-by: Justin Kerr --- nerfstudio/configs/base_config.py | 2 ++ nerfstudio/configs/experiment_config.py | 10 ++++++++-- nerfstudio/engine/trainer.py | 1 + nerfstudio/scripts/viewer/run_viewer.py | 8 ++++++-- nerfstudio/viewer_beta/viewer.py | 25 ++++++++++++++----------- pyproject.toml | 2 +- 6 files changed, 32 insertions(+), 16 deletions(-) diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 315659acb..fc906fae9 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -146,3 +146,5 @@ class ViewerConfig(PrintableConfig): """Image format viewer should use; jpeg is lossy compression, while png is lossless.""" jpeg_quality: int = 90 """Quality tradeoff to use for jpeg compression.""" + make_share_url: bool = False + """Viewer beta feature: print a shareable URL. `vis` must be set to viewer_beta; this flag is otherwise ignored.""" diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 0da8893ed..42cc4399e 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -22,7 +22,13 @@ from pathlib import Path from typing import Any, Dict, Literal, Optional import yaml -from nerfstudio.configs.base_config import InstantiateConfig, LoggingConfig, MachineConfig, ViewerConfig + +from nerfstudio.configs.base_config import ( + InstantiateConfig, + LoggingConfig, + MachineConfig, + ViewerConfig, +) from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.engine.optimizers import OptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig @@ -73,7 +79,7 @@ class ExperimentConfig(InstantiateConfig): relative_model_dir: Path = Path("nerfstudio_models/") """Relative path to save all checkpoints.""" load_scheduler: bool = True - """Whether to load the scheduler state_dict to resume training, if exists""" + """Whether to load the scheduler state_dict to resume training, if it exists.""" def is_viewer_enabled(self) -> bool: """Checks if a viewer is enabled.""" diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index c8ae50043..04fa8e05b 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -177,6 +177,7 @@ class Trainer: pipeline=self.pipeline, trainer=self, train_lock=self.train_lock, + share=self.config.viewer.make_share_url, ) banner_messages = [f"Viewer Beta at: {self.viewer_state.viewer_url}"] self._check_viewer_warnings() diff --git a/nerfstudio/scripts/viewer/run_viewer.py b/nerfstudio/scripts/viewer/run_viewer.py index 56a712026..de210a30d 100644 --- a/nerfstudio/scripts/viewer/run_viewer.py +++ b/nerfstudio/scripts/viewer/run_viewer.py @@ -55,6 +55,8 @@ class RunViewer: """Viewer configuration""" vis: Literal["viewer", "viewer_beta"] = "viewer" """Type of viewer""" + make_share_url: bool = False + """Viewer beta feature: print a shareable URL. `vis` must be set to viewer_beta; this flag is otherwise ignored.""" def main(self) -> None: """Main function.""" @@ -66,6 +68,7 @@ class RunViewer: num_rays_per_chunk = config.viewer.num_rays_per_chunk assert self.viewer.num_rays_per_chunk == -1 config.vis = self.vis + config.viewer.make_share_url = self.make_share_url config.viewer = self.viewer.as_viewer_config() config.viewer.num_rays_per_chunk = num_rays_per_chunk @@ -103,6 +106,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): log_filename=viewer_log_path, datapath=base_dir, pipeline=pipeline, + share=config.viewer.make_share_url, ) banner_messages = [f"Viewer Beta at: {viewer_state.viewer_url}"] @@ -126,11 +130,11 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): def entrypoint(): """Entrypoint for use with pyproject scripts.""" tyro.extras.set_accent_color("bright_yellow") - tyro.cli(RunViewer).main() + tyro.cli(tyro.conf.FlagConversionOff[RunViewer]).main() if __name__ == "__main__": entrypoint() # For sphinx docs -get_parser_fn = lambda: tyro.extras.get_parser(RunViewer) # noqa +get_parser_fn = lambda: tyro.extras.get_parser(tyro.conf.FlagConversionOff[RunViewer]) # noqa diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index 9e2314ad4..d2a53b198 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -26,6 +26,8 @@ import torchvision import viser import viser.theme import viser.transforms as vtf + +from nerfstudio.cameras.camera_optimizers import CameraOptimizer from nerfstudio.configs import base_config as cfg from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.models.base_model import Model @@ -40,8 +42,6 @@ from nerfstudio.viewer_beta.render_state_machine import RenderAction, RenderStat 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 @@ -59,6 +59,7 @@ class Viewer: datapath: path to data pipeline: pipeline object to use trainer: trainer object to use + share: print a shareable URL Attributes: viewer_url: url to open viewer @@ -77,6 +78,7 @@ class Viewer: pipeline: Pipeline, trainer: Optional[Trainer] = None, train_lock: Optional[threading.Lock] = None, + share: bool = False, ): self.config = config self.trainer = trainer @@ -103,7 +105,7 @@ class Viewer: self._prev_train_state: Literal["training", "paused", "completed"] = "training" self.client: Optional[viser.ClientHandle] = None - self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port) + self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port, share=share) buttons = ( viser.theme.TitlebarButton( text="Getting Started", @@ -191,14 +193,15 @@ class Viewer: @client.camera.on_update 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()) - 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)) + 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)) def set_camera_visibility(self, visible: bool) -> None: """Toggle the visibility of the training cameras.""" diff --git a/pyproject.toml b/pyproject.toml index add21cf6d..c64152675 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ dependencies = [ "torchvision>=0.14.1", "torchmetrics[image]>=1.0.1", "typing_extensions>=4.4.0", - "viser==0.1.2", + "viser==0.1.3", "nuscenes-devkit>=1.1.1", "wandb>=0.13.3", "xatlas",