From a15cf5b8ad18b4b825be9c02bf8048e924e73a30 Mon Sep 17 00:00:00 2001 From: akristoffersen Date: Tue, 27 Dec 2022 06:51:21 +0000 Subject: [PATCH] added pixel patch sampler. --- nerfstudio/data/pixel_samplers.py | 42 +++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 8f0878568..aa6469380 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -243,3 +243,45 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ ).long() return indices + + +class PatchPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods + """todo""" + + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + super().__init__(num_rays_per_batch, keep_full_image, **kwargs) + self.patch_size = kwargs["patch_size"] + + # 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[TensorType] = None, + device: Union[torch.device, str] = "cpu", + ) -> TensorType["batch_size", 3]: + + if mask: + # Note: if there is a mask, sampling reduces back to uniform sampling + indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + sub_bs = batch_size // (self.patch_size**2) + indices = torch.rand((sub_bs, 3), device=device) * torch.tensor( + [num_images, image_height - self.patch_size, image_width - self.patch_size], + device=device, + ) + + indices = indices.view(sub_bs, 1, 1, 3).broadcast_to(sub_bs, self.patch_size, self.patch_size, 3) + + yys, xxs = torch.meshgrid( + torch.arange(self.patch_size, device=device), torch.arange(self.patch_size, device=device) + ) + indices[:, ..., 1] += yys + indices[:, ..., 2] += xxs + + indices = torch.floor(indices).long() + indices = indices.flatten(0, 2) + + return indices