diff --git a/nerfactory/data/dataloader.py b/nerfactory/data/dataloader.py index 3f2063a34..c86a335ff 100644 --- a/nerfactory/data/dataloader.py +++ b/nerfactory/data/dataloader.py @@ -17,21 +17,26 @@ Data loader. """ import random +from abc import abstractmethod from typing import Dict, List, Optional, Tuple, Union from omegaconf import ListConfig +import torch +from torch import nn from torchtyping import TensorType + from nerfactory.cameras.cameras import Camera, get_camera from nerfactory.cameras.rays import RayBundle from nerfactory.data.image_dataset import ImageDataset from nerfactory.data.image_sampler import ImageSampler from nerfactory.data.pixel_sampler import PixelSampler -from nerfactory.data.structs import DatasetInputs +from nerfactory.data.structs import DatasetInputs, BaseDataContainer from nerfactory.data.utils import get_dataset_inputs_from_dataset_config +from nerfactory.graphs.modules.ray_generator import RayGenerator from nerfactory.utils import profiler from nerfactory.utils.config import DataConfig -from nerfactory.utils.misc import get_dict_to_torch, instantiate_from_dict_config +from nerfactory.utils.misc import IterableWrapper, get_dict_to_torch, instantiate_from_dict_config @profiler.time_function @@ -81,6 +86,132 @@ def setup_dataset_eval(config: DataConfig, test_mode: bool, device: str) -> Tupl return dataset_inputs_eval, dataloader_eval +class AbstractDataloader(nn.Module): + """Second version of the dataloader class (V2) + + This version of the dataloader is designed to subsume both the train and eval dataloaders, + 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 separatley and + this can be a combined train/eval if you want. + + The default should be for the train and eval dataloaders to be the same, but this can be + overridden. This is needed when there are learned parameters either in your data itself or in + the way some of the data (that we will pass to the renderer / Field) was generated. In these + cases, we want these parameters to be accessible by both the train and eval dataloaders, hence + why you would want them to be in the same dataloader. + + An instance where you may want to have only the eval dataloader is if you are doing evaluation + and don't have the dataset used to train the model. + + + Train Methods: + setup_train: sets up for being used as train + iter_train: returns an iterator of the train dataloader + next_train: will be called on __next__() for the training iterator + + Eval Methods: + setup_eval: sets up for being used as eval + iter_eval: returns an iterator of the eval dataloader + next_eval: will be called on __next__() for the eval iterator + + + Args: + image_sampler (ImageSampler): image sampler + pixel_sampler (PixelSampler): pixel sampler + ray_generator (RayGenerator): ray generator + use_train (bool): whether this is being used for training + use_eval (bool): whether this is being used for evaluation + + + + Attributes: + image_sampler (ImageSampler): image sampler + pixel_sampler (PixelSampler): pixel sampler + ray_generator (RayGenerator): ray generator + train_count (int): number of times train has been called + eval_count (int): number of times eval has been called + """ + + def __init__( + self, + image_sampler: ImageSampler, + pixel_sampler: PixelSampler, + ray_generator: RayGenerator, + use_train: bool, + use_eval: bool, + ): + super().__init__() + self.image_sampler = image_sampler + self.pixel_sampler = pixel_sampler + self.iter_image_sampler = iter(self.image_sampler) + self.ray_generator = ray_generator + 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 iter_train(self) -> IterableWrapper: + """Returns an iterator that executes the self.next_train function""" + self.train_count = 0 + return IterableWrapper(self, self.next_train) + + def iter_eval(self) -> IterableWrapper: + """Returns an iterator that executes the self.next_eval function""" + self.eval_count = 0 + return IterableWrapper(self, self.next_eval) + + @abstractmethod + def setup_train(self): + """Sets up the dataloader for training""" + + @abstractmethod + def setup_eval(self): + """Sets up the dataloader for evaluation""" + + @abstractmethod + def next_train(self) -> BaseDataContainer: + """Returns the next batch of data from the train dataloader""" + + @abstractmethod + def next_eval(self) -> BaseDataContainer: + """Returns the next batch of data from the eval dataloader""" + + +class AbstractStoredDataloader(AbstractDataloader): + """Subclass of the new V2 dataloader that is used for when things fit in memory, + and will be stored in the dataloader itself. + + Attributes: + camera_to_world (torch.Tensor): camera to world transformation + intrinsics (torch.Tensor): intrinsics matrix + """ + + camera_to_world: torch.Tensor + intrinsics: torch.Tensor + + # These abstract methods are only copied over to appease the linter + @abstractmethod + def setup_train(self): + """Sets up the dataloader for training""" + + @abstractmethod + def setup_eval(self): + """Sets up the dataloader for evaluation""" + + @abstractmethod + def next_train(self) -> BaseDataContainer: + """Returns the next batch of data from the train dataloader""" + + @abstractmethod + def next_eval(self) -> BaseDataContainer: + """Returns the next batch of data from the eval dataloader""" + + class TrainDataloader: """Training dataloader base class.""" diff --git a/nerfactory/data/structs.py b/nerfactory/data/structs.py index 9b6733d67..437d656cd 100644 --- a/nerfactory/data/structs.py +++ b/nerfactory/data/structs.py @@ -22,6 +22,8 @@ from typing import List, Optional import torch from torchtyping import TensorType +from nerfactory.cameras.rays import RayBundle + @dataclass class PointCloud: @@ -106,3 +108,20 @@ class DatasetInputs: def as_dict(self) -> dict: """Returns the dataclass as a dictionary.""" return vars(self) + + +@dataclass +class BaseDataContainer: + """A container for data that is not specific to any dataset. It should be everything + needed by the renderer and loss calculation. Different datasets and models will probably + need to subclass from this differently, since different Field modules will require different + data. + + Args: + rays (RayBundle): The rays for the image. + ground_truth_pixels (TensorType["num_pixels", 3]): The ground truth pixels for the image. + """ + + rays: RayBundle # Raybundle and the cameras will be merged into one thing in a later PR + cameras: torch.tensor + ground_truth_pixels: Optional[TensorType["num_pixels", 3]] = None diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index b27afd055..f478191c1 100644 --- a/nerfactory/utils/misc.py +++ b/nerfactory/utils/misc.py @@ -19,7 +19,7 @@ Miscellaneous helper code. import hashlib import json from pydoc import locate -from typing import Any, Dict +from typing import Any, Callable, Dict from math import floor, log import torch @@ -93,6 +93,33 @@ def get_hash_str_from_dict(dictionary: Dict[str, Any]) -> str: return dhash.hexdigest() +class IterableWrapper: # pylint: disable=too-few-public-methods + """A helper that will allow an instance of a class to return multiple kinds of iterables bound + to different functions of that class. + + To use this, pass in the instance of the class you want to have multiple kinds of iterables for, + and pass in a method belonging to that class that you want to be the __next__() method of the + iterable. The resulting instantiated object will be an iterable that will use the passed in + class method as the __next__() method. + + Args: + instance (object): instance class we are wrapping + new_next (callable): function that will be called instead as the __next__() + + + Attributes: + instance (object): instance class we are wrapping + new_next (callable): function that will be called instead as the __next__() + """ + + def __init__(self, instance: object, new_next: Callable): + self.instance = instance + self.new_next = new_next + + def __next__(self): + return self.new_next(self.instance) + + def human_format(num): """Format a number in a more human readable way""" units = ["", "K", "M", "B", "T", "P"]