tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-14 20:03:49 +08:00
Code Activity

Dataloading Revamp (#3216)

Browse Source 
* initial debugging and testing works

* pwais changes with RayBatchStream to alleviate training

* few bugs to iron out with multiprocessing, specifically pickled collate_fn

* working version of RayBatchStream

* additional docstrings

* cleanup

* much more documentation

* successfully trained AEA-script2_seq2 closed_loop without OOM

* porting over aria dataset-size feature

* added logic to handle eviction of a worker's cached_collated_batch

* antonio's implementation of stream batches

* training on a dataset with 4000 images works!

* some configuration speedups, loops aren't actually needed!

* quick fix adjustment to aria

* removed unnecessary looping

* much faster training when adding i variable to collate every 5 ray bundles

* cleanup unnecssary variables in Dataloader

* further cleanup

* adding caching of compressed images to RAM to reduce disk bottleneck

* added caching to RAM for masks

* found fast way to collate - many tricks applied

* quick update to aria to test on different datasets

* cleaned up the accelerated pil_to_numpy function

* cleaning up PR

* this commit was used to generate the time metrics and profiling metrics

* REAL commit used to run tests

* testing with nerfacto-big

* generated RayBundle collate and converting images from uint8s to float32 on GPU tests

* updating nerfacto to support uint8 easily, will need to figure out a way to contain this within the datamanager API

* datamanager updates, both splat and nerf

* must use writeable arrays because torch requires them

* cleaned up base_dataset, added pickle to utils, more code in full_image, and cleaner desc for base_datamanager

* lots of process on a parallel FullImageDatamanger

* can train big splats with pre-assertion hack or ROI hack and 0 workers

* fixed all undistortion issues with ParallelImageDatamanager

* adding some downsampling and parallel tests with splatfacto!

* deleted commented code in dataloaders.py and added bugfix to shuffling

* testing splatfacto-big

* cleaned up base_pipeline.py

* cleaned up base_pipeline.py ACTUALLY THIS TIME, forgot to save last time

* cleaned up a lot of code

* process_project_aria back to main branch and some cleanup in full_image_datamanager

* clarifying docstrings

* further PR cleanup

* updating models

* further cleanup

* removed caching of images into bytestrings

* adding caching of compressed images to RAM, forgot that hardware matters

* removing oom methods, adding the ability to add a flag to dataloading

* removed CacheDataloader, moved RayBatchStream to dataloaders.py, new vanilla_datamanager rewritten

* fixing base_piplines, deleting a weird datamanager_configs file that was accidently created

* cleaning up next_train

* replaced parallel datamanager with new datamanager

* reverted the original base_datamanager.py, new datamanager replaced parallel_datamanager.py

* modified VanillaConfig, but VanillaDataManager is the same as before

* cleaning up, 2 datamanagers now - original and new parallel one

* able to train with new nerfstudio dataloader now

* side by side datamanagers, moved tons of logic into dataloaders.py and created new files for our parallel datamangers

* added custom ray processing API to support implementations like LERF, cleaned up FullImageDatamanager to original because of new ParallelImageDatamanger

* adding functionality for ns-eval by adding FixedIndicesEvalDataloader to the setup_eval

* adding both ray API and image-view API to datamanagers for custom parallelization

* updating splatfacto config for 4k tests

* updating docstrings to be more descriptive

* new datamanager API breaks when setup_eval() has multiple workers, not sure why but single worker will have to do

* adding custom_view_processor to ImageBatchStream

* reverting full_images_datamanager to main branch

* removing nn.Module inheritance from Datamanager class

* don't need to move datamanger to device anymore since Datamanager is not a subclass of nn.Module

* finished integration test with nerfacto

* simplified config variables, integrated the parallelism/disk-data-loading all into one datamanager

* updated the splatfacto config to be simpler with the dataloading and now uses FullImageDatamanager (which has been changed)

* style checks and some cleanup

* new splatfacto test, cleaning up nerfacto integration test

* removing redundant parallel_full_images_datamaanger, as the OG full_image_datamanager now has full parallelized support

* ruff linting and pyright fixing

* further pyright fixing

* another pyright fixing

* fixing pyright error, camera optimization no longer part of datamanager

* fixing one pyright

* fixing dataloading error when camera is not undistorted with dataloader

* fixing comments and updating style

* undoing a style change i made

* undoing another style change i made by accident

* fixing slow runtime

* fixing a more general camera undistortion bug

* move images to device properly

* minor improvements

* add print statement about >500 images, cleanup method configs

* make method configs consistent across nerfacto models

* adding description comments

* updating description

* resolving some pyright issues with export.py, explained in PR desc

* fixing pyright issues in base_pipeline.py

* ran pyright on exporter and base_pipeline.py without issues

* adding a git ignore to a clearly checked pyright issue

* typo

* fixing most ns-dev-test cases

* cleanup, passing final ns-dev-test

* oops, accidentally pushed the deletion of a docstring, undoing that

* another cleanup

* some fixes to eval pipeline

* lint

* add asserts for spawn

* lint

* cleaning up import statements in parallel_datamanager.py

* adding new developer documentation if users would like to migrate their custom datamanagers to support new features

* removing unnecessary to_device no-op

* further updates to documentation

* lint

* more docs

* docs

* remove comment

* add docs, fix depth dataset with parallel datamanager, fix mask sampling bug

* remove profiling

* more profile removal

* custom_view_processor->custom_image_processor

* doc clarification

* datamanager doc nit

* whitespace

* nits

* remove stuff from __post_init__, tune num workers more, add random offset in raybatchstream

* removing unnecessary assertion, updating docstring because DataManager is no longer an nn.Module

* clarifying configuration with num_images_to_sample_from and num_times_to_repeat_images, cleaning up functions

* adding logic so that nerfacto users can load_from_disk and customize image batch sizes and repeat parameters

* ruff formatting! whoops forgot to format

* fixing logic, now if users set load_from_disk to true, datamanager will use 50 and 10. If users set it and specify their own values, we support that as well

* adding separate datamanager config so that target can be removed in method_configs

---------

Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com>
Co-authored-by: Brent Yi <yibrenth@gmail.com>
This commit is contained in:
Anthony ZhangJustin KerrBrent Yi
 authored and GitHub committed 2025-01-31 10:40:01 -08:00
1 parent 189328ec79
commit 4a3e3e6403
16 files changed
+1084 -471
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nerfstudio/configs/method_configs.py
+4 -4
View File
@@ -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,
nerfstudio/data/datamanagers/base_datamanager.py
+23 -47
View File
@@ -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."""
nerfstudio/data/datamanagers/full_images_datamanager.py
+85 -209
View File
@@ -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
nerfstudio/data/datamanagers/parallel_datamanager.py
+113 -170
View File
@@ -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
nerfstudio/data/datasets/base_dataset.py
+33 -8
View File
@@ -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]}"
nerfstudio/data/datasets/depth_dataset.py
+5 -3
View File
@@ -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)
nerfstudio/data/pixel_samplers.py
+1 -1
View File
@@ -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
nerfstudio/data/utils/data_utils.py
+41 -3
View File
@@ -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
nerfstudio/data/utils/dataloaders.py
+516 -10
View File
@@ -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
nerfstudio/engine/trainer.py
+1 -1
View File
@@ -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"
)
nerfstudio/pipelines/base_pipeline.py
+8 -5
View File
@@ -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:
nerfstudio/scripts/exporter.py
+2 -4
View File
@@ -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
nerfstudio/viewer/viewer.py
+1 -3
View File
@@ -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