Datamanagers#

Base#

Datamanager.

class nerfstudio.data.datamanagers.base_datamanager.DataManager[source]#

Generic data manager’s abstract class

This version of the data manager is designed be a monolithic way to load data and latents, especially since this may contain learnable parameters which need to be shared across the train and test data managers. The idea is that we have setup methods for train and eval separately and this can be a combined train/eval if you want.

Usage: To get data, use the next_train and next_eval functions. This data manager’s next_train and next_eval methods will return 2 things:

‘rays’: This will contain the rays or camera we are sampling, with latents and
conditionals attached (everything needed at inference)
A “batch” of auxiliary information: This will contain the mask, the ground truth
pixels, etc needed to actually train, score, etc the model

Rationale: Because of this abstraction we’ve added, we can support more NeRF paradigms beyond the vanilla nerf paradigm of single-scene, fixed-images, no-learnt-latents. We can now support variable scenes, variable number of images, and arbitrary latents.

Train Methods:: setup_train: sets up for being used as train iter_train: will be called on __iter__() for the train iterator next_train: will be called on __next__() for the training iterator get_train_iterable: utility that gets a clean pythonic iterator for your training data
Eval Methods:: setup_eval: sets up for being used as eval iter_eval: will be called on __iter__() for the eval iterator next_eval: will be called on __next__() for the eval iterator get_eval_iterable: utility that gets a clean pythonic iterator for your eval data

train_count#

the step number of our train iteration, needs to be incremented manually

Type: int

eval_count#

the step number of our eval iteration, needs to be incremented manually

Type: int

train_dataset#

the dataset for the train dataset

Type: Dataset

eval_dataset#

the dataset for the eval dataset

Type: Dataset

includes_time#

whether the dataset includes time information

Type: bool

Additional attributes specific to each subclass are defined in the setup_train and setup_eval

functions.

forward()[source]#

Blank forward method

This is an nn.Module, and so requires a forward() method normally, although in our case we do not need a forward() method

abstract get_datapath() → Path[source]#: Returns the path to the data. This is used to determine where to save camera paths.

get_eval_iterable(length=-1) → IterableWrapper[source]#

Gets a trivial pythonic iterator that will use the iter_eval and next_eval functions as __iter__ and __next__ methods respectively.

This basically is just a little utility if you want to do something like: | for ray_bundle, batch in datamanager.get_eval_iterable(): | <eval code here> since the returned IterableWrapper is just an iterator with the __iter__ and __next__ methods (methods bound to our DataManager instance in this case) specified in the constructor.

abstract get_eval_rays_per_batch() → int[source]#: Returns the number of rays per batch for evaluation.

abstract get_param_groups() → Dict[str, List[Parameter]][source]#

Get the param groups for the data manager.

Returns: A list of dictionaries containing the data manager’s param groups.

get_train_iterable(length=-1) → IterableWrapper[source]#

Gets a trivial pythonic iterator that will use the iter_train and next_train functions as __iter__ and __next__ methods respectively.

This basically is just a little utility if you want to do something like: | for ray_bundle, batch in datamanager.get_train_iterable(): | <eval code here> since the returned IterableWrapper is just an iterator with the __iter__ and __next__ methods (methods bound to our DataManager instance in this case) specified in the constructor.

abstract get_train_rays_per_batch() → int[source]#: Returns the number of rays per batch for training.

get_training_callbacks(training_callback_attributes: TrainingCallbackAttributes) → List[TrainingCallback][source]#: Returns a list of callbacks to be used during training.

iter_eval()[source]#

The __iter__ function for the eval iterator.

This only exists to assist the get_eval_iterable function, since we need to pass in an __iter__ function for our trivial iterable that we are making.

iter_train()[source]#

The __iter__ function for the train iterator.

This only exists to assist the get_train_iterable function, since we need to pass in an __iter__ function for our trivial iterable that we are making.

abstract next_eval(step: int) → Tuple[Union[RayBundle, Cameras], Dict][source]#

Returns the next batch of data from the eval data manager.

Parameters: step – the step number of the eval image to retrieve
Returns: A tuple of the ray/camera for the image, and a dictionary of additional batch information such as the groundtruth image.

abstract next_eval_image(step: int) → Tuple[Cameras, Dict][source]#

Retrieve the next eval image.

Parameters: step – the step number of the eval image to retrieve
Returns: A tuple of the step number, the ray/camera for the image, and a dictionary of additional batch information such as the groundtruth image.

abstract next_train(step: int) → Tuple[Union[RayBundle, Cameras], Dict][source]#

Returns the next batch of data from the train data manager.

Parameters: step – the step number of the eval image to retrieve
Returns: A tuple of the ray bundle for the image, and a dictionary of additional batch information such as the groundtruth image.

abstract setup_eval()[source]#: Sets up the data manager for evaluation

abstract setup_train()[source]#

Sets up the data manager for training.

Here you will define any subclass specific object attributes from the attribute

class nerfstudio.data.datamanagers.base_datamanager.DataManagerConfig(_target: ~typing.Type = <factory>, data: ~typing.Optional[~pathlib.Path] = None, masks_on_gpu: bool = False, images_on_gpu: bool = False)[source]#

Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers; After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked train/eval data at each iteration

data: Optional[Path] = None#: Source of data, may not be used by all models.

images_on_gpu: bool = False#: Process images on GPU for speed at the expense of memory, if True.

masks_on_gpu: bool = False#: Process masks on GPU for speed at the expense of memory, if True.

class nerfstudio.data.datamanagers.base_datamanager.VanillaDataManager(config: VanillaDataManagerConfig, device: Union[device, str] = 'cpu', test_mode: Literal['test', 'val', 'inference'] = 'val', world_size: int = 1, local_rank: int = 0, **kwargs)[source]#

Basic stored data manager implementation.

This is pretty much a port over from our old dataloading utilities, and is a little jank under the hood. We may clean this up a little bit under the hood with more standard dataloading components that can be strung together, but it can be just used as a black box for now since only the constructor is likely to change in the future, or maybe passing in step number to the next_train and next_eval functions.

Parameters: config – the DataManagerConfig used to instantiate class

create_eval_dataset() → TDataset[source]#: Sets up the data loaders for evaluation

create_train_dataset() → TDataset[source]#: Sets up the data loaders for training

property dataset_type: Type[TDataset]#: Returns the dataset type passed as the generic argument

get_datapath() → Path[source]#: Returns the path to the data. This is used to determine where to save camera paths.

get_eval_rays_per_batch() → int[source]#: Returns the number of rays per batch for evaluation.

get_param_groups() → Dict[str, List[Parameter]][source]#: Get the param groups for the data manager. :returns: A list of dictionaries containing the data manager’s param groups.

get_train_rays_per_batch() → int[source]#: Returns the number of rays per batch for training.

next_eval(step: int) → Tuple[RayBundle, Dict][source]#: Returns the next batch of data from the eval dataloader.

next_eval_image(step: int) → Tuple[Cameras, Dict][source]#

Retrieve the next eval image.

Parameters: step – the step number of the eval image to retrieve
Returns: A tuple of the step number, the ray/camera for the image, and a dictionary of additional batch information such as the groundtruth image.

next_train(step: int) → Tuple[RayBundle, Dict][source]#: Returns the next batch of data from the train dataloader.

setup_eval()[source]#: Sets up the data loader for evaluation

setup_train()[source]#: Sets up the data loaders for training

class nerfstudio.data.datamanagers.base_datamanager.VanillaDataManagerConfig(_target: ~typing.Type = <factory>, data: ~typing.Optional[~pathlib.Path] = None, masks_on_gpu: bool = False, images_on_gpu: bool = False, dataparser: ~typing.Union[~nerfstudio.data.dataparsers.nerfstudio_dataparser.NerfstudioDataParserConfig, ~nerfstudio.data.dataparsers.minimal_dataparser.MinimalDataParserConfig, ~nerfstudio.data.dataparsers.arkitscenes_dataparser.ARKitScenesDataParserConfig, ~nerfstudio.data.dataparsers.blender_dataparser.BlenderDataParserConfig, ~nerfstudio.data.dataparsers.instant_ngp_dataparser.InstantNGPDataParserConfig, ~nerfstudio.data.dataparsers.nuscenes_dataparser.NuScenesDataParserConfig, ~nerfstudio.data.dataparsers.dnerf_dataparser.DNeRFDataParserConfig, ~nerfstudio.data.dataparsers.phototourism_dataparser.PhototourismDataParserConfig, ~nerfstudio.data.dataparsers.dycheck_dataparser.DycheckDataParserConfig, ~nerfstudio.data.dataparsers.scannet_dataparser.ScanNetDataParserConfig, ~nerfstudio.data.dataparsers.sdfstudio_dataparser.SDFStudioDataParserConfig, ~nerfstudio.data.dataparsers.nerfosr_dataparser.NeRFOSRDataParserConfig, ~nerfstudio.data.dataparsers.sitcoms3d_dataparser.Sitcoms3DDataParserConfig, ~nerfstudio.data.dataparsers.scannetpp_dataparser.ScanNetppDataParserConfig, ~nerfstudio.data.dataparsers.colmap_dataparser.ColmapDataParserConfig] = <factory>, train_num_rays_per_batch: int = 1024, train_num_images_to_sample_from: int = -1, train_num_times_to_repeat_images: int = -1, eval_num_rays_per_batch: int = 1024, eval_num_images_to_sample_from: int = -1, eval_num_times_to_repeat_images: int = -1, eval_image_indices: ~typing.Optional[~typing.Tuple[int, ...]] = (0,), collate_fn: ~typing.Callable[[~typing.Any], ~typing.Any] = <function nerfstudio_collate>, camera_res_scale_factor: float = 1.0, patch_size: int = 1, camera_optimizer: ~typing.Optional[~nerfstudio.cameras.camera_optimizers.CameraOptimizerConfig] = None, pixel_sampler: ~nerfstudio.data.pixel_samplers.PixelSamplerConfig = <factory>)[source]#

A basic data manager for a ray-based model

__post_init__()[source]#: Warn user of camera optimizer change.

camera_optimizer: Optional[CameraOptimizerConfig] = None#: Deprecated, has been moved to the model config.

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

static collate_fn(batch: Any, extra_mappings: Optional[Dict[type, Callable]] = None) → Any#: Specifies the collate function to use for the train and eval dataloaders.

dataparser: Union[NerfstudioDataParserConfig, MinimalDataParserConfig, ARKitScenesDataParserConfig, BlenderDataParserConfig, InstantNGPDataParserConfig, NuScenesDataParserConfig, DNeRFDataParserConfig, PhototourismDataParserConfig, DycheckDataParserConfig, ScanNetDataParserConfig, SDFStudioDataParserConfig, NeRFOSRDataParserConfig, Sitcoms3DDataParserConfig, ScanNetppDataParserConfig, ColmapDataParserConfig]#: Specifies the dataparser used to unpack the data.

eval_image_indices: Optional[Tuple[int, ...]] = (0,)#: Specifies the image indices to use during eval; if None, uses all.

eval_num_images_to_sample_from: int = -1#: Number of images to sample during eval iteration.

eval_num_rays_per_batch: int = 1024#: Number of rays per batch to use per eval iteration.

eval_num_times_to_repeat_images: int = -1#: When not evaluating on all images, number of iterations before picking new images. If -1, never pick new images.

patch_size: int = 1#: Size of patch to sample from. If > 1, patch-based sampling will be used.

pixel_sampler: PixelSamplerConfig#: Specifies the pixel sampler used to sample pixels from images.

train_num_images_to_sample_from: int = -1#: Number of images to sample during training iteration.

train_num_rays_per_batch: int = 1024#: Number of rays per batch to use per 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.

nerfstudio.data.datamanagers.base_datamanager.variable_res_collate(batch: List[Dict]) → Dict[source]#

Default collate function for the cached dataloader. :param batch: Batch of samples from the dataset.

Returns: Collated batch.