diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 82e4f66d3..52764c468 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -54,6 +54,7 @@ from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import ( EquirectangularPixelSampler, PatchPixelSampler, + PairPixelSampler, PixelSampler, ) from nerfstudio.data.utils.dataloaders import ( @@ -432,7 +433,8 @@ class VanillaDataManager(DataManager, Generic[TDataset]): def _get_pixel_sampler(self, dataset: TDataset, *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) + return PairPixelSampler(*args, **kwargs, radius=self.config.patch_size) + # If all images are equirectangular, use equirectangular pixel sampler is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index a145dce46..15deb446c 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -300,3 +300,47 @@ class PatchPixelSampler(PixelSampler): indices = indices.flatten(0, 2) return indices + +class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods + """Samples pair of pixels from 'image_batch's. Samples pairs of pixels from + from the images randomly within a 'radius' distance apart. Useful for pair-based losses. + Args: + num_rays_per_batch: number of rays to sample per batch + keep_full_image: whether or not to include a reference to the full image in returned batch + radius: max distance between pairs of pixels + """ + + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + self.radius = kwargs["radius"] + self.rays_to_sample = num_rays_per_batch // 2 + super().__init__(num_rays_per_batch, keep_full_image, **kwargs) + + # overrides base method + def sample_method( # pylint: disable=no-self-use + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[Tensor] = None, + device: Union[torch.device, str] = "cpu", + ) -> Int[Tensor, "batch_size 3"]: + if mask: + # Note: if there is a mask, sampling reduces back to uniform sampling + raise NotImplementedError() + else: + s = (self.rays_to_sample, 1) + ns = torch.randint(0, num_images, s, dtype=torch.long, device=device) + hs = torch.randint(self.radius, image_height - self.radius, s, dtype=torch.long, device=device) + ws = torch.randint(self.radius, image_width - self.radius, s, dtype=torch.long, device=device) + indices = torch.concat((ns, hs, ws), dim=1) + + pair_indices = torch.hstack( + ( + torch.zeros(self.rays_to_sample, 1, device=device, dtype=torch.long), + torch.randint(-self.radius, self.radius, (self.rays_to_sample, 2), device=device, dtype=torch.long), + ) + ) + pair_indices += indices + indices = torch.hstack((indices, pair_indices)).view(self.rays_to_sample * 2, 3) + return indices