From ffb221a062bca0cd6fe343b8cda409c5f7ae86bf Mon Sep 17 00:00:00 2001 From: AdamRashid96 Date: Wed, 2 Aug 2023 00:37:46 +0000 Subject: [PATCH] adding patch_size back --- nerfstudio/data/datamanagers/base_datamanager.py | 7 +++++++ nerfstudio/data/pixel_samplers.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 62d1b41ef..e115949d5 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -57,6 +57,7 @@ from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import ( PixelSampler, PixelSamplerConfig, + PatchPixelSamplerConfig, ) from nerfstudio.data.utils.dataloaders import ( CacheDataloader, @@ -345,6 +346,8 @@ class VanillaDataManagerConfig(DataManagerConfig): """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.""" pixel_sampler: PixelSamplerConfig = PixelSamplerConfig() """Specifies the pixel sampler used to sample pixels from images.""" @@ -458,6 +461,10 @@ class VanillaDataManager(DataManager, Generic[TDataset]): def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> PixelSampler: """Infer pixel sampler to use.""" + if self.config.patch_size > 1: + 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.") diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 068bbb9d1..bb5e9d5c5 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -380,7 +380,7 @@ class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods device: Union[torch.device, str] = "cpu", ) -> Int[Tensor, "batch_size 3"]: if mask: - # Note: if there is a mask, should switch to the base PixelSampler class + # Note: if there is a mask, should switch to the base PixelSampler class raise NotImplementedError() else: