diff --git a/docs/developer_guides/pipelines/datamanagers.md b/docs/developer_guides/pipelines/datamanagers.md index 84722ba78..316514c64 100644 --- a/docs/developer_guides/pipelines/datamanagers.md +++ b/docs/developer_guides/pipelines/datamanagers.md @@ -14,7 +14,16 @@ ## What is a DataManager? -The DataManager returns RayBundle and RayGT objects. Let's first take a look at the most important abstract methods required by the DataManager. +The DataManager batches and returns two components from an input dataset: + +1. A representation of viewpoint (either cameras or rays). + - For splatting methods (`FullImageDataManager`): a `Cameras` object. + - For ray sampling methods (`VanillaDataManager`): a `RayBundle` object. +2. A dictionary of ground truth data. + - For splatting methods (`FullImageDataManager`): dictionary contains complete images. + - For ray sampling methods (`VanillaDataManager`): dictionary contains per-ray information. + +Behaviors are defined by implementing the abstract methods required by the DataManager: ```python class DataManager(nn.Module): @@ -22,11 +31,11 @@ class DataManager(nn.Module): """ @abstractmethod - def next_train(self, step: int) -> Tuple[RayBundle, Dict]: + def next_train(self, step: int) -> Tuple[Union[RayBundle, Cameras], Dict]: """Returns the next batch of data for train.""" @abstractmethod - def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: + def next_eval(self, step: int) -> Tuple[Union[RayBundle, Cameras], Dict]: """Returns the next batch of data for eval.""" @abstractmethod @@ -94,3 +103,98 @@ See the code! ## Creating Your Own We currently don't have other implementations because most papers follow the VanillaDataManager implementation. However, it should be straightforward to add a VanillaDataManager with logic that progressively adds cameras, for instance, by relying on the step and modifying RayBundle and RayGT generation logic. + +## Disk Caching for Large Datasets +As of January 2025, the FullImageDatamanager and ParallelImageDatamanager implementations now support parallelized dataloading and dataloading from disk to avoid Out-Of-Memory errors and support very large datasets. To train a NeRF-based method with a large dataset that's unable to fit in memory, please add the `load_from_disk` flag to your `ns-train` command. For example with nerfacto: +```bash +ns-train nerfacto --data {PROCESSED_DATA_DIR} --pipeline.datamanager.load-from-disk +``` + +To train splatfacto with a large dataset that's unable to fit in memory, please set the device of `cache_images` to `"disk"`. For example with splatfacto: +```bash +ns-train splatfacto --data {PROCESSED_DATA_DIR} --pipeline.datamanager.cache-images disk +``` + +## Migrating Your DataManager to the new DataManager +Many methods subclass a DataManager and add extra data to it. If you would like your custom datamanager to also support new parallel features, you can migrate any custom dataloading logic to the new `custom_ray_processor()` API. This function takes in a full training batch (either image or ray bundle) and allows the user to modify or add to it. Let's take a look at an example for the LERF method, which was built on Nerfstudio's VanillaDataManager. This API provides an interface to attach new information to the RayBundle (for ray based methods), Cameras object (for splatting based methods), or ground truth dictionary. It runs in a background process if disk caching is enabled, otherwise it runs in the main process. + +Naively transfering code to `custom_ray_processor` may still OOM on very large datasets if initialization code requires computing something over the whole dataset. To fully take advantage of parallelization make sure your subclassed datamanager computes new information inside the `custom_ray_processor`, or caches a subset of the whole dataset. This can also still be slow if pre-computation requires GPU-heavy steps on the same GPU used for training. + +**Note**: Because the parallel DataManager uses background processes, any member of the DataManager needs to be *picklable* to be used inside `custom_ray_processor`. + +```python +class LERFDataManager(VanillaDataManager): + """Subclass VanillaDataManager to add extra data processing + + Args: + config: the DataManagerConfig used to instantiate class + """ + + config: LERFDataManagerConfig + + def __init__( + self, + config: LERFDataManagerConfig, + device: Union[torch.device, str] = "cpu", + test_mode: Literal["test", "val", "inference"] = "val", + world_size: int = 1, + local_rank: int = 0, + **kwargs, + ): + super().__init__( + config=config, device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank, **kwargs + ) + # Some code to initialize all the CLIP and DINO feature encoders. + self.image_encoder: BaseImageEncoder = kwargs["image_encoder"] + self.dino_dataloader = ... + self.clip_interpolator = ... + + def next_train(self, step: int) -> Tuple[RayBundle, Dict]: + """Returns the next batch of data from the train dataloader. + + In this custom DataManager we need to add on the data that LERF needs, namely CLIP and DINO features. + """ + self.train_count += 1 + image_batch = next(self.iter_train_image_dataloader) + assert self.train_pixel_sampler is not None + batch = self.train_pixel_sampler.sample(image_batch) + ray_indices = batch["indices"] + ray_bundle = self.train_ray_generator(ray_indices) + batch["clip"], clip_scale = self.clip_interpolator(ray_indices) + batch["dino"] = self.dino_dataloader(ray_indices) + ray_bundle.metadata["clip_scales"] = clip_scale + # assume all cameras have the same focal length and image width + ray_bundle.metadata["fx"] = self.train_dataset.cameras[0].fx.item() + ray_bundle.metadata["width"] = self.train_dataset.cameras[0].width.item() + ray_bundle.metadata["fy"] = self.train_dataset.cameras[0].fy.item() + ray_bundle.metadata["height"] = self.train_dataset.cameras[0].height.item() + return ray_bundle, batch +``` + +To migrate this custom datamanager to the new datamanager, we'll subclass the new ParallelDataManager and shift the data customization process from `next_train()` to `custom_ray_processor()`. +The function `custom_ray_processor()` is called with a fully populated ray bundle and ground truth batch, just like the subclassed `next_train` in the above code. This code, however, is run in a background process. + +```python +class LERFDataManager(ParallelDataManager, Generic[TDataset]): + """ + __init__ stays the same + """ + + ... + + def custom_ray_processor( + self, ray_bundle: RayBundle, batch: Dict + ) -> Tuple[RayBundle, Dict]: + """An API to add latents, metadata, or other further customization to the RayBundle dataloading process that is parallelized.""" + ray_indices = batch["indices"] + batch["clip"], clip_scale = self.clip_interpolator(ray_indices) + batch["dino"] = self.dino_dataloader(ray_indices) + ray_bundle.metadata["clip_scales"] = clip_scale + + # Assume all cameras have the same focal length and image dimensions. + ray_bundle.metadata["fx"] = self.train_dataset.cameras[0].fx.item() + ray_bundle.metadata["width"] = self.train_dataset.cameras[0].width.item() + ray_bundle.metadata["fy"] = self.train_dataset.cameras[0].fy.item() + ray_bundle.metadata["height"] = self.train_dataset.cameras[0].height.item() + return ray_bundle, batch +``` diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 88b22eadd..bc1b4225a 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -28,7 +28,7 @@ from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import ExternalMethodDummyTrainerConfig, get_external_methods from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig -from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig +from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager, ParallelDataManagerConfig from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig @@ -220,7 +220,7 @@ method_configs["depth-nerfacto"] = TrainerConfig( mixed_precision=True, pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( - _target=VanillaDataManager[DepthDataset], + _target=ParallelDataManager[DepthDataset], dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, @@ -302,7 +302,7 @@ method_configs["instant-ngp-bounded"] = TrainerConfig( method_configs["mipnerf"] = TrainerConfig( method_name="mipnerf", pipeline=VanillaPipelineConfig( - datamanager=ParallelDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), + datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), model=VanillaModelConfig( _target=MipNerfModel, loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0}, @@ -375,7 +375,7 @@ method_configs["tensorf"] = TrainerConfig( max_num_iterations=30000, mixed_precision=False, pipeline=VanillaPipelineConfig( - datamanager=ParallelDataManagerConfig( + datamanager=VanillaDataManagerConfig( dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index cff03607b..833c6044c 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -19,7 +19,6 @@ Datamanager. 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 @@ -42,7 +41,6 @@ from typing import ( import torch import tyro -from torch import nn from torch.nn import Parameter from torch.utils.data.distributed import DistributedSampler from typing_extensions import TypeVar @@ -56,7 +54,12 @@ from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs 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 @@ -64,36 +67,6 @@ 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; @@ -111,7 +84,7 @@ class DataManagerConfig(InstantiateConfig): """Process images on GPU for speed at the expense of memory, if True.""" -class DataManager(nn.Module): +class DataManager: """Generic data manager's abstract class This version of the data manager is designed be a monolithic way to load data and latents, @@ -164,16 +137,16 @@ class DataManager(nn.Module): train_sampler: Optional[DistributedSampler] = None eval_sampler: Optional[DistributedSampler] = None includes_time: bool = False + test_mode: Literal["test", "val", "inference"] = "val" def __init__(self): """Constructor for the DataManager class. Subclassed DataManagers will likely need to override this constructor. - If you aren't manually calling the setup_train and setup_eval functions from an overriden - constructor, that you call super().__init__() BEFORE you initialize any - nn.Modules or nn.Parameters, but AFTER you've already set all the attributes you need - for the setup functions.""" + If you aren't manually calling the setup_train() and setup_eval() functions from an overridden + constructor, please call super().__init__() in your subclass' __init__() method after + you've already set all the attributes you need for the setup functions.""" super().__init__() self.train_count = 0 self.eval_count = 0 @@ -311,18 +284,22 @@ class VanillaDataManagerConfig(DataManagerConfig): """Target class to instantiate.""" dataparser: AnnotatedDataParserUnion = field(default_factory=BlenderDataParserConfig) """Specifies the dataparser used to unpack the data.""" + cache_images_type: Literal["uint8", "float32"] = "float32" + """The image type returned from manager, caching images in uint8 saves memory""" train_num_rays_per_batch: int = 1024 """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.""" + train_num_images_to_sample_from: Union[int, float] = float("inf") + """Number of images to load into CPU RAM to generate RayBundles from during training. If infinity, load + all images into CPU RAM to generate RayBundles.""" + train_num_times_to_repeat_images: Union[int, float] = float("inf") + """Number of RayBundles to generate for a batch of images loaded into CPU RAM before sampling new images. + If infinity, never sample new images. Note: decreasing train_num_images_to_sample_from and increasing + train_num_times_to_repeat_images alleviates CPU bottleneck.""" 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 + eval_num_images_to_sample_from: Union[int, float] = float("inf") """Number of images to sample during eval iteration.""" - eval_num_times_to_repeat_images: int = -1 + eval_num_times_to_repeat_images: Union[int, float] = float("inf") """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,) @@ -331,8 +308,7 @@ class VanillaDataManagerConfig(DataManagerConfig): """Specifies the collate function to use for the train and eval dataloaders.""" camera_res_scale_factor: float = 1.0 """The scale factor for scaling spatial data such as images, mask, semantics - along with relevant information about camera intrinsics - """ + along with relevant information about camera intrinsics""" patch_size: int = 1 """Size of patch to sample from. If > 1, patch-based sampling will be used.""" diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 6ffff9177..3ec06120c 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -29,21 +29,22 @@ from functools import cached_property from pathlib import Path from typing import Dict, ForwardRef, Generic, List, Literal, Optional, Tuple, Type, Union, cast, get_args, get_origin -import cv2 import fpsample import numpy as np import torch from rich.progress import track from torch.nn import Parameter +from torch.utils.data import DataLoader from typing_extensions import assert_never -from nerfstudio.cameras.camera_utils import fisheye624_project, fisheye624_unproject_helper -from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.cameras.cameras import Cameras from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig, TDataset from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.utils.data_utils import identity_collate +from nerfstudio.data.utils.dataloaders import ImageBatchStream, _undistort_image from nerfstudio.utils.misc import get_orig_class from nerfstudio.utils.rich_utils import CONSOLE @@ -61,10 +62,11 @@ class FullImageDatamanagerConfig(DataManagerConfig): 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.""" - cache_images: Literal["cpu", "gpu"] = "gpu" - """Whether to cache images in memory. If "cpu", caches on cpu. If "gpu", caches on device.""" + cache_images: Literal["cpu", "gpu", "disk"] = "gpu" + """Where to cache images in memory. + - If "cpu", caches images on cpu RAM as pytorch tensors. + - If "gpu", caches images on device as pytorch tensors. + - If "disk", keeps images on disk which conserves memory. Datamanager will use parallel dataloader""" cache_images_type: Literal["uint8", "float32"] = "float32" """The image type returned from manager, caching images in uint8 saves memory""" max_thread_workers: Optional[int] = None @@ -79,6 +81,14 @@ class FullImageDatamanagerConfig(DataManagerConfig): fps_reset_every: int = 100 """The number of iterations before one resets fps sampler repeatly, which is essentially drawing fps_reset_every samples from the pool of all training cameras without replacement before a new round of sampling starts.""" + dataloader_num_workers: int = 4 + """The number of workers performing the dataloading from either disk/RAM, which + includes collating, pixel sampling, unprojecting, ray generation etc.""" + prefetch_factor: int = 4 + """The limit number of batches a worker will start loading once an iterator is created. + More details are described here: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader""" + cache_compressed_images: bool = False + """If True, cache raw image files as byte strings to RAM.""" class FullImageDatamanager(DataManager, Generic[TDataset]): @@ -101,6 +111,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): local_rank: int = 0, **kwargs, ): + if config.cache_images == "disk": + try: + torch.multiprocessing.set_start_method("spawn") + except RuntimeError: + assert torch.multiprocessing.get_start_method() == "spawn", 'start method must be "spawn"' self.config = config self.device = device self.world_size = world_size @@ -117,27 +132,20 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): if test_mode == "inference": self.dataparser.downscale_factor = 1 # Avoid opening images self.includes_time = self.dataparser.includes_time - self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train") self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() if len(self.train_dataset) > 500 and self.config.cache_images == "gpu": CONSOLE.print( - "Train dataset has over 500 images, overriding cache_images to cpu", + "Train dataset has over 500 images, overriding cache_images to cpu. If you still get OOM errors or segfault, please consider seting cache_images to 'disk'", style="bold yellow", ) self.config.cache_images = "cpu" - self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device - if self.config.masks_on_gpu is True: - self.exclude_batch_keys_from_device.remove("mask") - if self.config.images_on_gpu is True: - self.exclude_batch_keys_from_device.remove("image") # Some logic to make sure we sample every camera in equal amounts self.train_unseen_cameras = self.sample_train_cameras() self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] assert len(self.train_unseen_cameras) > 0, "No data found in dataset" - super().__init__() def sample_train_cameras(self): @@ -169,7 +177,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): ) n = num_train_cameras kdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(data, n, h=3) - self.train_unsampled_epoch_count += 1 self.train_unsampled_epoch_count[kdline_fps_samples_idx] = 0 return kdline_fps_samples_idx.tolist() @@ -180,19 +187,20 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): def cached_train(self) -> List[Dict[str, torch.Tensor]]: """Get the training images. Will load and undistort the images the first time this (cached) property is accessed.""" + assert self.config.cache_images != "disk", "Can not call _load_images() with `disk` as input" return self._load_images("train", cache_images_device=self.config.cache_images) @cached_property def cached_eval(self) -> List[Dict[str, torch.Tensor]]: """Get the eval images. Will load and undistort the images the first time this (cached) property is accessed.""" + assert self.config.cache_images != "disk", "Can not call _load_images() with `disk` as input" return self._load_images("eval", cache_images_device=self.config.cache_images) def _load_images( self, split: Literal["train", "eval"], cache_images_device: Literal["cpu", "gpu"] ) -> List[Dict[str, torch.Tensor]]: undistorted_images: List[Dict[str, torch.Tensor]] = [] - # Which dataset? if split == "train": dataset = self.train_dataset @@ -240,7 +248,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): total=len(dataset), ) ) - # Move to device. if cache_images_device == "gpu": for cache in undistorted_images: @@ -258,7 +265,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): self.train_cameras = self.train_dataset.cameras else: assert_never(cache_images_device) - return undistorted_images def create_train_dataset(self) -> TDataset: @@ -266,6 +272,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return self.dataset_type( dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, + cache_compressed_images=self.config.cache_compressed_images, ) def create_eval_dataset(self) -> TDataset: @@ -273,6 +280,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return self.dataset_type( dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), scale_factor=self.config.camera_res_scale_factor, + cache_compressed_images=self.config.cache_compressed_images, ) @cached_property @@ -284,7 +292,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return default if orig_class is not None and get_origin(orig_class) is FullImageDatamanager: return get_args(orig_class)[0] - # For inherited classes, we need to find the correct type to instantiate for base in getattr(self, "__orig_bases__", []): if get_origin(base) is FullImageDatamanager: @@ -305,15 +312,54 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): def setup_train(self): """Sets up the data loaders for training""" + if self.config.cache_images == "disk": + self.train_imagebatch_stream = ImageBatchStream( + input_dataset=self.train_dataset, + sampling_seed=self.config.train_cameras_sampling_seed, + cache_images_type=self.config.cache_images_type, + device=self.device, + custom_image_processor=self.custom_image_processor, + ) + self.train_image_dataloader = DataLoader( + self.train_imagebatch_stream, + batch_size=1, + num_workers=self.config.dataloader_num_workers, + collate_fn=identity_collate, + ) + self.iter_train_image_dataloader = iter(self.train_image_dataloader) def setup_eval(self): """Sets up the data loader for evaluation""" + if self.config.cache_images == "disk": + self.eval_imagebatch_stream = ImageBatchStream( + input_dataset=self.eval_dataset, + sampling_seed=self.config.train_cameras_sampling_seed, + cache_images_type=self.config.cache_images_type, + device=self.device, + custom_image_processor=self.custom_image_processor, + ) + self.eval_image_dataloader = DataLoader( + self.eval_imagebatch_stream, + batch_size=1, + num_workers=0, # This must be 0 otherwise there is a crash when trying to pickle custom_image_processor + collate_fn=identity_collate, + ) + self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) @property def fixed_indices_eval_dataloader(self) -> List[Tuple[Cameras, Dict]]: """ Pretends to be the dataloader for evaluation, it returns a list of (camera, data) tuples """ + if self.config.cache_images == "disk": + dataloader = DataLoader( + self.eval_imagebatch_stream, + batch_size=1, + num_workers=0, + collate_fn=identity_collate, + ) + return [batch[0] for batch in dataloader] + image_indices = [i for i in range(len(self.eval_dataset))] data = [d.copy() for d in self.cached_eval] _cameras = deepcopy(self.eval_dataset.cameras).to(self.device) @@ -331,19 +377,19 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): """ return {} - def get_train_rays_per_batch(self): + def get_train_rays_per_batch(self) -> int: """Returns resolution of the image returned from datamanager.""" - if len(self.cached_train) != 0: - h = self.cached_train[0]["image"].shape[0] - w = self.cached_train[0]["image"].shape[1] - return h * w - else: - return 800 * 800 + camera = self.train_dataset.cameras[0].reshape(()) + return int(camera.width[0].item() * camera.height[0].item()) def next_train(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next training batch - Returns a Camera instead of raybundle""" + self.train_count += 1 + if self.config.cache_images == "disk": + camera, data = next(self.iter_train_image_dataloader)[0] + return camera, data + image_idx = self.train_unseen_cameras.pop(0) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.train_unseen_cameras) == 0: @@ -364,16 +410,21 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): def next_eval(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next evaluation batch - Returns a Camera instead of raybundle""" + self.eval_count += 1 + if self.config.cache_images == "disk": + camera, data = next(self.iter_eval_image_dataloader)[0] + return camera, data + return self.next_eval_image(step=step) def next_eval_image(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next evaluation batch - Returns a Camera instead of raybundle - TODO: Make sure this logic is consistent with the vanilladatamanager""" + if self.config.cache_images == "disk": + camera, data = next(self.iter_eval_image_dataloader)[0] + return camera, data image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras) - 1)) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: @@ -385,181 +436,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): camera = self.eval_dataset.cameras[image_idx : image_idx + 1].to(self.device) return camera, data - -def _undistort_image( - camera: Cameras, distortion_params: np.ndarray, data: dict, image: np.ndarray, K: np.ndarray -) -> Tuple[np.ndarray, np.ndarray, Optional[torch.Tensor]]: - mask = None - if camera.camera_type.item() == CameraType.PERSPECTIVE.value: - assert distortion_params[3] == 0, ( - "We doesn't support the 4th Brown parameter for image undistortion, " - "Only k1, k2, k3, p1, p2 can be non-zero." - ) - # because OpenCV expects the order of distortion parameters to be (k1, k2, p1, p2, k3), we need to reorder them - # see https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html - distortion_params = np.array( - [ - distortion_params[0], - distortion_params[1], - distortion_params[4], - distortion_params[5], - distortion_params[2], - distortion_params[3], - 0, - 0, - ] - ) - # because OpenCV expects the pixel coord to be top-left, we need to shift the principal point by 0.5 - # see https://github.com/nerfstudio-project/nerfstudio/issues/3048 - K[0, 2] = K[0, 2] - 0.5 - K[1, 2] = K[1, 2] - 0.5 - if np.any(distortion_params): - newK, roi = cv2.getOptimalNewCameraMatrix(K, distortion_params, (image.shape[1], image.shape[0]), 0) - image = cv2.undistort(image, K, distortion_params, None, newK) # type: ignore - else: - newK = K - roi = 0, 0, image.shape[1], image.shape[0] - # crop the image and update the intrinsics accordingly - x, y, w, h = roi - image = image[y : y + h, x : x + w] - # update the principal point based on our cropped region of interest (ROI) - newK[0, 2] -= x - newK[1, 2] -= y - if "depth_image" in data: - data["depth_image"] = data["depth_image"][y : y + h, x : x + w] - if "mask" in data: - mask = data["mask"].numpy() - mask = mask.astype(np.uint8) * 255 - if np.any(distortion_params): - mask = cv2.undistort(mask, K, distortion_params, None, newK) # type: ignore - mask = mask[y : y + h, x : x + w] - mask = torch.from_numpy(mask).bool() - if len(mask.shape) == 2: - mask = mask[:, :, None] - newK[0, 2] = newK[0, 2] + 0.5 - newK[1, 2] = newK[1, 2] + 0.5 - K = newK - - elif camera.camera_type.item() == CameraType.FISHEYE.value: - K[0, 2] = K[0, 2] - 0.5 - K[1, 2] = K[1, 2] - 0.5 - distortion_params = np.array( - [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] - ) - newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify( - K, distortion_params, (image.shape[1], image.shape[0]), np.eye(3), balance=0 - ) - map1, map2 = cv2.fisheye.initUndistortRectifyMap( - K, distortion_params, np.eye(3), newK, (image.shape[1], image.shape[0]), cv2.CV_32FC1 - ) - # and then remap: - image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR) - if "mask" in data: - mask = data["mask"].numpy() - mask = mask.astype(np.uint8) * 255 - mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK) - mask = torch.from_numpy(mask).bool() - if len(mask.shape) == 2: - mask = mask[:, :, None] - newK[0, 2] = newK[0, 2] + 0.5 - newK[1, 2] = newK[1, 2] + 0.5 - K = newK - elif camera.camera_type.item() == CameraType.FISHEYE624.value: - fisheye624_params = torch.cat( - [camera.fx, camera.fy, camera.cx, camera.cy, torch.from_numpy(distortion_params)], dim=0 - ) - assert fisheye624_params.shape == (16,) - assert ( - "mask" not in data - and camera.metadata is not None - and "fisheye_crop_radius" in camera.metadata - and isinstance(camera.metadata["fisheye_crop_radius"], float) - ) - fisheye_crop_radius = camera.metadata["fisheye_crop_radius"] - - # Approximate the FOV of the unmasked region of the camera. - upper, lower, left, right = fisheye624_unproject_helper( - torch.tensor( - [ - [camera.cx, camera.cy - fisheye_crop_radius], - [camera.cx, camera.cy + fisheye_crop_radius], - [camera.cx - fisheye_crop_radius, camera.cy], - [camera.cx + fisheye_crop_radius, camera.cy], - ], - dtype=torch.float32, - )[None], - params=fisheye624_params[None], - ).squeeze(dim=0) - fov_radians = torch.max( - torch.acos(torch.sum(upper * lower / torch.linalg.norm(upper) / torch.linalg.norm(lower))), - torch.acos(torch.sum(left * right / torch.linalg.norm(left) / torch.linalg.norm(right))), - ) - - # Heuristics to determine parameters of an undistorted image. - undist_h = int(fisheye_crop_radius * 2) - undist_w = int(fisheye_crop_radius * 2) - undistort_focal = undist_h / (2 * torch.tan(fov_radians / 2.0)) - undist_K = torch.eye(3) - undist_K[0, 0] = undistort_focal # fx - undist_K[1, 1] = undistort_focal # fy - undist_K[0, 2] = (undist_w - 1) / 2.0 # cx; for a 1x1 image, center should be at (0, 0). - undist_K[1, 2] = (undist_h - 1) / 2.0 # cy - - # Undistorted 2D coordinates -> rays -> reproject to distorted UV coordinates. - undist_uv_homog = torch.stack( - [ - *torch.meshgrid( - torch.arange(undist_w, dtype=torch.float32), - torch.arange(undist_h, dtype=torch.float32), - ), - torch.ones((undist_w, undist_h), dtype=torch.float32), - ], - dim=-1, - ) - assert undist_uv_homog.shape == (undist_w, undist_h, 3) - dist_uv = ( - fisheye624_project( - xyz=( - torch.einsum( - "ij,bj->bi", - torch.linalg.inv(undist_K), - undist_uv_homog.reshape((undist_w * undist_h, 3)), - )[None] - ), - params=fisheye624_params[None, :], - ) - .reshape((undist_w, undist_h, 2)) - .numpy() - ) - map1 = dist_uv[..., 1] - map2 = dist_uv[..., 0] - - # Use correspondence to undistort image. - image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR) - - # Compute undistorted mask as well. - dist_h = camera.height.item() - dist_w = camera.width.item() - mask = np.mgrid[:dist_h, :dist_w] - mask[0, ...] -= dist_h // 2 - mask[1, ...] -= dist_w // 2 - mask = np.linalg.norm(mask, axis=0) < fisheye_crop_radius - mask = torch.from_numpy( - cv2.remap( - mask.astype(np.uint8) * 255, - map1, - map2, - interpolation=cv2.INTER_LINEAR, - borderMode=cv2.BORDER_CONSTANT, - borderValue=0, - ) - / 255.0 - ).bool()[..., None] - if len(mask.shape) == 2: - mask = mask[:, :, None] - assert mask.shape == (undist_h, undist_w, 1) - K = undist_K.numpy() - else: - raise NotImplementedError("Only perspective and fisheye cameras are supported") - - return K, image, mask + def custom_image_processor(self, camera: Cameras, data: Dict) -> Tuple[Cameras, Dict]: + """An API to add latents, metadata, or other further customization an camera-and-image view dataloading process that is parallelized""" + return camera, data diff --git a/nerfstudio/data/datamanagers/parallel_datamanager.py b/nerfstudio/data/datamanagers/parallel_datamanager.py index 0b56c988c..fe3a62f3c 100644 --- a/nerfstudio/data/datamanagers/parallel_datamanager.py +++ b/nerfstudio/data/datamanagers/parallel_datamanager.py @@ -18,32 +18,28 @@ Parallel data manager that generates training data in multiple python processes. from __future__ import annotations -import concurrent.futures -import queue -import time from dataclasses import dataclass, field from functools import cached_property from pathlib import Path from typing import Dict, ForwardRef, Generic, List, Literal, Optional, Tuple, Type, Union, cast, get_args, get_origin import torch -from pathos.helpers import mp -from rich.progress import track from torch.nn import Parameter +from torch.utils.data import DataLoader -from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle -from nerfstudio.data.datamanagers.base_datamanager import ( - DataManager, - TDataset, - VanillaDataManagerConfig, - variable_res_collate, -) +from nerfstudio.data.datamanagers.base_datamanager import DataManager, TDataset, VanillaDataManagerConfig from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs 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.model_components.ray_generators import RayGenerator +from nerfstudio.data.pixel_samplers import PixelSampler +from nerfstudio.data.utils.data_utils import identity_collate +from nerfstudio.data.utils.dataloaders import ( + FixedIndicesEvalDataloader, + RandIndicesEvalDataloader, + RayBatchStream, + variable_res_collate, +) from nerfstudio.utils.misc import get_orig_class from nerfstudio.utils.rich_utils import CONSOLE @@ -54,79 +50,17 @@ class ParallelDataManagerConfig(VanillaDataManagerConfig): _target: Type = field(default_factory=lambda: ParallelDataManager) """Target class to instantiate.""" - num_processes: int = 1 - """Number of processes to use for train data loading. More than 1 doesn't result in that much better performance""" - queue_size: int = 2 - """Size of shared data queue containing generated ray bundles and batches. - If queue_size <= 0, the queue size is infinite.""" - max_thread_workers: Optional[int] = None - """Maximum number of threads to use in thread pool executor. If None, use ThreadPool default.""" - - -class DataProcessor(mp.Process): # type: ignore - """Parallel dataset batch processor. - - This class is responsible for generating ray bundles from an input dataset - in parallel python processes. - - Args: - out_queue: the output queue for storing the processed data - config: configuration object for the parallel data manager - dataparser_outputs: outputs from the dataparser - dataset: input dataset - pixel_sampler: The pixel sampler for sampling rays - """ - - def __init__( - self, - out_queue: mp.Queue, # type: ignore - config: ParallelDataManagerConfig, - dataparser_outputs: DataparserOutputs, - dataset: TDataset, - pixel_sampler: PixelSampler, - ): - super().__init__() - self.daemon = True - self.out_queue = out_queue - self.config = config - self.dataparser_outputs = dataparser_outputs - self.dataset = dataset - self.exclude_batch_keys_from_device = self.dataset.exclude_batch_keys_from_device - self.pixel_sampler = pixel_sampler - self.ray_generator = RayGenerator(self.dataset.cameras) - - def run(self): - """Append out queue in parallel with ray bundles and batches.""" - self.cache_images() - while True: - batch = self.pixel_sampler.sample(self.img_data) - ray_indices = batch["indices"] - ray_bundle: RayBundle = self.ray_generator(ray_indices) - # check that GPUs are available - if torch.cuda.is_available(): - ray_bundle = ray_bundle.pin_memory() - while True: - try: - self.out_queue.put((ray_bundle, batch)) - break - except queue.Full: - time.sleep(0.0001) - except Exception: - CONSOLE.print_exception() - CONSOLE.print("[bold red]Error: Error occurred in parallel datamanager queue.") - - def cache_images(self): - """Caches all input images into a NxHxWx3 tensor.""" - indices = range(len(self.dataset)) - batch_list = [] - results = [] - with concurrent.futures.ThreadPoolExecutor(max_workers=self.config.max_thread_workers) as executor: - for idx in indices: - res = executor.submit(self.dataset.__getitem__, idx) - results.append(res) - for res in track(results, description="Loading data batch", transient=False): - batch_list.append(res.result()) - self.img_data = self.config.collate_fn(batch_list) + load_from_disk: bool = False + """If True, conserves RAM memory by loading images from disk. + If False, caches all the images as tensors to RAM and loads from RAM.""" + dataloader_num_workers: int = 4 + """The number of workers performing the dataloading from either disk/RAM, which + includes collating, pixel sampling, unprojecting, ray generation etc.""" + prefetch_factor: int = 10 + """The limit number of batches a worker will start loading once an iterator is created. + More details are described here: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader""" + cache_compressed_images: bool = False + """If True, cache raw image files as byte strings to RAM.""" class ParallelDataManager(DataManager, Generic[TDataset]): @@ -136,6 +70,13 @@ class ParallelDataManager(DataManager, Generic[TDataset]): config: the DataManagerConfig used to instantiate class """ + config: ParallelDataManagerConfig + train_dataset: TDataset + eval_dataset: TDataset + train_dataparser_outputs: DataparserOutputs + train_pixel_sampler: Optional[PixelSampler] = None + eval_pixel_sampler: Optional[PixelSampler] = None + def __init__( self, config: ParallelDataManagerConfig, @@ -145,6 +86,10 @@ class ParallelDataManager(DataManager, Generic[TDataset]): local_rank: int = 0, **kwargs, ): + try: + torch.multiprocessing.set_start_method("spawn") + except RuntimeError: + assert torch.multiprocessing.get_start_method() == "spawn", 'start method must be "spawn"' self.config = config self.device = device self.world_size = world_size @@ -161,21 +106,22 @@ class ParallelDataManager(DataManager, Generic[TDataset]): self.dataparser.downscale_factor = 1 # Avoid opening images self.includes_time = self.dataparser.includes_time self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train") - self.eval_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split=self.test_split) - cameras = self.train_dataparser_outputs.cameras - if len(cameras) > 1: - for i in range(1, len(cameras)): - if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height: - CONSOLE.print("Variable resolution, using variable_res_collate") - self.config.collate_fn = variable_res_collate - break self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() self.exclude_batch_keys_from_device = self.train_dataset.exclude_batch_keys_from_device - # Spawn is critical for not freezing the program (PyTorch compatability issue) - # check if spawn is already set - if mp.get_start_method(allow_none=True) is None: # type: ignore - mp.set_start_method("spawn") # type: ignore + if self.config.masks_on_gpu is True and "mask" in self.exclude_batch_keys_from_device: + self.exclude_batch_keys_from_device.remove("mask") + if self.config.images_on_gpu is True and "image" in self.exclude_batch_keys_from_device: + self.exclude_batch_keys_from_device.remove("image") + # Setup our collate function (same as base_datamanager.py) + if self.train_dataparser_outputs is not None: + cameras = self.train_dataparser_outputs.cameras + if len(cameras) > 1: + for i in range(1, len(cameras)): + if cameras[0].width != cameras[i].width or cameras[0].height != cameras[i].height or True: + CONSOLE.print("Variable resolution, using variable_res_collate") + self.config.collate_fn = variable_res_collate + break super().__init__() @cached_property @@ -208,6 +154,7 @@ class ParallelDataManager(DataManager, Generic[TDataset]): return self.dataset_type( dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, + cache_compressed_images=self.config.cache_compressed_images, ) def create_eval_dataset(self) -> TDataset: @@ -215,97 +162,96 @@ class ParallelDataManager(DataManager, Generic[TDataset]): return self.dataset_type( dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split), scale_factor=self.config.camera_res_scale_factor, - ) - - def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> PixelSampler: - """Infer pixel sampler to use.""" - if self.config.patch_size > 1 and type(self.config.pixel_sampler) is PixelSamplerConfig: - return PatchPixelSamplerConfig().setup( - patch_size=self.config.patch_size, num_rays_per_batch=num_rays_per_batch - ) - is_equirectangular = (dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value).all() - if is_equirectangular.any(): - CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - - fisheye_crop_radius = None - if dataset.cameras.metadata is not None: - fisheye_crop_radius = dataset.cameras.metadata.get("fisheye_crop_radius") - - return self.config.pixel_sampler.setup( - is_equirectangular=is_equirectangular, - num_rays_per_batch=num_rays_per_batch, - fisheye_crop_radius=fisheye_crop_radius, + cache_compressed_images=self.config.cache_compressed_images, ) def setup_train(self): - """Sets up parallel python data processes for training.""" - assert self.train_dataset is not None - self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) # type: ignore - self.data_queue = mp.Queue(maxsize=self.config.queue_size) # type: ignore - self.data_procs = [ - DataProcessor( - out_queue=self.data_queue, # type: ignore - config=self.config, - dataparser_outputs=self.train_dataparser_outputs, - dataset=self.train_dataset, - pixel_sampler=self.train_pixel_sampler, - ) - for i in range(self.config.num_processes) - ] - for proc in self.data_procs: - proc.start() - print("Started threads") + self.train_raybatchstream = RayBatchStream( + input_dataset=self.train_dataset, + num_rays_per_batch=self.config.train_num_rays_per_batch, + num_images_to_sample_from=( + 50 + if self.config.load_from_disk and self.config.train_num_images_to_sample_from == float("inf") + else self.config.train_num_images_to_sample_from + ), + num_times_to_repeat_images=( + 10 + if self.config.load_from_disk and self.config.train_num_times_to_repeat_images == float("inf") + else self.config.train_num_times_to_repeat_images + ), + device=self.device, + collate_fn=variable_res_collate, + load_from_disk=self.config.load_from_disk, + custom_ray_processor=self.custom_ray_processor, + ) + self.train_ray_dataloader = DataLoader( + self.train_raybatchstream, + batch_size=1, + num_workers=self.config.dataloader_num_workers, + prefetch_factor=self.config.prefetch_factor, + shuffle=False, + collate_fn=identity_collate, # Our dataset handles batching / collation of rays + ) + self.iter_train_raybundles = iter(self.train_ray_dataloader) 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, + self.eval_raybatchstream = RayBatchStream( + input_dataset=self.eval_dataset, + num_rays_per_batch=self.config.train_num_rays_per_batch, + num_images_to_sample_from=( + 50 + if self.config.load_from_disk and self.config.eval_num_images_to_sample_from == float("inf") + else self.config.eval_num_images_to_sample_from + ), + num_times_to_repeat_images=( + 10 + if self.config.load_from_disk and self.config.eval_num_times_to_repeat_images == float("inf") + else self.config.eval_num_times_to_repeat_images + ), device=self.device, - num_workers=self.world_size * 4, - pin_memory=True, - collate_fn=self.config.collate_fn, - exclude_batch_keys_from_device=self.exclude_batch_keys_from_device, + collate_fn=variable_res_collate, + load_from_disk=True, + custom_ray_processor=self.custom_ray_processor, ) - self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) - self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch) # type: ignore - self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device)) - # for loading full images - self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( + self.eval_ray_dataloader = DataLoader( + self.eval_raybatchstream, + batch_size=1, + num_workers=0, # This must be 0 otherwise there is a crash when trying to pickle custom_ray_processor + shuffle=False, + collate_fn=identity_collate, # Our dataset handles batching / collation of rays + ) + self.iter_eval_raybundles = iter(self.eval_ray_dataloader) + self.image_eval_dataloader = RandIndicesEvalDataloader( # this is used for ns-eval input_dataset=self.eval_dataset, device=self.device, num_workers=self.world_size * 4, ) - self.eval_dataloader = RandIndicesEvalDataloader( + self.eval_dataloader = RandIndicesEvalDataloader( # this is used by ns-render + input_dataset=self.eval_dataset, + device=self.device, + num_workers=self.world_size * 4, + ) + self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( input_dataset=self.eval_dataset, device=self.device, num_workers=self.world_size * 4, ) def next_train(self, step: int) -> Tuple[RayBundle, Dict]: - """Returns the next batch of data from the parallel training processes.""" + """Returns the next batch of data from the train dataloader.""" self.train_count += 1 - bundle, batch = self.data_queue.get() - ray_bundle = bundle.to(self.device) + ray_bundle, batch = next(self.iter_train_raybundles)[0] return ray_bundle, batch def next_eval(self, step: int) -> Tuple[RayBundle, Dict]: """Returns the next batch of data from the eval dataloader.""" self.eval_count += 1 - image_batch = next(self.iter_eval_image_dataloader) - assert self.eval_pixel_sampler is not None - assert isinstance(image_batch, dict) - batch = self.eval_pixel_sampler.sample(image_batch) - ray_indices = batch["indices"] - ray_bundle = self.eval_ray_generator(ray_indices) + ray_bundle, batch = next(self.iter_train_raybundles)[0] return ray_bundle, batch def next_eval_image(self, step: int) -> Tuple[Cameras, Dict]: """Retrieve the next eval image.""" - for camera, batch in self.eval_dataloader: + for camera, batch in self.image_eval_dataloader: assert camera.shape[0] == 1 return camera, batch raise ValueError("No more eval images") @@ -333,9 +279,6 @@ class ParallelDataManager(DataManager, Generic[TDataset]): """ return {} - def __del__(self): - """Clean up the parallel data processes.""" - if hasattr(self, "data_procs"): - for proc in self.data_procs: - proc.terminate() - proc.join() + def custom_ray_processor(self, ray_bundle: RayBundle, batch: Dict) -> Tuple[RayBundle, Dict]: + """An API to add latents, metadata, or other further customization to the RayBundle dataloading process that is parallelized""" + return ray_bundle, batch diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index e9493dcbc..16194a82d 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/nerfstudio/data/datasets/base_dataset.py @@ -18,6 +18,7 @@ Dataset. from __future__ import annotations +import io from copy import deepcopy from pathlib import Path from typing import Dict, List, Literal @@ -32,7 +33,7 @@ from torch.utils.data import Dataset from nerfstudio.cameras.cameras import Cameras from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs -from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path +from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path, pil_to_numpy class InputDataset(Dataset): @@ -46,7 +47,9 @@ class InputDataset(Dataset): exclude_batch_keys_from_device: List[str] = ["image", "mask"] cameras: Cameras - def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): + def __init__( + self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0, cache_compressed_images: bool = False + ): super().__init__() self._dataparser_outputs = dataparser_outputs self.scale_factor = scale_factor @@ -55,6 +58,18 @@ class InputDataset(Dataset): self.cameras = deepcopy(dataparser_outputs.cameras) self.cameras.rescale_output_resolution(scaling_factor=scale_factor) self.mask_color = dataparser_outputs.metadata.get("mask_color", None) + self.cache_compressed_images = cache_compressed_images + """If cache_compressed_images == True, cache all the image files into RAM in their compressed form (jpeg, png, etc. but not as pytorch tensors)""" + if cache_compressed_images: + self.binary_images = [] + self.binary_masks = [] + for image_filename in self._dataparser_outputs.image_filenames: + with open(image_filename, "rb") as f: + self.binary_images.append(io.BytesIO(f.read())) + if self._dataparser_outputs.mask_filenames is not None: + for mask_filename in self._dataparser_outputs.mask_filenames: + with open(mask_filename, "rb") as f: + self.binary_masks.append(io.BytesIO(f.read())) def __len__(self): return len(self._dataparser_outputs.image_filenames) @@ -66,17 +81,20 @@ class InputDataset(Dataset): image_idx: The image index in the dataset. """ image_filename = self._dataparser_outputs.image_filenames[image_idx] - pil_image = Image.open(image_filename) + if self.cache_compressed_images: + pil_image = Image.open(self.binary_images[image_idx]) + else: + pil_image = Image.open(image_filename) if self.scale_factor != 1.0: width, height = pil_image.size newsize = (int(width * self.scale_factor), int(height * self.scale_factor)) pil_image = pil_image.resize(newsize, resample=Image.Resampling.BILINEAR) - image = np.array(pil_image, dtype="uint8") # shape is (h, w) or (h, w, 3 or 4) + image = pil_to_numpy(pil_image) # shape is (h, w) or (h, w, 3 or 4) and dtype == "uint8" if len(image.shape) == 2: image = image[:, :, None].repeat(3, axis=2) assert len(image.shape) == 3 assert image.dtype == np.uint8 - assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." + assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is incorrect." return image def get_image_float32(self, image_idx: int) -> Float[Tensor, "image_height image_width num_channels"]: @@ -85,7 +103,9 @@ class InputDataset(Dataset): Args: image_idx: The image index in the dataset. """ - image = torch.from_numpy(self.get_numpy_image(image_idx).astype("float32") / 255.0) + image = self.get_numpy_image(image_idx) + image = image / np.float32(255) + image = torch.from_numpy(image) if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4: assert (self._dataparser_outputs.alpha_color >= 0).all() and ( self._dataparser_outputs.alpha_color <= 1 @@ -99,7 +119,9 @@ class InputDataset(Dataset): Args: image_idx: The image index in the dataset. """ - image = torch.from_numpy(self.get_numpy_image(image_idx)) + image = torch.from_numpy( + self.get_numpy_image(image_idx) + ) # removed astype(np.uint8) because get_numpy_image returns uint8 if self._dataparser_outputs.alpha_color is not None and image.shape[-1] == 4: assert (self._dataparser_outputs.alpha_color >= 0).all() and ( self._dataparser_outputs.alpha_color <= 1 @@ -126,7 +148,10 @@ class InputDataset(Dataset): data = {"image_idx": image_idx, "image": image} if self._dataparser_outputs.mask_filenames is not None: - mask_filepath = self._dataparser_outputs.mask_filenames[image_idx] + if self.cache_compressed_images: + mask_filepath = self.binary_masks[image_idx] + else: + mask_filepath = self._dataparser_outputs.mask_filenames[image_idx] data["mask"] = get_image_mask_tensor_from_path(filepath=mask_filepath, scale_factor=self.scale_factor) assert data["mask"].shape[:2] == data["image"].shape[:2], ( f"Mask and image have different shapes. Got {data['mask'].shape[:2]} and {data['image'].shape[:2]}" diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index d4e22e3b6..fb0b5dfaa 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/nerfstudio/data/datasets/depth_dataset.py @@ -43,7 +43,9 @@ class DepthDataset(InputDataset): exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"] - def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): + def __init__( + self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0, cache_compressed_images: bool = False + ): super().__init__(dataparser_outputs, scale_factor) # if there are no depth images than we want to generate them all with zoe depth @@ -77,7 +79,7 @@ class DepthDataset(InputDataset): filenames = dataparser_outputs.image_filenames repo = "isl-org/ZoeDepth" - self.zoe = torch_compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).to(device)) + zoe = torch_compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).to(device)) for i in track(range(len(filenames)), description="Generating depth images"): image_filename = filenames[i] @@ -91,7 +93,7 @@ class DepthDataset(InputDataset): image = torch.permute(image, (2, 0, 1)).unsqueeze(0).to(device) if image.shape[1] == 4: image = image[:, :3, :, :] - depth_tensor = self.zoe.infer(image).squeeze().unsqueeze(-1) + depth_tensor = zoe.infer(image).squeeze().unsqueeze(-1) depth_tensors.append(depth_tensor) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 009d40571..f2bc6d96e 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -106,7 +106,7 @@ class PixelSampler: num_valid = 0 for _ in range(self.config.max_num_iterations): c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - chosen_indices_validity = mask.squeeze()[c, y, x].bool() + chosen_indices_validity = mask.squeeze(-1)[c, y, x].bool() num_valid = int(torch.sum(chosen_indices_validity).item()) if num_valid == num_samples: break diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index 51d57da67..e79f16903 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/nerfstudio/data/utils/data_utils.py @@ -15,15 +15,48 @@ """Utility functions to allow easy re-use of common operations across dataloaders""" from pathlib import Path -from typing import List, Tuple, Union +from typing import IO, List, Tuple, Union import cv2 import numpy as np import torch from PIL import Image +from PIL.Image import Image as PILImage -def get_image_mask_tensor_from_path(filepath: Path, scale_factor: float = 1.0) -> torch.Tensor: +def pil_to_numpy(im: PILImage) -> np.ndarray: + """Converts a PIL Image object to a NumPy array. + + Args: + im (PIL.Image.Image): The input PIL Image object. + + Returns: + numpy.ndarray representing the image data. + """ + # Load in image completely (PIL defaults to lazy loading) + im.load() + + # Unpack data + e = Image._getencoder(im.mode, "raw", im.mode) + e.setimage(im.im) + + # NumPy buffer for the result + shape, typestr = Image._conv_type_shape(im) + data = np.empty(shape, dtype=np.dtype(typestr)) + mem = data.data.cast("B", (data.data.nbytes,)) + + bufsize, s, offset = 65536, 0, 0 + while not s: + _, s, d = e.encode(bufsize) + mem[offset : offset + len(d)] = d + offset += len(d) + if s < 0: + raise RuntimeError("encoder error %d in tobytes" % s) + + return data + + +def get_image_mask_tensor_from_path(filepath: Union[Path, IO[bytes]], scale_factor: float = 1.0) -> torch.Tensor: """ Utility function to read a mask image from the given path and return a boolean tensor """ @@ -32,7 +65,7 @@ def get_image_mask_tensor_from_path(filepath: Path, scale_factor: float = 1.0) - width, height = pil_mask.size newsize = (int(width * scale_factor), int(height * scale_factor)) pil_mask = pil_mask.resize(newsize, resample=Image.Resampling.NEAREST) - mask_tensor = torch.from_numpy(np.array(pil_mask)).unsqueeze(-1).bool() + mask_tensor = torch.from_numpy(pil_to_numpy(pil_mask)).unsqueeze(-1).bool() if len(mask_tensor.shape) != 3: raise ValueError("The mask image should have 1 channel") return mask_tensor @@ -85,3 +118,8 @@ def get_depth_image_from_path( image = image.astype(np.float32) * scale_factor image = cv2.resize(image, (width, height), interpolation=interpolation) return torch.from_numpy(image[:, :, np.newaxis]) + + +def identity_collate(x): + """This function does nothing but serves to help our dataloaders have a pickleable function, as lambdas are not pickleable""" + return x diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index 6a64ba738..9fa2faaf7 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -18,32 +18,281 @@ Code for sampling images from a dataset of images. # for multithreading import concurrent.futures +import math import multiprocessing import random from abc import abstractmethod -from typing import Any, Callable, Dict, List, Optional, Sized, Tuple, Union +from collections import defaultdict +from typing import Any, Callable, Dict, List, Literal, Optional, Sized, Tuple, Union, cast +import cv2 +import numpy as np import torch from rich.progress import track -from torch.utils.data import Dataset +from torch.utils.data import Dataset, IterableDataset, get_worker_info from torch.utils.data.dataloader import DataLoader -from nerfstudio.cameras.cameras import Cameras +from nerfstudio.cameras.camera_utils import fisheye624_project, fisheye624_unproject_helper +from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datasets.base_dataset import InputDataset +from nerfstudio.data.pixel_samplers import PatchPixelSamplerConfig, PixelSampler, PixelSamplerConfig from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate +from nerfstudio.model_components.ray_generators import RayGenerator from nerfstudio.utils.misc import get_dict_to_torch from nerfstudio.utils.rich_utils import CONSOLE +def variable_res_collate(batch: List[Dict]) -> Dict: + """Default collate function for our 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 + + +def _undistort_image( + camera: Cameras, distortion_params: np.ndarray, data: dict, image: np.ndarray, K: np.ndarray +) -> Tuple[np.ndarray, np.ndarray, Optional[torch.Tensor]]: + mask = None + if camera.camera_type.item() == CameraType.PERSPECTIVE.value: + assert distortion_params[3] == 0, ( + "We don't support the 4th Brown parameter for image undistortion, Only k1, k2, k3, p1, p2 can be non-zero." + ) + # we rearrange the distortion parameters because OpenCV expects the order (k1, k2, p1, p2, k3) + # see https://docs.opencv.org/4.x/dc/dbb/tutorial_py_calibration.html + distortion_params = np.array( + [ + distortion_params[0], + distortion_params[1], + distortion_params[4], + distortion_params[5], + distortion_params[2], + distortion_params[3], + 0, + 0, + ] + ) + # because OpenCV expects the pixel coord to be top-left, we need to shift the principal point by 0.5 + # see https://github.com/nerfstudio-project/nerfstudio/issues/3048 + K[0, 2] = K[0, 2] - 0.5 + K[1, 2] = K[1, 2] - 0.5 + if np.any(distortion_params): + newK, roi = cv2.getOptimalNewCameraMatrix(K, distortion_params, (image.shape[1], image.shape[0]), 0) + image = cv2.undistort(image, K, distortion_params, None, newK) # type: ignore + else: + newK = K + roi = 0, 0, image.shape[1], image.shape[0] + # crop the image and update the intrinsics accordingly + x, y, w, h = roi + image = image[y : y + h, x : x + w] + newK[0, 2] -= x + newK[1, 2] -= y + + if "depth_image" in data: + data["depth_image"] = data["depth_image"][y : y + h, x : x + w] + if "mask" in data: + mask = data["mask"].numpy() + mask = mask.astype(np.uint8) * 255 + if np.any(distortion_params): + mask = cv2.undistort(mask, K, distortion_params, None, newK) # type: ignore + mask = mask[y : y + h, x : x + w] + mask = torch.from_numpy(mask).bool() + if len(mask.shape) == 2: + mask = mask[:, :, None] + newK[0, 2] = newK[0, 2] + 0.5 + newK[1, 2] = newK[1, 2] + 0.5 + K = newK + + elif camera.camera_type.item() == CameraType.FISHEYE.value: + K[0, 2] = K[0, 2] - 0.5 + K[1, 2] = K[1, 2] - 0.5 + distortion_params = np.array( + [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] + ) + newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify( + K, distortion_params, (image.shape[1], image.shape[0]), np.eye(3), balance=0 + ) + map1, map2 = cv2.fisheye.initUndistortRectifyMap( + K, distortion_params, np.eye(3), newK, (image.shape[1], image.shape[0]), cv2.CV_32FC1 + ) + # and then remap: + image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR) + if "mask" in data: + mask = data["mask"].numpy() + mask = mask.astype(np.uint8) * 255 + mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK) + mask = torch.from_numpy(mask).bool() + if len(mask.shape) == 2: + mask = mask[:, :, None] + newK[0, 2] = newK[0, 2] + 0.5 + newK[1, 2] = newK[1, 2] + 0.5 + K = newK + elif camera.camera_type.item() == CameraType.FISHEYE624.value: + fisheye624_params = torch.cat( + [camera.fx, camera.fy, camera.cx, camera.cy, torch.from_numpy(distortion_params)], dim=0 + ) + assert fisheye624_params.shape == (16,) + assert ( + "mask" not in data + and camera.metadata is not None + and "fisheye_crop_radius" in camera.metadata + and isinstance(camera.metadata["fisheye_crop_radius"], float) + ) + fisheye_crop_radius = camera.metadata["fisheye_crop_radius"] + + # Approximate the FOV of the unmasked region of the camera. + upper, lower, left, right = fisheye624_unproject_helper( + torch.tensor( + [ + [camera.cx, camera.cy - fisheye_crop_radius], + [camera.cx, camera.cy + fisheye_crop_radius], + [camera.cx - fisheye_crop_radius, camera.cy], + [camera.cx + fisheye_crop_radius, camera.cy], + ], + dtype=torch.float32, + )[None], + params=fisheye624_params[None], + ).squeeze(dim=0) + fov_radians = torch.max( + torch.acos(torch.sum(upper * lower / torch.linalg.norm(upper) / torch.linalg.norm(lower))), + torch.acos(torch.sum(left * right / torch.linalg.norm(left) / torch.linalg.norm(right))), + ) + + # Heuristics to determine parameters of an undistorted image. + undist_h = int(fisheye_crop_radius * 2) + undist_w = int(fisheye_crop_radius * 2) + undistort_focal = undist_h / (2 * torch.tan(fov_radians / 2.0)) + undist_K = torch.eye(3) + undist_K[0, 0] = undistort_focal # fx + undist_K[1, 1] = undistort_focal # fy + undist_K[0, 2] = (undist_w - 1) / 2.0 # cx; for a 1x1 image, center should be at (0, 0). + undist_K[1, 2] = (undist_h - 1) / 2.0 # cy + + # Undistorted 2D coordinates -> rays -> reproject to distorted UV coordinates. + undist_uv_homog = torch.stack( + [ + *torch.meshgrid( + torch.arange(undist_w, dtype=torch.float32), + torch.arange(undist_h, dtype=torch.float32), + ), + torch.ones((undist_w, undist_h), dtype=torch.float32), + ], + dim=-1, + ) + assert undist_uv_homog.shape == (undist_w, undist_h, 3) + dist_uv = ( + fisheye624_project( + xyz=( + torch.einsum( + "ij,bj->bi", + torch.linalg.inv(undist_K), + undist_uv_homog.reshape((undist_w * undist_h, 3)), + )[None] + ), + params=fisheye624_params[None, :], + ) + .reshape((undist_w, undist_h, 2)) + .numpy() + ) + map1 = dist_uv[..., 1] + map2 = dist_uv[..., 0] + + # Use correspondence to undistort image. + image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR) + + # Compute undistorted mask as well. + dist_h = camera.height.item() + dist_w = camera.width.item() + mask = np.mgrid[:dist_h, :dist_w] + mask[0, ...] -= dist_h // 2 + mask[1, ...] -= dist_w // 2 + mask = np.linalg.norm(mask, axis=0) < fisheye_crop_radius + mask = torch.from_numpy( + cv2.remap( + mask.astype(np.uint8) * 255, + map1, + map2, + interpolation=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + ) + / 255.0 + ).bool()[..., None] + if len(mask.shape) == 2: + mask = mask[:, :, None] + assert mask.shape == (undist_h, undist_w, 1) + K = undist_K.numpy() + else: + raise NotImplementedError("Only perspective and fisheye cameras are supported") + return K, image, mask + + +def undistort_view( + idx: int, dataset: InputDataset, image_type: Literal["uint8", "float32"] = "float32" +) -> Tuple[Cameras, Dict]: + """Undistorts an image to one taken by a linear (pinhole) camera model and returns a new Camera with these updated intrinsics + Note: this method does not modify the dataset's attributes at all. + + Returns: The undistorted data (image, depth, mask, etc.) and the new linear Camera object + """ + data = dataset.get_data(idx, image_type) + camera = dataset.cameras[idx].reshape(()) + assert data["image"].shape[1] == camera.width.item() and data["image"].shape[0] == camera.height.item(), ( + f"The size of image ({data['image'].shape[1]}, {data['image'].shape[0]}) loaded " + f"does not match the camera parameters ({camera.width.item(), camera.height.item()}), idx = {idx}" + ) + if camera.distortion_params is None or torch.all(camera.distortion_params == 0): + return camera.reshape((1,)), data + K = camera.get_intrinsics_matrices().numpy() + distortion_params = camera.distortion_params.numpy() + image = data["image"].numpy() + K, image, mask = _undistort_image(camera, distortion_params, data, image, K) + data["image"] = torch.from_numpy(image) + if mask is not None: + data["mask"] = mask + + # create a new Camera with the rectified / undistorted intrinsics + new_camera = Cameras( + camera_to_worlds=camera.camera_to_worlds.unsqueeze(0), + fx=torch.Tensor([[float(K[0, 0])]]), + fy=torch.Tensor([[float(K[1, 1])]]), + cx=torch.Tensor([[float(K[0, 2])]]), + cy=torch.Tensor([[float(K[1, 2])]]), + width=torch.Tensor([[image.shape[1]]]).to(torch.int32), + height=torch.Tensor([[image.shape[0]]]).to(torch.int32), + ) + return new_camera, data + + class CacheDataloader(DataLoader): """Collated image dataset that implements caching of default-pytorch-collatable data. Creates batches of the InputDataset return type. 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. -1 to never pick new images. + num_samples_to_collate: How many images to sample rays for each batch. -1 or infinity for all images. + num_times_to_repeat_images: How often to yield an image batch before resampling. -1 or infinity to never pick new images. device: Device to perform computation. collate_fn: The function we will use to collate our training data """ @@ -51,8 +300,8 @@ class CacheDataloader(DataLoader): def __init__( self, dataset: Dataset, - num_images_to_sample_from: int = -1, - num_times_to_repeat_images: int = -1, + num_images_to_sample_from: Union[int, float] = float("inf"), + num_times_to_repeat_images: Union[int, float] = float("inf"), device: Union[torch.device, str] = "cpu", collate_fn: Callable[[Any], Any] = nerfstudio_collate, exclude_batch_keys_from_device: Optional[List[str]] = None, @@ -83,7 +332,7 @@ class CacheDataloader(DataLoader): "[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: + elif self.num_times_to_repeat_images == float("inf"): CONSOLE.print( f"Caching {self.num_images_to_sample_from} out of {len(self.dataset)} images, without resampling." ) @@ -100,7 +349,14 @@ class CacheDataloader(DataLoader): """Returns a list of batches from the dataset attribute.""" assert isinstance(self.dataset, Sized) - indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) + indices = random.sample( + range(len(self.dataset)), + k=( + len(self.dataset) + if self.num_images_to_sample_from == float("inf") + else int(self.num_images_to_sample_from) + ), + ) batch_list = [] results = [] @@ -119,7 +375,7 @@ class CacheDataloader(DataLoader): return batch_list def _get_collated_batch(self): - """Returns a collated batch.""" + """Returns a collated batch of images.""" batch_list = self._get_batch_list() collated_batch = self.collate_fn(batch_list) collated_batch = get_dict_to_torch( @@ -146,6 +402,256 @@ class CacheDataloader(DataLoader): yield collated_batch +class RayBatchStream(IterableDataset): + """Wrapper around Pytorch's IterableDataset to generate the next batch of rays (next RayBundle) and corresponding labels + with multiple parallel workers. + + Each worker samples a small batch of images, pixel samples those images, and generates rays for one training step. + The same batch of images can be pixel sampled multiple times hasten ray generation, as retrieving images is process + bottlenecked by disk read speed. To avoid Out-Of-Memory (OOM) errors, this batch of images is small and regenerated + by resampling the worker's partition of images to maintain sampling diversity. + """ + + def __init__( + self, + input_dataset: InputDataset, + sampling_seed: int = 3301, + num_rays_per_batch: int = 1024, + num_images_to_sample_from: Union[int, float] = float("inf"), + num_times_to_repeat_images: Union[int, float] = float("inf"), + device: Union[torch.device, str] = "cpu", + # variable_res_collate avoids np.stack'ing images, which allows it to be much faster than `nerfstudio_collate` + collate_fn: Callable[[Any], Any] = cast(Any, staticmethod(variable_res_collate)), + num_image_load_threads: int = 4, + exclude_batch_keys_from_device: Optional[List[str]] = None, + load_from_disk: bool = False, + patch_size: int = 1, + custom_ray_processor: Optional[Callable[[RayBundle, Dict], Tuple[RayBundle, Dict]]] = None, + ): + if exclude_batch_keys_from_device is None: + exclude_batch_keys_from_device = ["image"] + self.input_dataset = input_dataset + self.sampling_seed = sampling_seed + assert isinstance(self.input_dataset, Sized) + self.num_rays_per_batch = num_rays_per_batch + """Number of rays per batch to user per training iteration.""" + self.num_images_to_sample_from = num_images_to_sample_from + """How many images to sample to generate a RayBundle. More images means greater sampling diversity at expense of increased RAM usage.""" + self.num_times_to_repeat_images = num_times_to_repeat_images + """How many RayBundles to generate from this batch of images after sampling `num_images_to_sample_from` images.""" + self.device = device + """If a CUDA GPU is present, self.device will be set to use that GPU.""" + self.collate_fn = collate_fn + """What collate function is used to batch images to be used for pixel sampling and ray generation. """ + self.num_image_load_threads = num_image_load_threads + """Number of threads created to read images from disk and form collated batches.""" + self.exclude_batch_keys_from_device = exclude_batch_keys_from_device + """Which key of the batch (such as 'image', 'mask','depth') to prevent from moving to the device. + For instance, if you would like to conserve GPU memory, don't move the image tensors to the GPU, + which comes at a cost of total training time. The default value is ['image'].""" + self.load_from_disk = load_from_disk + """If True, conserves RAM memory by loading images from disk. + If False, each worker caches all the images in its dataset partition as tensors to RAM and loads from RAM.""" + self.patch_size = patch_size + """Size of patch to sample from. If > 1, patch-based sampling will be used.""" + self._cached_collated_batch = None + """Each worker has a self._cached_collated_batch contains a collated batch of images cached in RAM for a specific worker that's ready for pixel sampling.""" + self.pixel_sampler_config: PixelSamplerConfig = PixelSamplerConfig() + """Specifies the pixel sampler config used to sample pixels from images. Each worker will have its own pixel sampler""" + self.ray_generator: Optional[RayGenerator] = None + """Each worker will have its own ray generator, so this is set to None for now.""" + self.custom_ray_processor = custom_ray_processor + + def _get_pixel_sampler(self, dataset: InputDataset, num_rays_per_batch: int) -> PixelSampler: + """copied from VanillaDataManager.""" + from nerfstudio.cameras.cameras import CameraType + + if self.patch_size > 1 and type(self.pixel_sampler_config) is PixelSamplerConfig: + return PatchPixelSamplerConfig().setup(patch_size=self.patch_size, num_rays_per_batch=num_rays_per_batch) + is_equirectangular = (dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value).all() + if is_equirectangular.any(): + CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") + + fisheye_crop_radius = None + if dataset.cameras.metadata is not None: + fisheye_crop_radius = dataset.cameras.metadata.get("fisheye_crop_radius") + + return self.pixel_sampler_config.setup( + is_equirectangular=is_equirectangular, + num_rays_per_batch=num_rays_per_batch, + fisheye_crop_radius=fisheye_crop_radius, + ) + + def _get_batch_list(self, indices): + """Returns a list representing a single batch from the dataset attribute. + Each item of the list is a dictionary with dict_keys(['image_idx', 'image']) representing 1 image. + This function is used to sample and load images from disk/RAM and is only called in _get_collated_batch() + The length of the list is equal to the (# of training images) / (num_workers) + + Note: The `indices` given to _get_collated_batch() are the `indices` passed to _get_batch_list(). These `indices` + are either set to the entire dataset if we are not loading from disk or some partiton of dataset whose size + is dependent on self.num_images_to_sample_from if we are loading from disk. + """ + assert isinstance(self.input_dataset, Sized) + batch_list = [] + results = [] + + num_threads = int(self.num_image_load_threads) if self.load_from_disk else 4 * int(self.num_image_load_threads) + num_threads = min(num_threads, multiprocessing.cpu_count() - 1) + num_threads = max(num_threads, 1) + + # NB: this is I/O heavy because we are going to disk and reading an image filename + # hence multi-threaded inside the worker + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + for idx in indices: + res = executor.submit(self.input_dataset.__getitem__, idx) + results.append(res) + for res in results: + batch_list.append(res.result()) + + return batch_list + + def _get_collated_batch(self, indices): + """Takes the output of _get_batch_list and collates them with nerfstudio_collate() or variable_res_collate() + Note: dict is an instance of collections.abc.Mapping + + The resulting output is collated_batch: a dictionary with dict_keys(['image_idx', 'image']) + - collated_batch['image_idx'] is tensor with shape torch.Size([per_worker]) + - collated_batch['image'] is tensor with shape torch.Size([per_worker, height, width, 3]) + """ + batch_list = self._get_batch_list(indices=indices) + collated_batch = self.collate_fn(batch_list) + collated_batch = get_dict_to_torch( + collated_batch, device=self.device, exclude=self.exclude_batch_keys_from_device + ) + return collated_batch + + 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 + 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(self.sampling_seed) + num_rays_per_loop = self.num_rays_per_batch # default train_num_rays_per_batch is 4096 + + # each worker has its own pixel sampler + worker_pixel_sampler = self._get_pixel_sampler(self.input_dataset, num_rays_per_loop) + + # the generated RayBundles will be on the same device as self.input_dataset.cameras (CPU) + self.ray_generator = RayGenerator(self.input_dataset.cameras) + + i = 0 + true_random = random.Random(worker_info.id) if worker_info is not None else r + # We offset the value of repeat so that they're not all running out of images at once + repeat_offset_max = 10 if worker_info is not None else 1 + repeat_offset = true_random.randint(0, repeat_offset_max) + while True: + if not self.load_from_disk: + collated_batch = self._cached_collated_batch + elif i % (self.num_times_to_repeat_images + repeat_offset) == 0: + r.shuffle(worker_indices) + repeat_offset = true_random.randint(0, repeat_offset_max) + if self.num_images_to_sample_from == float("inf"): + # if infinity, the worker gets all available indices in its partition + image_indices = worker_indices + else: + # get a total of 'num_images_to_sample_from' image indices + image_indices = worker_indices[: self.num_images_to_sample_from] + + collated_batch = self._get_collated_batch(image_indices) + i += 1 + """ + 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_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 + # Note: collated_batch["image"].get_device() will return CPU if self.exclude_batch_keys_from_device contains 'image' + ray_indices = batch["indices"] + # the ray_bundle is on the GPU; batch["image"] is on the CPU, here we move it to the GPU + ray_bundle = self.ray_generator(ray_indices).to(self.device) + if self.custom_ray_processor: + ray_bundle, batch = self.custom_ray_processor(ray_bundle, batch) + + yield ray_bundle, batch + + +class ImageBatchStream(IterableDataset): + """ + A wrapper of InputDataset that outputs undistorted full images and cameras. This makes the + datamanager more lightweight since we don't have to do generate rays. Useful for full-image + training e.g. rasterization pipelines + """ + + def __init__( + self, + input_dataset: InputDataset, + sampling_seed: int = 3301, + cache_images_type: Literal["uint8", "float32"] = "float32", + device: Union[torch.device, str] = "cpu", + custom_image_processor: Optional[Callable[[Cameras, Dict], Tuple[Cameras, Dict]]] = None, + ): + self.input_dataset = input_dataset + self.sampling_seed = sampling_seed + self.cache_images_type = cache_images_type + self.device = device + self.custom_image_processor = custom_image_processor + + def __iter__(self): + dataset_indices = list(range(len(self.input_dataset))) + worker_info = get_worker_info() + if worker_info is not None: # if we have multiple processes + per_worker = int(math.ceil(len(dataset_indices) / 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 + worker_indices = dataset_indices[ + slice_start : slice_start + per_worker + ] # the indices of the datapoints in the dataset this worker will load + r = random.Random(self.sampling_seed) + r.shuffle(worker_indices) + i = 0 # i refers to what image index we are outputting: i=0 => we are yielding our first image,camera + + while True: + if i >= len(worker_indices): + # if we've iterated through all the worker's partition of images, we need to reshuffle + r.shuffle(worker_indices) + i = 0 + idx = worker_indices[i] # idx refers to the actual datapoint index this worker will retrieve + camera, data = undistort_view(idx, self.input_dataset, self.cache_images_type) # type: ignore + if camera.metadata is None: + camera.metadata = {} + camera.metadata["cam_idx"] = idx + + # Apply custom processing if provided + if self.custom_image_processor: + camera, data = self.custom_image_processor(camera, data) + + 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) + yield camera, data + + class EvalDataloader(DataLoader): """Evaluation dataloader base class diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 12796a540..a653e1de8 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -234,7 +234,7 @@ class Trainer: """Train the model.""" assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs" if hasattr(self.pipeline.datamanager, "train_dataparser_outputs"): - self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform( + self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform( # type: ignore self.base_dir / "dataparser_transforms.json" ) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 6695deb8f..00df5f71b 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -258,12 +258,11 @@ class VanillaPipeline(Pipeline): seed_pts = None if ( hasattr(self.datamanager, "train_dataparser_outputs") - and "points3D_xyz" in self.datamanager.train_dataparser_outputs.metadata + and "points3D_xyz" in self.datamanager.train_dataparser_outputs.metadata # type: ignore ): - pts = self.datamanager.train_dataparser_outputs.metadata["points3D_xyz"] - pts_rgb = self.datamanager.train_dataparser_outputs.metadata["points3D_rgb"] + pts = self.datamanager.train_dataparser_outputs.metadata["points3D_xyz"] # type: ignore + pts_rgb = self.datamanager.train_dataparser_outputs.metadata["points3D_rgb"] # type: ignore seed_pts = (pts, pts_rgb) - self.datamanager.to(device) # TODO(ethan): get rid of scene_bounds from the model assert self.datamanager.train_dataset is not None, "Missing input dataset" @@ -427,7 +426,11 @@ class VanillaPipeline(Pipeline): ) image_prefix = "eval" return self.get_average_image_metrics( - self.datamanager.fixed_indices_eval_dataloader, image_prefix, step, output_path, get_std + self.datamanager.fixed_indices_eval_dataloader, # type: ignore + image_prefix, + step, + output_path, + get_std, ) def load_pipeline(self, loaded_state: Dict[str, Any], step: int) -> None: diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index b6404af37..be8d48a24 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -36,9 +36,7 @@ from typing_extensions import Annotated, Literal from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager -from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanager from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager -from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManager from nerfstudio.data.scene_box import OrientedBox from nerfstudio.exporter import texture_utils, tsdf_utils from nerfstudio.exporter.exporter_utils import collect_camera_poses, generate_point_cloud, get_mesh_from_filename @@ -141,7 +139,7 @@ class ExportPointCloud(Exporter): # Increase the batchsize to speed up the evaluation. assert isinstance( pipeline.datamanager, - (VanillaDataManager, ParallelDataManager, FullImageDatamanager, RandomCamerasDataManager), + (VanillaDataManager, ParallelDataManager), ) assert pipeline.datamanager.train_pixel_sampler is not None pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch @@ -329,7 +327,7 @@ class ExportPoissonMesh(Exporter): # Increase the batchsize to speed up the evaluation. assert isinstance( pipeline.datamanager, - (VanillaDataManager, ParallelDataManager, FullImageDatamanager, RandomCamerasDataManager), + (VanillaDataManager, ParallelDataManager), ) assert pipeline.datamanager.train_pixel_sampler is not None pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index bc58043aa..8db05c0dd 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -372,9 +372,7 @@ class Viewer: # TODO this fn accounts for like ~5% of total train time # Update the train camera locations based on optimization assert self.camera_handles is not None - if hasattr(self.pipeline.datamanager, "train_camera_optimizer"): - camera_optimizer = self.pipeline.datamanager.train_camera_optimizer - elif hasattr(self.pipeline.model, "camera_optimizer"): + if hasattr(self.pipeline.model, "camera_optimizer"): camera_optimizer = self.pipeline.model.camera_optimizer else: return diff --git a/tests/test_nerfacto_integration.py b/tests/test_nerfacto_integration.py new file mode 100644 index 000000000..64b2e528f --- /dev/null +++ b/tests/test_nerfacto_integration.py @@ -0,0 +1,89 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import subprocess +import sys +from typing import Literal + + +def run_command_with_console_output(command, stop_on_output=None): + try: + process = subprocess.Popen( + command, + shell=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, # Redirect stderr to stdout + text=True, + bufsize=1, # Line buffered + universal_newlines=True, + ) + + # Read and print output in real-time + for line in process.stdout: + print(line, end="") + sys.stdout.flush() # Ensure output is printed immediately + + if stop_on_output and stop_on_output in line: + print(f"\nDetected '{stop_on_output}'. Stopping the process.") + process.terminate() + break + + return_code = process.wait() + if return_code != 0: + print(f"Command failed with return code {return_code}") + except Exception as e: + print(f"An error occurred: {e}") + + +def run_ns_download_data(scene: Literal["poster", "dozer", "desolation"]): + command = f"ns-download-data nerfstudio --capture-name={scene}" + run_command_with_console_output(command) + + +def run_ns_train_nerfacto(scene: Literal["poster", "dozer", "desolation"]): + dataset_path = f"data/nerfstudio/{scene}" + command = f"ns-train nerfacto --data {dataset_path}" + run_command_with_console_output(command, stop_on_output="Checkpoint Directory") + + +def run_ns_eval(scene: Literal["poster", "dozer", "desolation"]): + timestamp = sorted(os.listdir(f"outputs/{scene}/nerfacto/"))[-1] + config_filename = f"outputs/{scene}/nerfacto/{timestamp}/config.yml" + command = f"ns-eval --load-config {config_filename} --output-path nerfacto_integration_eval.json" + run_command_with_console_output(command) + + with open("nerfacto_integration_eval.json", "r") as f: + results = json.load(f) + + assert results["results"]["psnr"] > 20.0, "PSNR was lower than 20" + assert results["results"]["ssim"] > 0.7, "SSIM was lower than 0.7" + + +def main(): + scene = "dozer" # You can change this to "poster" or "desolation" + + print("Starting data download...") + run_ns_download_data(scene) + + print("\nStarting training...") + run_ns_train_nerfacto(scene) + + print("\nStarting evaluation...") + run_ns_eval(scene) + + +if __name__ == "__main__": + main() diff --git a/tests/test_splatfacto_integration.py b/tests/test_splatfacto_integration.py new file mode 100644 index 000000000..c798f914f --- /dev/null +++ b/tests/test_splatfacto_integration.py @@ -0,0 +1,55 @@ +# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from typing import Literal + +from test_nerfacto_integration import run_command_with_console_output, run_ns_download_data + + +def run_ns_train_splatfacto(scene: Literal["poster", "dozer", "desolation"]): + dataset_path = f"data/nerfstudio/{scene}" + command = f"ns-train splatfacto --data {dataset_path}" + run_command_with_console_output(command, stop_on_output="Checkpoint Directory") + + +def run_ns_eval(scene: Literal["poster", "dozer", "desolation"]): + timestamp = sorted(os.listdir(f"outputs/{scene}/splatfacto/"))[-1] + config_filename = f"outputs/{scene}/splatfacto/{timestamp}/config.yml" + command = f"ns-eval --load-config {config_filename} --output-path splatfacto_integration_eval.json" + run_command_with_console_output(command) + + with open("splatfacto_integration_eval.json", "r") as f: + results = json.load(f) + + assert results["results"]["psnr"] > 20.0, "PSNR was lower than 20" + assert results["results"]["ssim"] > 0.7, "SSIM was lower than 0.7" + + +def main(): + scene = "dozer" # You can change this to "poster" or "desolation" + + print("Starting data download...") + run_ns_download_data(scene) + + print("\nStarting training...") + run_ns_train_splatfacto(scene) + + print("\nStarting evaluation...") + run_ns_eval(scene) + + +if __name__ == "__main__": + main()