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 <justin.g.kerr@gmail.com>

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

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

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()

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

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

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