pair pixel sampler

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

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