viewer now uses dummy camera frustums, new logic to detect mocking (needs reviewing), 3 TODOs

Author: AntonioMacaronio

nerfstudio/scripts/viewer/run_viewer.py

@@ -87,8 +87,18 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int):
         pipeline: Pipeline instance of which to load weights
         step: Step at which the pipeline was saved
     """
-    base_dir = config.get_base_dir()
-    viewer_log_path = base_dir / config.viewer.relative_log_filename
+    # Check if we're using a shared checkpoint (load_dir is outside the standard experiment structure)
+    try:
+        base_dir = config.get_base_dir()
+        # If get_base_dir() would create a path that doesn't exist, we're likely using a shared checkpoint
+        if not base_dir.parent.exists():
+            # Use the checkpoint directory as the base for log files
+            viewer_log_path = config.load_dir / config.viewer.relative_log_filename
+        else:
+            viewer_log_path = base_dir / config.viewer.relative_log_filename
+    except (FileNotFoundError, OSError):
+        # Fallback to using the checkpoint directory for shared checkpoints
+        viewer_log_path = config.load_dir / config.viewer.relative_log_filename
     banner_messages = None
     viewer_state = None
     viewer_callback_lock = Lock()

nerfstudio/utils/eval_utils.py

@@ -49,6 +49,7 @@ def patch_config_for_mock_data(config: TrainerConfig) -> TrainerConfig:
     datamanager_data = getattr(config.pipeline.datamanager, 'data', None)
 
     # The dataparser will use its own data field if it's meaningful, otherwise it inherits from datamanager
+    # TODO: this is a hack, but I really need to change this
     if dataparser_data and str(dataparser_data) != "." and dataparser_data.name != "":
         actual_data_path = dataparser_data
     else:

nerfstudio/viewer/viewer.py

@@ -444,45 +444,76 @@ def init_scene(
         # draw the training cameras and images
         self.camera_handles: Dict[int, viser.CameraFrustumHandle] = {}
         self.original_c2w: Dict[int, np.ndarray] = {}
-        image_indices = self._pick_drawn_image_idxs(len(train_dataset))
-        for idx in image_indices:
-            image = train_dataset[idx]["image"]
-            camera = train_dataset.cameras[idx]
-            image_uint8 = (image * 255).detach().type(torch.uint8)
-            image_uint8 = image_uint8.permute(2, 0, 1)
-
-            # torchvision can be slow to import, so we do it lazily.
-            import torchvision
-
-            image_uint8 = torchvision.transforms.functional.resize(image_uint8, 100, antialias=None)  # type: ignore
-            image_uint8 = image_uint8.permute(1, 2, 0)
-            image_uint8 = image_uint8.cpu().numpy()
-            c2w = camera.camera_to_worlds.cpu().numpy()
-            R = vtf.SO3.from_matrix(c2w[:3, :3])
-            R = R @ vtf.SO3.from_x_radians(np.pi)
-            camera_handle = self.viser_server.scene.add_camera_frustum(
-                name=f"/cameras/camera_{idx:05d}",
-                fov=float(2 * np.arctan((camera.cx / camera.fx[0]).cpu())),
-                scale=self.config.camera_frustum_scale,
-                aspect=float((camera.cx[0] / camera.cy[0]).cpu()),
-                image=image_uint8,
-                wxyz=R.wxyz,
-                position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,
-            )
+        
+        # Check if we're using mock data - if so, skip image loading to avoid file errors
+        is_mock_data = (
+            len(train_dataset) > 0 and 
+            hasattr(train_dataset, '_dataparser_outputs') and
+            len(train_dataset._dataparser_outputs.image_filenames) > 0 and
+            str(train_dataset._dataparser_outputs.image_filenames[0]).startswith("mock_image")
+        )
+        
+        if is_mock_data:
+            # For mock data, just draw camera frustums without images
+            image_indices = self._pick_drawn_image_idxs(len(train_dataset))
+            for idx in image_indices:
+                camera = train_dataset.cameras[idx]
+                c2w = camera.camera_to_worlds.cpu().numpy()
+                R = vtf.SO3.from_matrix(c2w[:3, :3])
+                R = R @ vtf.SO3.from_x_radians(np.pi)
+                camera_handle = self.viser_server.scene.add_camera_frustum(
+                    name=f"/cameras/camera_{idx:05d}",
+                    fov=2 * np.arctan(camera.height / (2 * camera.fy)).item(),
+                    aspect=camera.width / camera.height,
+                    scale=0.1,
+                    color=(255, 255, 255),
+                    wxyz=R.wxyz,
+                    position=c2w[:3, 3],
+                    visible=False,
+                )
+                self.camera_handles[idx] = camera_handle
+                self.original_c2w[idx] = c2w
+        else:
+            # Normal image loading for real data
+            image_indices = self._pick_drawn_image_idxs(len(train_dataset))
+            for idx in image_indices:
+                image = train_dataset[idx]["image"]
+                camera = train_dataset.cameras[idx]
+                image_uint8 = (image * 255).detach().type(torch.uint8)
+                image_uint8 = image_uint8.permute(2, 0, 1)
+
+                # torchvision can be slow to import, so we do it lazily.
+                import torchvision
+
+                image_uint8 = torchvision.transforms.functional.resize(image_uint8, 100, antialias=None)  # type: ignore
+                image_uint8 = image_uint8.permute(1, 2, 0)
+                image_uint8 = image_uint8.cpu().numpy()
+                c2w = camera.camera_to_worlds.cpu().numpy()
+                R = vtf.SO3.from_matrix(c2w[:3, :3])
+                R = R @ vtf.SO3.from_x_radians(np.pi)
+                camera_handle = self.viser_server.scene.add_camera_frustum(
+                    name=f"/cameras/camera_{idx:05d}",
+                    fov=float(2 * np.arctan((camera.cx / camera.fx[0]).cpu())),
+                    scale=self.config.camera_frustum_scale,
+                    aspect=float((camera.cx[0] / camera.cy[0]).cpu()),
+                    image=image_uint8,
+                    wxyz=R.wxyz,
+                    position=c2w[:3, 3] * VISER_NERFSTUDIO_SCALE_RATIO,
+                )
 
-            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
+                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
+                    return on_click_callback
 
-            camera_handle.on_click(create_on_click_callback(idx))
+                camera_handle.on_click(create_on_click_callback(idx))
 
-            self.camera_handles[idx] = camera_handle
-            self.original_c2w[idx] = c2w
+                self.camera_handles[idx] = camera_handle
+                self.original_c2w[idx] = c2w
 
         self.train_state = train_state
         self.train_util = 0.9