tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Fix datamanager (#1481)

Browse Source
This commit is contained in:
Matthew Tancik 2023-02-22 16:08:53 -08:00 committed by GitHub
parent ba161a2234
commit 9e68a46abe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
nerfstudio/data/datamanagers/base_datamanager.py
View File

@ -49,8 +49,13 @@ from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserCo
from nerfstudio.data.dataparsers.phototourism_dataparser import (
PhototourismDataParserConfig,
)
from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
from nerfstudio.data.datasets.base_dataset import InputDataset
from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler
from nerfstudio.data.pixel_samplers import (
EquirectangularPixelSampler,
PatchPixelSampler,
PixelSampler,
)
from nerfstudio.data.utils.dataloaders import (
CacheDataloader,
FixedIndicesEvalDataloader,
@ -75,6 +80,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of
"dnerf-data": DNeRFDataParserConfig(),
"phototourism-data": PhototourismDataParserConfig(),
"dycheck-data": DycheckDataParserConfig(),
"sdfstudio-data": SDFStudioDataParserConfig(),
},
prefix_names=False, # Omit prefixes in subcommands themselves.
)
@ -281,6 +287,8 @@ class VanillaDataManagerConfig(InstantiateConfig):
"""The scale factor for scaling spatial data such as images, mask, semantics
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."""
class VanillaDataManager(DataManager): # pylint: disable=abstract-method
@ -344,6 +352,9 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
self, dataset: InputDataset, *args: Any, **kwargs: Any
) -> PixelSampler:
"""Infer pixel sampler to use."""
if self.config.patch_size > 1:
return PatchPixelSampler(*args, **kwargs, patch_size=self.config.patch_size)
# If all images are equirectangular, use equirectangular pixel sampler
is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value
if is_equirectangular.all():
@ -391,12 +402,9 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method
)
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)
self.eval_camera_optimizer = self.config.camera_optimizer.setup(
num_cameras=self.eval_dataset.cameras.size, device=self.device
)
self.eval_ray_generator = RayGenerator(
self.eval_dataset.cameras.to(self.device),
self.eval_camera_optimizer,
self.train_camera_optimizer, # should be shared between train and eval.
)
# for loading full images
self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(