pair pixel sampling with a mask

nerfstudio/data/datamanagers/base_datamanager.py

@@ -461,7 +461,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
 
     def _get_pixel_sampler(self, dataset: TDataset, num_rays_per_batch: int) -> PixelSampler:
         """Infer pixel sampler to use."""
-        if self.config.patch_size > 1:
+        if self.config.patch_size > 1 and type(self.config.pixel_sampler) is PixelSamplerConfig:
             return PatchPixelSamplerConfig().setup(
                 patch_size=self.config.patch_size, num_rays_per_batch=num_rays_per_batch
             )

nerfstudio/data/pixel_samplers.py

@@ -23,7 +23,7 @@ from jaxtyping import Int
 from torch import Tensor
 
 from dataclasses import dataclass, field
-from infinitescenes.utils.mesh_utils import erode
+from nerfstudio.data.utils.pixel_sampling_utils import erode_mask
 from typing import (
     Dict,
     Optional,
@@ -319,8 +319,9 @@ class PatchPixelSampler(PixelSampler):
         device: Union[torch.device, str] = "cpu",
     ) -> Int[Tensor, "batch_size 3"]:
         if isinstance(mask, Tensor):
-            # Note: if there is a mask, should switch to the base PixelSampler class
-            raise NotImplementedError()
+            raise NotImplementedError(
+                "Masked sampling not implemented for PairPixelSampler. Change Config to PixelSamplerConfig instead."
+            )
         else:
             sub_bs = batch_size // (self.config.patch_size**2)
             indices = torch.rand((sub_bs, 3), device=device) * torch.tensor(
@@ -328,20 +329,20 @@ class PatchPixelSampler(PixelSampler):
                 device=device,
             )
 
-            indices = (
-                indices.view(sub_bs, 1, 1, 3)
-                .broadcast_to(sub_bs, self.config.patch_size, self.config.patch_size, 3)
-                .clone()
-            )
+        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
-            indices[:, ..., 2] += xxs
+        yys, xxs = torch.meshgrid(
+            torch.arange(self.config.patch_size, device=device), torch.arange(self.config.patch_size, device=device)
+        )
+        indices[:, ..., 1] += yys
+        indices[:, ..., 2] += xxs
 
-            indices = torch.floor(indices).long()
-            indices = indices.flatten(0, 2)
+        indices = torch.floor(indices).long()
+        indices = indices.flatten(0, 2)
 
         return indices
 
@@ -380,20 +381,12 @@ class PairPixelSampler(PixelSampler):  # pylint: disable=too-few-public-methods
         mask: Optional[Tensor] = None,
         device: Union[torch.device, str] = "cpu",
     ) -> Int[Tensor, "batch_size 3"]:
-<<<<<<< Updated upstream
-        if mask:
-            # Note: if there is a mask, should switch to the base PixelSampler class
-
-            raise NotImplementedError()
-=======
-        if mask is not None:
-            m = mask.permute(0,3,1,2).float()
-            for i in range(self.radius):
-                m = erode(m, kernel_size=3)
+        if isinstance(mask, Tensor):
+            m = mask.permute(0, 3, 1, 2).float()
+            m = erode_mask(m, 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]
->>>>>>> Stashed changes
         else:
             s = (self.rays_to_sample, 1)
             ns = torch.randint(0, num_images, s, dtype=torch.long, device=device)

nerfstudio/data/utils/pixel_sampling_utils.py

@@ -0,0 +1,66 @@
+# 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 - tensor
+    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)

nerfstudio/models/nerfacto.py

@@ -26,7 +26,6 @@ import torch
 from torch.nn import Parameter
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image import PeakSignalNoiseRatio
-from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
 from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation