Limit number of images in viewer (#845)

nerfstudio/configs/base_config.py

@@ -175,6 +175,9 @@ class ViewerConfig(PrintableConfig):
     """the ip address where the bridge server is running"""
     num_rays_per_chunk: int = 32768
     """number of rays per chunk to render with viewer"""
+    max_num_display_images: int = 512
+    """Maximum number of training images to display in the viewer, to avoid lag. This does not change which images are
+    actually used in training/evaluation. If -1, display all."""
 
 
 from nerfstudio.engine.optimizers import OptimizerConfig

nerfstudio/viewer/server/viewer_utils.py

@@ -274,6 +274,22 @@ class ViewerState:
         self.webrtc_thread = None
         self.kill_webrtc_signal = False
 
+    def _pick_drawn_image_idxs(self, total_num: int) -> list[int]:
+        """Determine indicies of images to display in viewer.
+
+        Args:
+            total_num: total number of training images.
+
+        Returns:
+            List of indices from [0, total_num-1].
+        """
+        if self.config.max_num_display_images < 0:
+            num_display_images = total_num
+        else:
+            num_display_images = min(self.config.max_num_display_images, total_num)
+        # draw indices, roughly evenly spaced
+        return np.linspace(0, total_num - 1, num_display_images, dtype=np.int32).tolist()
+
     def init_scene(self, dataset: InputDataset, start_train=True) -> None:
         """Draw some images and the scene aabb in the viewer.
 
@@ -290,7 +306,7 @@ class ViewerState:
         self.vis["sceneState/cameras"].delete()
 
         # draw the training cameras and images
-        image_indices = range(len(dataset))
+        image_indices = self._pick_drawn_image_idxs(len(dataset))
         for idx in image_indices:
             image = dataset[idx]["image"]
             bgr = image[..., [2, 1, 0]]