Patch Sampling (#1458)

* patch sampling.

* changed str to literal.

* clean up sample selector.

nerfstudio/data/datamanagers/base_datamanager.py

@@ -51,7 +51,11 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import (
 )
 from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
 from nerfstudio.data.datasets.base_dataset import InputDataset
-from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler
+from nerfstudio.data.pixel_samplers import (
+    EquirectangularPixelSampler,
+    PatchPixelSampler,
+    PixelSampler,
+)
 from nerfstudio.data.utils.dataloaders import (
     CacheDataloader,
     FixedIndicesEvalDataloader,
@@ -283,6 +287,8 @@ class VanillaDataManagerConfig(InstantiateConfig):
     """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."""
 
 
 class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
@@ -346,6 +352,9 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self, dataset: InputDataset, *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)
+
         # If all images are equirectangular, use equirectangular pixel sampler
         is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value
         if is_equirectangular.all():

nerfstudio/data/pixel_samplers.py

@@ -243,3 +243,54 @@ class EquirectangularPixelSampler(PixelSampler):  # pylint: disable=too-few-publ
             ).long()
 
         return indices
+
+
+class PatchPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
+    """Samples 'pixel_batch's from 'image_batch's. Samples square patches
+    from the images randomly. Useful for patch-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
+        patch_size: side length of patch. This must be consistent in the method
+        config in order for samples to be reshaped into patches correctly.
+    """
+
+    def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None:
+        self.patch_size = kwargs["patch_size"]
+        num_rays = (num_rays_per_batch // (self.patch_size**2)) * (self.patch_size**2)
+        super().__init__(num_rays, 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[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).clone()
+
+            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