Add control for resampling images (#847)

nerfstudio/data/datamanagers.py

@@ -15,6 +15,7 @@
 """
 Data loader.
 """
+
 from __future__ import annotations
 
 from abc import abstractmethod
@@ -23,6 +24,7 @@ from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch
 import tyro
+from rich.progress import Console
 from torch import nn
 from torch.nn import Parameter
 from torch.utils.data import Dataset
@@ -49,6 +51,8 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
 
+CONSOLE = Console(width=120)
+
 AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of flags in subcommands.
     tyro.extras.subcommand_type_from_defaults(
         {
@@ -242,10 +246,16 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """Number of rays per batch to use per training iteration."""
     train_num_images_to_sample_from: int = -1
     """Number of images to sample during training iteration."""
+    train_num_times_to_repeat_images: int = -1
+    """When not training on all images, number of iterations before picking new
+    images. If -1, never pick new images."""
     eval_num_rays_per_batch: int = 1024
     """Number of rays per batch to use per eval iteration."""
     eval_num_images_to_sample_from: int = -1
     """Number of images to sample during eval iteration."""
+    eval_num_times_to_repeat_images: int = -1
+    """When not evaluating on all images, number of iterations before picking
+    new images. If -1, never pick new images."""
     eval_image_indices: Optional[Tuple[int, ...]] = (0,)
     """Specifies the image indices to use during eval; if None, uses all."""
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
@@ -294,9 +304,11 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     def setup_train(self):
         """Sets up the data loaders for training"""
         assert self.train_dataset is not None
+        CONSOLE.print("Setting up training dataset...")
         self.train_image_dataloader = CacheDataloader(
             self.train_dataset,
             num_images_to_sample_from=self.config.train_num_images_to_sample_from,
+            num_times_to_repeat_images=self.config.train_num_times_to_repeat_images,
             device=self.device,
             num_workers=self.world_size * 4,
             pin_memory=True,
@@ -314,9 +326,11 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
     def setup_eval(self):
         """Sets up the data loader for evaluation"""
         assert self.eval_dataset is not None
+        CONSOLE.print("Setting up evaluation dataset...")
         self.eval_image_dataloader = CacheDataloader(
             self.eval_dataset,
             num_images_to_sample_from=self.config.eval_num_images_to_sample_from,
+            num_times_to_repeat_images=self.config.eval_num_times_to_repeat_images,
             device=self.device,
             num_workers=self.world_size * 4,
             pin_memory=True,
@@ -340,9 +354,6 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             num_workers=self.world_size * 4,
         )
 
-        # TODO: eval dataloader should be separate from train
-        self.iter_eval_dataloader = iter(self.eval_image_dataloader)
-
     def next_train(self, step: int) -> Tuple[RayBundle, Dict]:
         """Returns the next batch of data from the train dataloader."""
         self.train_count += 1

nerfstudio/data/dataparsers/record3d_dataparser.py

@@ -51,8 +51,8 @@ class Record3DDataParserConfig(DataParserConfig):
     orientation_method: Literal["pca", "up"] = "up"
     """The method to use for orientation"""
     max_dataset_size: int = 300
-    """Max number of images to train on. If the dataset has
-    more, images will be sampled approximately evenly."""
+    """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1,
+    use all images."""
 
 
 @dataclass
@@ -86,7 +86,7 @@ class Record3D(DataParser):
             axis=-1,
         ).astype(np.float32)
 
-        if num_images > self.config.max_dataset_size:
+        if self.config.max_dataset_size != -1 and num_images > self.config.max_dataset_size:
             # Evenly select max_dataset_size images from dataset, including first
             # and last indices.
             idx = np.round(np.linspace(0, num_images - 1, self.config.max_dataset_size)).astype(int)

nerfstudio/data/utils/dataloaders.py

@@ -21,7 +21,7 @@ from abc import abstractmethod
 from typing import Dict, Optional, Tuple, Union
 
 import torch
-from rich.progress import track
+from rich.progress import Console, track
 from torch.utils.data import Dataset, default_collate
 from torch.utils.data.dataloader import DataLoader
 
@@ -30,6 +30,8 @@ from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.utils.datasets import InputDataset
 from nerfstudio.utils.misc import get_dict_to_torch
 
+CONSOLE = Console(width=120)
+
 
 class CacheDataloader(DataLoader):
     """Collated image dataset that implements caching of default-pytorch-collatable data.
@@ -38,7 +40,7 @@ class CacheDataloader(DataLoader):
     Args:
         dataset: Dataset to sample from.
         num_samples_to_collate: How many images to sample rays for each batch. -1 for all images.
-        num_times_to_repeat_images: How often to collate new images.
+        num_times_to_repeat_images: How often to collate new images. -1 to never pick new images.
         device: Device to perform computation.
     """
 
@@ -46,13 +48,13 @@ class CacheDataloader(DataLoader):
         self,
         dataset: Dataset,
         num_images_to_sample_from: int = -1,
-        num_times_to_repeat_images: int = 0,
+        num_times_to_repeat_images: int = -1,
         device: Union[torch.device, str] = "cpu",
         **kwargs,
     ):
         self.dataset = dataset
         self.num_times_to_repeat_images = num_times_to_repeat_images
-        self.cache_all_images = num_images_to_sample_from == -1
+        self.cache_all_images = (num_images_to_sample_from == -1) or (num_images_to_sample_from >= len(self.dataset))
         self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from
         self.device = device
 
@@ -61,7 +63,21 @@ class CacheDataloader(DataLoader):
 
         self.cached_collated_batch = None
         if self.cache_all_images:
+            CONSOLE.print(f"Caching all {len(self.dataset)} images.")
+            if len(self.dataset) > 500:
+                CONSOLE.print(
+                    "[bold yellow]Warning: If you run out of memory, try reducing the number of images to sample from."
+                )
             self.cached_collated_batch = self._get_collated_batch()
+        elif self.num_times_to_repeat_images == -1:
+            CONSOLE.print(
+                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, without resampling."
+            )
+        else:
+            CONSOLE.print(
+                f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, "
+                f"resampling every {self.num_times_to_repeat_images} iters."
+            )
         super().__init__(dataset=dataset, **kwargs)
 
     def __getitem__(self, idx):