diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 41d457458..d58f7d5a5 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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 diff --git a/nerfstudio/data/utils/pixel_sampling_utils.py b/nerfstudio/data/utils/pixel_sampling_utils.py new file mode 100644 index 000000000..464381f4c --- /dev/null +++ b/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)