diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 1c8c3e08d..e48218c39 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -51,7 +51,11 @@ from nerfstudio.data.dataparsers.phototourism_dataparser import ( ) from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset -from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler +from nerfstudio.data.pixel_samplers import ( + EquirectangularPixelSampler, + PatchPixelSampler, + PixelSampler, +) from nerfstudio.data.utils.dataloaders import ( CacheDataloader, FixedIndicesEvalDataloader, @@ -283,6 +287,8 @@ class VanillaDataManagerConfig(InstantiateConfig): """The scale factor for scaling spatial data such as images, mask, semantics along with relevant information about camera intrinsics """ + patch_size: int = 1 + """Size of patch to sample from. If >1, patch-based sampling will be used.""" class VanillaDataManager(DataManager): # pylint: disable=abstract-method @@ -346,6 +352,9 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self, dataset: InputDataset, *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) + # If all images are equirectangular, use equirectangular pixel sampler is_equirectangular = dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value if is_equirectangular.all(): diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 91fc6301d..8ac8e183d 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -243,3 +243,54 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ ).long() return indices + + +class PatchPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods + """Samples 'pixel_batch's from 'image_batch's. Samples square patches + from the images randomly. Useful for patch-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 + patch_size: side length of patch. This must be consistent in the method + config in order for samples to be reshaped into patches correctly. + """ + + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + self.patch_size = kwargs["patch_size"] + num_rays = (num_rays_per_batch // (self.patch_size**2)) * (self.patch_size**2) + super().__init__(num_rays, 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[TensorType] = None, + device: Union[torch.device, str] = "cpu", + ) -> TensorType["batch_size", 3]: + + if mask: + # Note: if there is a mask, sampling reduces back to uniform sampling + indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + sub_bs = batch_size // (self.patch_size**2) + indices = torch.rand((sub_bs, 3), device=device) * torch.tensor( + [num_images, image_height - self.patch_size, image_width - self.patch_size], + device=device, + ) + + indices = indices.view(sub_bs, 1, 1, 3).broadcast_to(sub_bs, self.patch_size, self.patch_size, 3).clone() + + yys, xxs = torch.meshgrid( + torch.arange(self.patch_size, device=device), torch.arange(self.patch_size, device=device) + ) + indices[:, ..., 1] += yys + indices[:, ..., 2] += xxs + + indices = torch.floor(indices).long() + indices = indices.flatten(0, 2) + + return indices