working with nerfstudio 1.0 viewer

javieryu · Jan 24, 2024 · f40b0da · f40b0da
1 parent 2f089cc
commit f40b0da
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/nerfbridge/ros_trainer.py b/nerfbridge/ros_trainer.py
@@ -78,7 +78,7 @@ def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None:
                 train_lock=self.train_lock,
                 share=self.config.viewer.make_share_url,
             )
-            banner_messages = [f"Viewer at: {self.viewer_state.viewer_url}"]
+            banner_messages = self.viewer_state.viewer_info
 
         self._check_viewer_warnings()
 

diff --git a/nerfbridge/ros_viewer.py b/nerfbridge/ros_viewer.py
@@ -1,4 +1,4 @@
-from typing import Optional, Dict, Literal, List
+from typing import Optional, Dict, Literal, Set
 
 import numpy as np
 
@@ -8,6 +8,7 @@
 
 from nerfstudio.viewer.viewer import Viewer
 from nerfstudio.data.datasets.base_dataset import InputDataset
+import pdb
 
 VISER_NERFSTUDIO_SCALE_RATIO: float = 10.0
 
@@ -32,7 +33,7 @@ def init_scene(
         # draw the training cameras and images
         self.camera_handles: Dict[int, viser.CameraFrustumHandle] = {}
         self.original_c2w: Dict[int, np.ndarray] = {}
-        self.cameras_drawn: List[int] = []
+        self.cameras_drawn: Set[int] = set()
         self.dataset = train_dataset
         image_indices = self._pick_drawn_image_idxs(len(train_dataset))
         cameras = train_dataset.cameras.to("cpu")
@@ -49,7 +50,7 @@ def init_scene(
                 image=None,
                 wxyz=R.wxyz,
                 position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,
-                visible=True,
+                visible=False,
             )
 
             @camera_handle.on_click
@@ -70,10 +71,10 @@ def update_camera_poses(self):
         """Updates the camera poses in the scene."""
         image_indices = self.dataset.updated_indices
         for idx in image_indices:
-            if not idx in self.cameras_drawn:
+            if not idx in self.cameras_drawn and idx in self.camera_handles:
                 self.original_c2w[idx] = (
                     self.dataset.cameras.camera_to_worlds[idx].cpu().numpy()
                 )
                 self.camera_handles[idx].visible = True
-                self.cameras_drawn.append(idx)
+                self.cameras_drawn.add(idx)
         super().update_camera_poses()