cleanup, passing final ns-dev-test

Author: AntonioMacaronio

nerfstudio/data/datamanagers/base_datamanager.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections import defaultdict
 from dataclasses import dataclass, field
 from functools import cached_property
 from pathlib import Path
@@ -55,44 +54,19 @@
 from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.pixel_samplers import PatchPixelSamplerConfig, PixelSampler, PixelSamplerConfig
-from nerfstudio.data.utils.dataloaders import CacheDataloader, FixedIndicesEvalDataloader, RandIndicesEvalDataloader
+from nerfstudio.data.utils.dataloaders import (
+    CacheDataloader,
+    FixedIndicesEvalDataloader,
+    RandIndicesEvalDataloader,
+    variable_res_collate,
+)
 from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper, get_orig_class
 from nerfstudio.utils.rich_utils import CONSOLE
 
 
-def variable_res_collate(batch: List[Dict]) -> Dict:
-    """Default collate function for the cached dataloader.
-    Args:
-        batch: Batch of samples from the dataset.
-    Returns:
-        Collated batch.
-    """
-    images = []
-    imgdata_lists = defaultdict(list)
-    for data in batch:
-        image = data.pop("image")
-        images.append(image)
-        topop = []
-        for key, val in data.items():
-            if isinstance(val, torch.Tensor):
-                # if the value has same height and width as the image, assume that it should be collated accordingly.
-                if len(val.shape) >= 2 and val.shape[:2] == image.shape[:2]:
-                    imgdata_lists[key].append(val)
-                    topop.append(key)
-        # now that iteration is complete, the image data items can be removed from the batch
-        for key in topop:
-            del data[key]
-
-    new_batch = nerfstudio_collate(batch)
-    new_batch["image"] = images
-    new_batch.update(imgdata_lists)
-
-    return new_batch
-
-
 @dataclass
 class DataManagerConfig(InstantiateConfig):
     """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
@@ -305,8 +279,6 @@ def get_param_groups(self) -> Dict[str, List[Parameter]]:
 
 @dataclass
 class VanillaDataManagerConfig(DataManagerConfig):
-    """A basic data manager for a ray-based model"""
-
     _target: Type = field(default_factory=lambda: VanillaDataManager)
     """Target class to instantiate."""
     dataparser: AnnotatedDataParserUnion = field(default_factory=BlenderDataParserConfig)

nerfstudio/data/utils/dataloaders.py

@@ -524,14 +524,16 @@ def __iter__(self):
         """This implementation allows every worker only cache the indices of the images they will use to generate rays to conserve RAM memory."""
         worker_info = get_worker_info()
         if worker_info is not None:  # if we have multiple processes
-            per_worker = int(math.ceil(len(self.input_dataset) / float(worker_info.num_workers)))
-            slice_start = worker_info.id * per_worker
-        else:  # we only have a single process
-            per_worker = len(self.input_dataset)
-            slice_start = 0
-        dataset_indices = list(range(len(self.input_dataset)))
-        # the indices of the datapoints in the dataset this worker will load
-        worker_indices = dataset_indices[slice_start : slice_start + per_worker]
+            if len(self.input_dataset) < worker_info.num_workers:
+                # if there's fewer datapoints than workers, each worker receives all datapoints
+                worker_indices = list(range(len(self.input_dataset)))
+            else:
+                per_worker = int(math.ceil(len(self.input_dataset) / float(worker_info.num_workers)))
+                slice_start = worker_info.id * per_worker
+                dataset_indices = list(range(len(self.input_dataset)))
+                worker_indices = dataset_indices[slice_start : slice_start + per_worker]
+        else:  # if we only have a single process
+            worker_indices = list(range(len(self.input_dataset)))
         if not self.load_from_disk:
             self._cached_collated_batch = self._get_collated_batch(worker_indices)
         r = random.Random(3301)
@@ -549,7 +551,6 @@ def __iter__(self):
                 collated_batch = self._cached_collated_batch
             elif i % self.num_times_to_repeat_images == 0:
                 r.shuffle(worker_indices)
-
                 if self.num_images_to_sample_from == -1:
                     # if -1, the worker gets all available indices in its partition
                     image_indices = worker_indices
@@ -562,10 +563,12 @@ def __iter__(self):
             """
             Here, the variable 'batch' refers to the output of our pixel sampler.
                 - batch is a dict_keys(['image', 'indices'])
-                - batch['image'] returns a pytorch tensor with shape `torch.Size([4096, 3])` , where 4096 = num_rays_per_batch. Note: each row in this tensor represents the RGB values as floats in [0, 1] of the pixel the ray goes through. The info of what specific image index that pixel belongs to is stored within batch[’indices’]
-                - batch['indices'] returns a pytorch tensor `torch.Size([4096, 3])` tensor where each row represents (image_index=camera_index, pixelRow, pixelCol)
-            What the pixel_sampler does (for variable_res_collate) is that it loops though each image, samples pixel within the mask, 
-            and returns them as the variable `indices` which has shape torch.Size([4096, 3]), where each row represents a pixel (image_idx, pixelRow, pixelCol)
+                - batch['image'] returns a pytorch tensor with shape `torch.Size([4096, 3])` , where 4096 = num_rays_per_batch. 
+                    - Note: each row in this tensor represents the RGB values as floats in [0, 1] of the pixel the ray goes through. 
+                    - The info of what specific image index that pixel belongs to is stored within batch[’indices’]
+                - batch['indices'] returns a pytorch tensor `torch.Size([4096, 3])` tensor where each row represents (image_idx, pixelRow, pixelCol)
+            pixel_sampler (for variable_res_collate) will loop though each image, samples pixel within the mask, and returns 
+            them as the variable `indices` which has shape torch.Size([4096, 3]), where each row represents a pixel (image_idx, pixelRow, pixelCol)
             """
             batch = worker_pixel_sampler.sample(collated_batch)  # type: ignore
             # collated_batch["image"].get_device() will return CPU if self.exclude_batch_keys_from_device contains 'image'
@@ -632,9 +635,9 @@ def __iter__(self):
 
             i += 1
             camera = camera.to(self.device)
-            for k in data.keys():
-                if isinstance(data[k], torch.Tensor):
-                    data[k] = data[k].to(self.device)
+            # for k in data.keys():
+            #     if isinstance(data[k], torch.Tensor):
+            #         data[k] = data[k].to(self.device)
             yield camera, data