diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index 3f1aff36a..340806577 100644 --- a/nerfstudio/configs/base_config.py +++ b/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 diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 70848d08d..4824262b7 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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]]