mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
cdacfdaf2e
* Pushing more documentation * Adding max length option to IterableWrapper, fixing nitpicks * Removing use of word dummy from dataclass
293 lines
13 KiB
Python
293 lines
13 KiB
Python
# Copyright 2022 The Plenoptix Team. 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.
|
|
|
|
"""
|
|
Data loader.
|
|
"""
|
|
|
|
import logging
|
|
from abc import abstractmethod
|
|
from typing import Dict, List, Tuple, Union
|
|
|
|
import torch
|
|
from omegaconf import DictConfig, ListConfig
|
|
from torch import nn
|
|
from torch.nn import Parameter
|
|
|
|
from nerfactory.cameras.rays import RayBundle
|
|
from nerfactory.dataloaders.eval import FixedIndicesEvalDataloader
|
|
from nerfactory.dataloaders.image_dataset import ImageDataset, PanopticImageDataset
|
|
from nerfactory.dataloaders.image_sampler import CacheImageSampler
|
|
from nerfactory.dataloaders.pixel_sampler import PixelSampler
|
|
from nerfactory.dataloaders.structs import DatasetInputs
|
|
from nerfactory.models.modules.ray_generator import RayGenerator
|
|
from nerfactory.utils import profiler
|
|
from nerfactory.utils.config import DataloaderConfig
|
|
from nerfactory.utils.misc import IterableWrapper, instantiate_from_dict_config
|
|
|
|
|
|
class Dataloader(nn.Module):
|
|
"""Generic dataloader's abstract class
|
|
|
|
This version of the dataloader 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 dataloaders. 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 dataloader's next_train and next_eval methods will return 2 things:
|
|
1. A Raybundle: This will contain the rays we are sampling, with latents and
|
|
conditionals attached (everything needed at inference)
|
|
2. A "batch" of auxilury 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
|
|
|
|
Args:
|
|
use_train (bool): whether this dataloader is being used for training
|
|
use_eval (bool): whether this dataloader is being used for evaluation
|
|
|
|
Attributes:
|
|
use_train (bool): whether or not we are using train
|
|
use_eval (bool): whether or not we are using eval
|
|
train_count (int): the step number of our train iteration, needs to be incremented manually
|
|
eval_count (int): the step number of our eval iteration, needs to be incremented manually
|
|
|
|
Additional attributes specific to each subclass are defined in the setuo_train and setup_eval
|
|
functions.
|
|
|
|
"""
|
|
|
|
def __init__(self, use_train: bool, use_eval: bool):
|
|
"""Constructor for the Dataloader class.
|
|
|
|
Subclassed Dataloaders 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__(use_train, use_eval) BEFORE you initialize any
|
|
nn.Modules or nn.Parameters, but AFTER you've already set all the attributes you need
|
|
for the setup functions."""
|
|
super().__init__()
|
|
self.use_train = use_train
|
|
self.use_eval = use_eval
|
|
self.train_count = 0
|
|
self.eval_count = 0
|
|
assert use_train or use_eval
|
|
if use_train:
|
|
self.setup_train()
|
|
if use_eval:
|
|
self.setup_eval()
|
|
|
|
def forward(self):
|
|
"""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"""
|
|
raise NotImplementedError
|
|
|
|
def iter_train(self):
|
|
"""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."""
|
|
self.train_count = 0
|
|
|
|
def iter_eval(self):
|
|
"""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."""
|
|
self.eval_count = 0
|
|
|
|
def get_train_iterable(self, length=-1) -> IterableWrapper:
|
|
"""Gets a trivial pythonic iterator that will use the iter_train and next_train functions
|
|
as __iter__ and __next__ methods respectivley.
|
|
|
|
This basically is just a little utility if you want to do something like:
|
|
| for ray_bundle, batch in dataloader.get_train_iterable():
|
|
| <eval code here>
|
|
since the returned IterableWrapper is just an iterator with the __iter__ and __next__
|
|
methods (methods bound to our Dataloader instance in this case) specified in the constructor.
|
|
"""
|
|
return IterableWrapper(self.iter_train, self.next_train, length)
|
|
|
|
def get_eval_iterable(self, length=-1) -> IterableWrapper:
|
|
"""Gets a trivial pythonic iterator that will use the iter_eval and next_eval functions
|
|
as __iter__ and __next__ methods respectivley.
|
|
|
|
This basically is just a little utility if you want to do something like:
|
|
| for ray_bundle, batch in dataloader.get_eval_iterable():
|
|
| <eval code here>
|
|
since the returned IterableWrapper is just an iterator with the __iter__ and __next__
|
|
methods (methods bound to our Dataloader instance in this case) specified in the constructor.
|
|
"""
|
|
return IterableWrapper(self.iter_eval, self.next_eval, length)
|
|
|
|
@abstractmethod
|
|
def setup_train(self):
|
|
"""Sets up the dataloader for training.
|
|
|
|
Here you will define any subclass specific object attributes from the attribute"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def setup_eval(self):
|
|
"""Sets up the dataloader for evaluation"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def next_train(self) -> Tuple:
|
|
"""Returns the next batch of data from the train dataloader.
|
|
|
|
This will be a tuple of all the information that this dataloader outputs.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def next_eval(self) -> Tuple:
|
|
"""Returns the next batch of data from the eval dataloader.
|
|
|
|
This will be a tuple of all the information that this dataloader outputs.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use
|
|
"""Get the param groups for the dataloader.
|
|
|
|
Returns:
|
|
A list of dictionaries containing the dataloader's param groups.
|
|
"""
|
|
return {}
|
|
|
|
|
|
class VanillaDataloader(Dataloader): # pylint: disable=abstract-method
|
|
"""Basic stored dataloader 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."""
|
|
|
|
def __init__(
|
|
self,
|
|
image_dataset_type: str,
|
|
train_datasetinputs: DatasetInputs,
|
|
train_num_rays_per_batch: int,
|
|
train_num_images_to_sample_from: int,
|
|
eval_datasetinputs: DatasetInputs,
|
|
eval_image_indices: Union[List[int], ListConfig], # TODO(ethan): get rid of this hydra ListConfig nonsense
|
|
eval_num_rays_per_chunk: int,
|
|
device: Union[torch.device, str] = "cpu",
|
|
**kwargs # pylint: disable=unused-argument
|
|
):
|
|
self.image_dataset_type = image_dataset_type
|
|
self.train_datasetinputs = train_datasetinputs
|
|
self.train_num_rays_per_batch = train_num_rays_per_batch
|
|
self.train_num_images_to_sample_from = train_num_images_to_sample_from
|
|
self.eval_datasetinputs = eval_datasetinputs
|
|
self.eval_image_indices = eval_image_indices
|
|
self.eval_num_rays_per_chunk = eval_num_rays_per_chunk
|
|
self.device = device
|
|
use_train = self.train_datasetinputs is not None
|
|
use_eval = self.eval_datasetinputs is not None
|
|
super().__init__(use_train, use_eval)
|
|
|
|
def setup_train(self):
|
|
"""Sets up the dataloader for training"""
|
|
if self.image_dataset_type == "rgb":
|
|
self.train_image_dataset = ImageDataset(**self.train_datasetinputs.as_dict())
|
|
elif self.image_dataset_type == "panoptic":
|
|
self.train_image_dataset = PanopticImageDataset(**self.train_datasetinputs.as_dict())
|
|
self.train_image_sampler = CacheImageSampler(
|
|
self.train_image_dataset, num_images_to_sample_from=self.train_num_images_to_sample_from, device=self.device
|
|
) # TODO(ethan): pass this in
|
|
self.iter_train_image_sampler = iter(self.train_image_sampler)
|
|
self.train_pixel_sampler = PixelSampler(self.train_num_rays_per_batch)
|
|
self.train_ray_generator = RayGenerator(
|
|
self.train_datasetinputs.intrinsics, self.train_datasetinputs.camera_to_world
|
|
)
|
|
|
|
def setup_eval(self):
|
|
"""Sets up the dataloader for evaluation"""
|
|
if self.image_dataset_type == "rgb":
|
|
self.eval_image_dataset = ImageDataset(**self.eval_datasetinputs.as_dict())
|
|
elif self.image_dataset_type == "panoptic":
|
|
self.eval_image_dataset = PanopticImageDataset(**self.eval_datasetinputs.as_dict())
|
|
self.eval_dataloader = FixedIndicesEvalDataloader(
|
|
image_dataset=self.eval_image_dataset,
|
|
intrinsics=self.eval_datasetinputs.intrinsics,
|
|
camera_to_world=self.eval_datasetinputs.camera_to_world,
|
|
num_rays_per_chunk=self.eval_num_rays_per_chunk,
|
|
image_indices=self.eval_image_indices,
|
|
device=self.device,
|
|
)
|
|
|
|
def next_train(self) -> Tuple[RayBundle, Dict]:
|
|
"""Returns the next batch of data from the train dataloader.
|
|
The RayBundle can be shaped in whatever way."""
|
|
self.train_count += 1
|
|
image_batch = next(self.iter_train_image_sampler)
|
|
batch = self.train_pixel_sampler.sample(image_batch)
|
|
ray_indices = batch["indices"]
|
|
ray_bundle = self.train_ray_generator.forward(ray_indices)
|
|
return ray_bundle, batch
|
|
|
|
def next_eval(self) -> Tuple[RayBundle, Dict]:
|
|
"""Returns the next batch of data from the eval dataloader.
|
|
The RayBundle should be shaped like an image."""
|
|
self.eval_count += 1
|
|
camera_ray_bundle, batch = next(self.eval_dataloader)
|
|
return camera_ray_bundle, batch
|
|
|
|
|
|
@profiler.time_function
|
|
def setup_dataloader(config: DataloaderConfig, device: str, test_mode=False) -> Dataloader:
|
|
"""Setup the dataloader."""
|
|
dataset_train = instantiate_from_dict_config(config.train_dataset)
|
|
train_datasetinputs = dataset_train.get_dataset_inputs(split="train")
|
|
if config.eval_dataset is not None:
|
|
dataset_eval = instantiate_from_dict_config(config.eval_dataset)
|
|
else:
|
|
logging.info("No eval dataset specified so using train dataset for eval.")
|
|
dataset_eval = dataset_train
|
|
eval_datasetinputs = dataset_eval.get_dataset_inputs(split="val" if not test_mode else "test")
|
|
|
|
dataloader: Dataloader = instantiate_from_dict_config(
|
|
DictConfig(config),
|
|
train_datasetinputs=train_datasetinputs,
|
|
eval_datasetinputs=eval_datasetinputs,
|
|
device=device,
|
|
)
|
|
dataloader.to(device)
|
|
return dataloader
|