4
0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00

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 <tancik@berkeley.edu>
Co-authored-by: Ethan Weber <ethanweber@berkeley.edu>
Co-authored-by: Ethan Weber <eweb0124@gmail.com>
This commit is contained in:
AdamRashid96 2023-08-04 07:27:35 -07:00 committed by GitHub
parent 9b03299d40
commit 7e90cb7989
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 7 deletions

View File

@ -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
)

View File

@ -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)