mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 10:55:57 +08:00
Julian Knodt
35610cb9e2
* Add dnerf dataparser * Switch time to be stored per camera I realized it would be better downstream to store times on the ray bundle, and in order to store them on the ray bundle it's much easier to store them on the camera. Rather than having to revert this change, I just made it here.
399 lines
17 KiB
Python
399 lines
17 KiB
Python
# Copyright 2022 The Nerfstudio 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.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import abstractmethod
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, List, Optional, Tuple, Type, Union
|
|
|
|
import torch
|
|
import tyro
|
|
from rich.progress import Console
|
|
from torch import nn
|
|
from torch.nn import Parameter
|
|
from torch.utils.data import Dataset
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
|
|
from nerfstudio.cameras.rays import RayBundle
|
|
from nerfstudio.configs.base_config import InstantiateConfig
|
|
from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
|
|
from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
|
|
from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
|
|
from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
|
|
InstantNGPDataParserConfig,
|
|
)
|
|
from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
|
|
from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig
|
|
from nerfstudio.data.pixel_samplers import PixelSampler
|
|
from nerfstudio.data.utils.dataloaders import (
|
|
CacheDataloader,
|
|
FixedIndicesEvalDataloader,
|
|
RandIndicesEvalDataloader,
|
|
)
|
|
from nerfstudio.data.utils.datasets import InputDataset
|
|
from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes
|
|
from nerfstudio.model_components.ray_generators import RayGenerator
|
|
from nerfstudio.utils.misc import IterableWrapper
|
|
|
|
CONSOLE = Console(width=120)
|
|
|
|
AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands.
|
|
tyro.extras.subcommand_type_from_defaults(
|
|
{
|
|
"nerfstudio-data": NerfstudioDataParserConfig(),
|
|
"blender-data": BlenderDataParserConfig(),
|
|
"friends-data": FriendsDataParserConfig(),
|
|
"instant-ngp-data": InstantNGPDataParserConfig(),
|
|
"record3d-data": Record3DDataParserConfig(),
|
|
"dnerf-data": DNeRFDataParserConfig(),
|
|
},
|
|
prefix_names=False, # Omit prefixes in subcommands themselves.
|
|
)
|
|
]
|
|
"""Union over possible dataparser types, annotated with metadata for tyro. This is the
|
|
same as the vanilla union, but results in shorter subcommand names."""
|
|
|
|
|
|
class DataManager(nn.Module):
|
|
"""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:
|
|
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
|
|
|
|
|
|
Attributes:
|
|
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
|
|
train_dataset (Dataset): the dataset for the train dataset
|
|
eval_dataset (Dataset): the dataset for the eval dataset
|
|
|
|
Additional attributes specific to each subclass are defined in the setup_train and setup_eval
|
|
functions.
|
|
|
|
"""
|
|
|
|
train_dataset: Optional[Dataset] = None
|
|
eval_dataset: Optional[Dataset] = None
|
|
train_sampler: Optional[DistributedSampler] = None
|
|
eval_sampler: Optional[DistributedSampler] = None
|
|
|
|
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."""
|
|
super().__init__()
|
|
self.train_count = 0
|
|
self.eval_count = 0
|
|
if self.train_dataset:
|
|
self.setup_train()
|
|
if self.eval_dataset:
|
|
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 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.
|
|
"""
|
|
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 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.
|
|
"""
|
|
return IterableWrapper(self.iter_eval, self.next_eval, length)
|
|
|
|
@abstractmethod
|
|
def setup_train(self):
|
|
"""Sets up the data manager for training.
|
|
|
|
Here you will define any subclass specific object attributes from the attribute"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def setup_eval(self):
|
|
"""Sets up the data manager for evaluation"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def next_train(self, step: int) -> Tuple:
|
|
"""Returns the next batch of data from the train data manager.
|
|
|
|
This will be a tuple of all the information that this data manager outputs.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def next_eval(self, step: int) -> Tuple:
|
|
"""Returns the next batch of data from the eval data manager.
|
|
|
|
This will be a tuple of all the information that this data manager outputs.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def next_eval_image(self, step: int) -> Tuple:
|
|
"""Returns the next eval image."""
|
|
raise NotImplementedError
|
|
|
|
def get_training_callbacks( # pylint:disable=no-self-use
|
|
self, training_callback_attributes: TrainingCallbackAttributes # pylint: disable=unused-argument
|
|
) -> List[TrainingCallback]:
|
|
"""Returns a list of callbacks to be used during training."""
|
|
return []
|
|
|
|
@abstractmethod
|
|
def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use
|
|
"""Get the param groups for the data manager.
|
|
|
|
Returns:
|
|
A list of dictionaries containing the data manager's param groups.
|
|
"""
|
|
return {}
|
|
|
|
|
|
@dataclass
|
|
class VanillaDataManagerConfig(InstantiateConfig):
|
|
"""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
|
|
"""
|
|
|
|
_target: Type = field(default_factory=lambda: VanillaDataManager)
|
|
"""Target class to instantiate."""
|
|
dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig()
|
|
"""Specifies the dataparser used to unpack the data."""
|
|
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."""
|
|
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
|
|
"""Number of images to sample during 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."""
|
|
eval_image_indices: Optional[Tuple[int, ...]] = (0,)
|
|
"""Specifies the image indices to use during eval; if None, uses all."""
|
|
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
|
|
"""Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
|
|
Record3D."""
|
|
|
|
|
|
class VanillaDataManager(DataManager): # pylint: disable=abstract-method
|
|
"""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.
|
|
|
|
Args:
|
|
config: the DataManagerConfig used to instantiate class
|
|
"""
|
|
|
|
config: VanillaDataManagerConfig
|
|
train_dataset: InputDataset
|
|
eval_dataset: InputDataset
|
|
|
|
def __init__(
|
|
self,
|
|
config: VanillaDataManagerConfig,
|
|
device: Union[torch.device, str] = "cpu",
|
|
test_mode: bool = False,
|
|
world_size: int = 1,
|
|
local_rank: int = 0,
|
|
**kwargs, # pylint: disable=unused-argument
|
|
):
|
|
self.config = config
|
|
self.device = device
|
|
self.world_size = world_size
|
|
self.local_rank = local_rank
|
|
self.sampler = None
|
|
|
|
self.train_dataset = InputDataset(config.dataparser.setup().get_dataparser_outputs(split="train"))
|
|
self.eval_dataset = InputDataset(
|
|
config.dataparser.setup().get_dataparser_outputs(split="val" if not test_mode else "test")
|
|
)
|
|
super().__init__()
|
|
|
|
def setup_train(self):
|
|
"""Sets up the data loaders for training"""
|
|
assert self.train_dataset is not None
|
|
CONSOLE.print("Setting up training dataset...")
|
|
self.train_image_dataloader = CacheDataloader(
|
|
self.train_dataset,
|
|
num_images_to_sample_from=self.config.train_num_images_to_sample_from,
|
|
num_times_to_repeat_images=self.config.train_num_times_to_repeat_images,
|
|
device=self.device,
|
|
num_workers=self.world_size * 4,
|
|
pin_memory=True,
|
|
)
|
|
self.iter_train_image_dataloader = iter(self.train_image_dataloader)
|
|
self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch)
|
|
self.train_camera_optimizer = self.config.camera_optimizer.setup(
|
|
num_cameras=self.train_dataset.dataparser_outputs.cameras.size, device=self.device
|
|
)
|
|
self.train_ray_generator = RayGenerator(
|
|
self.train_dataset.dataparser_outputs.cameras.to(self.device),
|
|
self.train_camera_optimizer,
|
|
)
|
|
|
|
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,
|
|
device=self.device,
|
|
num_workers=self.world_size * 4,
|
|
pin_memory=True,
|
|
)
|
|
self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
|
|
self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch)
|
|
self.eval_ray_generator = RayGenerator(
|
|
self.eval_dataset.dataparser_outputs.cameras.to(self.device),
|
|
self.train_camera_optimizer, # should be shared between train and eval.
|
|
)
|
|
# for loading full images
|
|
self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
|
|
input_dataset=self.eval_dataset,
|
|
device=self.device,
|
|
num_workers=self.world_size * 4,
|
|
)
|
|
self.eval_dataloader = RandIndicesEvalDataloader(
|
|
input_dataset=self.eval_dataset,
|
|
image_indices=self.config.eval_image_indices,
|
|
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 train dataloader."""
|
|
self.train_count += 1
|
|
image_batch = next(self.iter_train_image_dataloader)
|
|
batch = self.train_pixel_sampler.sample(image_batch)
|
|
ray_indices = batch["indices"]
|
|
ray_bundle = self.train_ray_generator(ray_indices)
|
|
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)
|
|
batch = self.eval_pixel_sampler.sample(image_batch)
|
|
ray_indices = batch["indices"]
|
|
ray_bundle = self.eval_ray_generator(ray_indices)
|
|
return ray_bundle, batch
|
|
|
|
def next_eval_image(self, step: int) -> Tuple[int, RayBundle, Dict]:
|
|
for camera_ray_bundle, batch in self.eval_dataloader:
|
|
assert camera_ray_bundle.camera_indices is not None
|
|
image_idx = int(camera_ray_bundle.camera_indices[0, 0, 0])
|
|
return image_idx, camera_ray_bundle, batch
|
|
raise ValueError("No more eval images")
|
|
|
|
def get_param_groups(self) -> Dict[str, List[Parameter]]: # pylint: disable=no-self-use
|
|
"""Get the param groups for the data manager.
|
|
Returns:
|
|
A list of dictionaries containing the data manager's param groups.
|
|
"""
|
|
param_groups = {}
|
|
|
|
camera_opt_params = list(self.train_camera_optimizer.parameters())
|
|
if self.config.camera_optimizer.mode != "off":
|
|
assert len(camera_opt_params) > 0
|
|
param_groups[self.config.camera_optimizer.param_group] = camera_opt_params
|
|
else:
|
|
assert len(camera_opt_params) == 0
|
|
|
|
return param_groups
|