Adding support for PairPixelSampler and PatchPixelSampler masks (#2305)

* changes

* pair pixel sampling with a mask

* reverting changes

* revert changes

* remove line in erode

* support for patch pixel samplerr masks

* fix

---------

Co-authored-by: machenmusik <machenmusik@users.noreply.github.com>

nerfstudio/data/pixel_samplers.py

@@ -23,6 +23,7 @@ from jaxtyping import Int
 from torch import Tensor
 
 from dataclasses import dataclass, field
+from nerfstudio.data.utils.pixel_sampling_utils import erode_mask
 from typing import (
     Dict,
     Optional,
@@ -284,7 +285,7 @@ class PatchPixelSamplerConfig(PixelSamplerConfig):
 
     _target: Type = field(default_factory=lambda: PatchPixelSampler)
     """Target class to instantiate."""
-    patch_size: int = 2
+    patch_size: int = 32
     """Side length of patch. This must be consistent in the method
     config in order for samples to be reshaped into patches correctly."""
 
@@ -318,9 +319,27 @@ class PatchPixelSampler(PixelSampler):
         device: Union[torch.device, str] = "cpu",
     ) -> Int[Tensor, "batch_size 3"]:
         if isinstance(mask, Tensor):
-            raise NotImplementedError(
-                "Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead."
+            sub_bs = batch_size // (self.config.patch_size**2)
+            half_patch_size = int(self.config.patch_size / 2)
+            m = erode_mask(mask.permute(0, 3, 1, 2).float(), pixel_radius=half_patch_size)
+            nonzero_indices = torch.nonzero(m[:, 0], as_tuple=False).to(device)
+            chosen_indices = random.sample(range(len(nonzero_indices)), k=sub_bs)
+            indices = nonzero_indices[chosen_indices]
+
+            indices = (
+                indices.view(sub_bs, 1, 1, 3)
+                .broadcast_to(sub_bs, self.config.patch_size, self.config.patch_size, 3)
+                .clone()
             )
+
+            yys, xxs = torch.meshgrid(
+                torch.arange(self.config.patch_size, device=device), torch.arange(self.config.patch_size, device=device)
+            )
+            indices[:, ..., 1] += yys - half_patch_size
+            indices[:, ..., 2] += xxs - half_patch_size
+
+            indices = torch.floor(indices).long()
+            indices = indices.flatten(0, 2)
         else:
             sub_bs = batch_size // (self.config.patch_size**2)
             indices = torch.rand((sub_bs, 3), device=device) * torch.tensor(
@@ -381,9 +400,10 @@ class PairPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
         device: Union[torch.device, str] = "cpu",
     ) -> Int[Tensor, "batch_size 3"]:
         if isinstance(mask, Tensor):
-            raise NotImplementedError(
-                "Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead."
-            )
+            m = erode_mask(mask.permute(0, 3, 1, 2).float(), pixel_radius=self.radius)
+            nonzero_indices = torch.nonzero(m[:, 0], as_tuple=False).to(device)
+            chosen_indices = random.sample(range(len(nonzero_indices)), k=self.rays_to_sample)
+            indices = nonzero_indices[chosen_indices]
         else:
             s = (self.rays_to_sample, 1)
             ns = torch.randint(0, num_images, s, dtype=torch.long, device=device)
@@ -391,12 +411,12 @@ class PairPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
             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 = 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)
+        )
+        pair_indices += indices
+        indices = torch.hstack((indices, pair_indices)).view(self.rays_to_sample * 2, 3)
         return indices

nerfstudio/data/utils/pixel_sampling_utils.py

@@ -0,0 +1,65 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" Pixel sampling utils such as eroding of valid masks that we sample from. """
+
+import torch
+from torch import Tensor
+from jaxtyping import Float
+
+
+def dilate(tensor: Float[Tensor, "bs 1 H W"], kernel_size=3) -> Float[Tensor, "bs 1 H W"]:
+    """Dilate a tensor with 0s and 1s. 0s will be be expanded based on the kernel size.
+
+    Args:
+        kernel_size: Size of the pooling region. Dilates/contracts 1 pixel if kernel_size is 3.
+    """
+
+    unique_vals = torch.unique(tensor)
+    if any(val not in (0, 1) for val in unique_vals) or tensor.dtype != torch.float32:
+        raise ValueError("Input tensor should contain only values 0 and 1, and should have dtype torch.float32.")
+
+    return torch.nn.functional.max_pool2d(tensor, kernel_size=kernel_size, stride=1, padding=(kernel_size - 1) // 2)
+
+
+def erode(tensor: Float[Tensor, "bs 1 H W"], kernel_size=3) -> Float[Tensor, "bs 1 H W"]:
+    """Erode a tensor with 0s and 1s. 1s will be expanded based on the kernel size.
+
+    Args:
+        kernel_size: Size of the pooling region. Erodes/expands 1 pixel if kernel_size is 3.
+    """
+
+    unique_vals = torch.unique(tensor)
+    if any(val not in (0, 1) for val in unique_vals) or tensor.dtype != torch.float32:
+        raise ValueError("Input tensor should contain only values 0 and 1, and should have dtype torch.float32.")
+
+    x = 1 - dilate(1 - tensor, kernel_size=kernel_size)
+    # set edge pixels to 0
+    p = (kernel_size - 1) // 2
+    x[:, :, :p, :] *= 0
+    x[:, :, :, :p] *= 0
+    x[:, :, -p:, :] *= 0
+    x[:, :, :, -p:] *= 0
+    return x
+
+
+def erode_mask(tensor: Float[Tensor, "bs 1 H W"], pixel_radius: int = 1):
+    """Erode a mask. Expands 1 values to nearby pixels with a max pooling operation.
+    A pixel radius of 1 will also extend the 1s along the diagonal.
+
+    Args:
+        pixel_radius: The number of pixels away from valid pixels (1s) that we may sample.
+    """
+    kernel_size = 1 + 2 * pixel_radius
+    return erode(tensor, kernel_size=kernel_size)