diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index e115949d5..52fd0af0f 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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 ) diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 8c992c19f..1bb7130b9 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/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) diff --git a/nerfstudio/data/utils/pixel_sampling_utils.py b/nerfstudio/data/utils/pixel_sampling_utils.py new file mode 100644 index 000000000..f869dbaa2 --- /dev/null +++ b/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) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index ed976f92e..123b3d152 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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