Pass current camera idx in metadata to get output from viser (nerfstu…

…dio-project#3293)

allows viser camera to pass clicked training camera idx into metadata

nerfstudio/viewer/utils.py

@@ -42,6 +42,8 @@ class CameraState:
     """Type of camera to render."""
     time: float = 0.0
     """The rendering time of the camera state."""
+    idx: int = 0
+    """The index of the current camera."""
 
 
 def get_camera(
@@ -78,6 +80,7 @@ def get_camera(
         camera_type=camera_state.camera_type,
         camera_to_worlds=camera_state.c2w.to(torch.float32)[None, ...],
         times=torch.tensor([camera_state.time], dtype=torch.float32),
+        metadata={"cam_idx": camera_state.idx},
     )
     return camera
 

nerfstudio/viewer/viewer.py

@@ -106,6 +106,8 @@ def __init__(
         self.train_btn_state: Literal["training", "paused", "completed"] = "training"
         self._prev_train_state: Literal["training", "paused", "completed"] = "training"
         self.last_move_time = 0
+        # track the camera index that last being clicked
+        self.current_camera_idx = 0
 
         self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port)
         # Set the name of the URL either to the share link if available, or the localhost
@@ -325,13 +327,15 @@ def get_camera_state(self, client: viser.ClientHandle) -> CameraState:
                 else CameraType.EQUIRECTANGULAR
                 if camera_type == "Equirectangular"
                 else assert_never(camera_type),
+                idx=self.current_camera_idx,
             )
         else:
             camera_state = CameraState(
                 fov=client.camera.fov,
                 aspect=client.camera.aspect,
                 c2w=c2w,
                 camera_type=CameraType.PERSPECTIVE,
+                idx=self.current_camera_idx,
             )
         return camera_state
 
@@ -462,11 +466,16 @@ def init_scene(
                 position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,
             )
 
-            @camera_handle.on_click
-            def _(event: viser.SceneNodePointerEvent[viser.CameraFrustumHandle]) -> None:
-                with event.client.atomic():
-                    event.client.camera.position = event.target.position
-                    event.client.camera.wxyz = event.target.wxyz
+            def create_on_click_callback(capture_idx):
+                def on_click_callback(event: viser.SceneNodePointerEvent[viser.CameraFrustumHandle]) -> None:
+                    with event.client.atomic():
+                        event.client.camera.position = event.target.position
+                        event.client.camera.wxyz = event.target.wxyz
+                        self.current_camera_idx = capture_idx
+
+                return on_click_callback
+
+            camera_handle.on_click(create_on_click_callback(idx))
 
             self.camera_handles[idx] = camera_handle
             self.original_c2w[idx] = c2w