From 7e90cb798938313cd91eca7aaf67b220b1348089 Mon Sep 17 00:00:00 2001 From: AdamRashid96 <71362382+AdamRashid96@users.noreply.github.com> Date: Fri, 4 Aug 2023 07:27:35 -0700 Subject: [PATCH] PixelSampler Fixes (#2298) * monodepth implemented * device * pair pixel sampler * docstring * removing unused imports * fix bugs * cache in cache dir, no longer edit json file * moving to depth nerfacto * fixed bug using median depth by accident * removing depthdataset from nerfacto * added pixelsamplerconfig * added equirectangular warning back * bug fix * delete sample_config.py + small fixes * rename to pixel_sampler * cleaning * test * passing num_rays_per_batch properly * fixed transforms typo * renaming in nerfacto * remove depth files from test data * removed defaulting to base pixel sampler with mask * adding patch_size back * Warning for notimplemented error and patch_size fix * fix config.pixel_sampler check --------- Co-authored-by: Matthew Tancik Co-authored-by: Ethan Weber Co-authored-by: Ethan Weber --- nerfstudio/data/datamanagers/base_datamanager.py | 2 +- nerfstudio/data/pixel_samplers.py | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index e115949d5..52fd0af0f 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -461,7 +461,7 @@ 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: + 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 ) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index bb5e9d5c5..41d457458 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -318,8 +318,9 @@ class PatchPixelSampler(PixelSampler): device: Union[torch.device, str] = "cpu", ) -> Int[Tensor, "batch_size 3"]: if isinstance(mask, Tensor): - # Note: if there is a mask, should switch to the base PixelSampler class - raise NotImplementedError() + raise NotImplementedError( + "Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead." + ) else: sub_bs = batch_size // (self.config.patch_size**2) indices = torch.rand((sub_bs, 3), device=device) * torch.tensor( @@ -379,10 +380,10 @@ class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods mask: Optional[Tensor] = None, 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 - - raise NotImplementedError() + if isinstance(mask, Tensor): + raise NotImplementedError( + "Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead." + ) else: s = (self.rays_to_sample, 1) ns = torch.randint(0, num_images, s, dtype=torch.long, device=device)