diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index fad3503a5..027b01352 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -171,7 +171,6 @@ method_configs["depth-nerfacto"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( _target=VanillaDataManager[DepthDataset], - patch_size=2, dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=4096, diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index aec500ddb..ff6933959 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -52,9 +52,9 @@ from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import ( - EquirectangularPixelSampler, - PairPixelSampler, PixelSampler, + PixelSamplerConfig, + PairPixelSamplerConfig, ) from nerfstudio.data.utils.dataloaders import ( CacheDataloader, @@ -338,8 +338,8 @@ class VanillaDataManagerConfig(DataManagerConfig): """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.""" + sampling_config: PixelSamplerConfig = PairPixelSamplerConfig() + """Specifies the pixel sampler used to sample pixels from images.""" TDataset = TypeVar("TDataset", bound=InputDataset, default=InputDataset) @@ -431,17 +431,10 @@ class VanillaDataManager(DataManager, Generic[TDataset]): def _get_pixel_sampler(self, dataset: TDataset, *args: Any, **kwargs: Any) -> PixelSampler: """Infer pixel sampler to use.""" - if self.config.patch_size > 1: - return PairPixelSampler(*args, **kwargs, radius=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(): - return EquirectangularPixelSampler(*args, **kwargs) - # Otherwise, use the default pixel sampler - if is_equirectangular.any(): - CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.") - return PixelSampler(*args, **kwargs) + is_equirectangular = (dataset.cameras.camera_type == CameraType.EQUIRECTANGULAR.value).all() + return self.config.sampling_config.setup( + is_equirectangular=is_equirectangular, num_rays_per_batch=self.config.train_num_rays_per_batch + ) def setup_train(self): """Sets up the data loaders for training""" diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 5157e05a2..780c4b44f 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,25 +17,52 @@ Code for sampling pixels. """ import random -from typing import Dict, Optional, Union import torch from jaxtyping import Int from torch import Tensor +from dataclasses import dataclass, field +from typing import ( + Dict, + Optional, + Type, + Union, +) + +from nerfstudio.configs.base_config import ( + InstantiateConfig, +) + + +@dataclass +class PixelSamplerConfig(InstantiateConfig): + """Configuration for pixel sampler instantiation.""" + + _target: Type = field(default_factory=lambda: PixelSampler) + """Target class to instantiate.""" + num_rays_per_batch: int = 4096 + """Number of rays to sample per batch.""" + keep_full_image: bool = False + """Whether or not to include a reference to the full image in returned batch.""" + is_equirectangular: bool = False + """List of whether or not camera i is equirectangular.""" + class PixelSampler: """Samples 'pixel_batch's from 'image_batch's. 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 + config: the DataManagerConfig used to instantiate class """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + config: PixelSamplerConfig + + def __init__(self, config: PixelSamplerConfig, **kwargs) -> None: self.kwargs = kwargs - self.num_rays_per_batch = num_rays_per_batch - self.keep_full_image = keep_full_image + self.config = config + self.num_rays_per_batch = self.config.num_rays_per_batch + self.keep_full_image = self.config.keep_full_image def set_num_rays_per_batch(self, num_rays_per_batch: int): """Set the number of rays to sample per batch. @@ -74,6 +101,36 @@ class PixelSampler: return indices + def sample_method_equirectangular( + self, + batch_size: int, + num_images: int, + image_height: int, + image_width: int, + mask: Optional[Tensor] = None, + device: Union[torch.device, str] = "cpu", + ) -> Int[Tensor, "batch_size 3"]: + if isinstance(mask, torch.Tensor): + # Note: if there is a mask, sampling reduces back to uniform sampling, which gives more + # sampling weight to the poles of the image than the equators. + # TODO(kevinddchen): implement the correct mask-sampling method. + + indices = self.sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) + else: + # We sample theta uniformly in [0, 2*pi] + # We sample phi in [0, pi] according to the PDF f(phi) = sin(phi) / 2. + # This is done by inverse transform sampling. + # http://corysimon.github.io/articles/uniformdistn-on-sphere/ + num_images_rand = torch.rand(batch_size, device=device) + phi_rand = torch.acos(1 - 2 * torch.rand(batch_size, device=device)) / torch.pi + theta_rand = torch.rand(batch_size, device=device) + indices = torch.floor( + torch.stack((num_images_rand, phi_rand, theta_rand), dim=-1) + * torch.tensor([num_images, image_height, image_width], device=device) + ).long() + + return indices + def collate_image_dataset_batch(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): """ Operates on a batch of images and samples pixels to use for generating rays. @@ -90,11 +147,21 @@ class PixelSampler: num_images, image_height, image_width, _ = batch["image"].shape if "mask" in batch: - indices = self.sample_method( - num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device - ) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device + ) + else: + indices = self.sample_method( + num_rays_per_batch, num_images, image_height, image_width, mask=batch["mask"], device=device + ) else: - indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_per_batch, num_images, image_height, image_width, device=device + ) + else: + indices = self.sample_method(num_rays_per_batch, num_images, image_height, image_width, device=device) c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) c, y, x = c.cpu(), y.cpu(), x.cpu() @@ -157,6 +224,12 @@ class PixelSampler: if i == num_images - 1: num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch indices = self.sample_method(num_rays_in_batch, 1, image_height, image_width, device=device) + if self.config.is_equirectangular: + indices = self.sample_method_equirectangular( + num_rays_in_batch, 1, image_height, image_width, device=device + ) + else: + indices = self.sample_method(num_rays_in_batch, 1, image_height, image_width, device=device) indices[:, 0] = i all_indices.append(indices) all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) @@ -203,45 +276,15 @@ class PixelSampler: return pixel_batch -class EquirectangularPixelSampler(PixelSampler): - """Samples 'pixel_batch's from 'image_batch's. Assumes images are - equirectangular and the sampling is done uniformly on the sphere. +@dataclass +class PatchPixelSamplerConfig(PixelSamplerConfig): + """Config dataclass for PatchPixelSampler.""" - 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 - """ - - # overrides base method - def sample_method( - self, - batch_size: int, - num_images: int, - image_height: int, - image_width: int, - mask: Optional[Tensor] = None, - device: Union[torch.device, str] = "cpu", - ) -> Int[Tensor, "batch_size 3"]: - if isinstance(mask, torch.Tensor): - # Note: if there is a mask, sampling reduces back to uniform sampling, which gives more - # sampling weight to the poles of the image than the equators. - # TODO(kevinddchen): implement the correct mask-sampling method. - - indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) - else: - # We sample theta uniformly in [0, 2*pi] - # We sample phi in [0, pi] according to the PDF f(phi) = sin(phi) / 2. - # This is done by inverse transform sampling. - # http://corysimon.github.io/articles/uniformdistn-on-sphere/ - num_images_rand = torch.rand(batch_size, device=device) - phi_rand = torch.acos(1 - 2 * torch.rand(batch_size, device=device)) / torch.pi - theta_rand = torch.rand(batch_size, device=device) - indices = torch.floor( - torch.stack((num_images_rand, phi_rand, theta_rand), dim=-1) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() - - return indices + _target: Type = field(default_factory=lambda: PatchPixelSampler) + """Target class to instantiate.""" + patch_size: int = 2 + """Side length of patch. This must be consistent in the method + config in order for samples to be reshaped into patches correctly.""" class PatchPixelSampler(PixelSampler): @@ -249,16 +292,16 @@ class PatchPixelSampler(PixelSampler): 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. + config: the PatchPixelSamplerConfig used to instantiate class """ - 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) + def __init__(self, config: PatchPixelSamplerConfig, **kwargs) -> None: + self.config = config + self.patch_size = self.config.patch_size + self.config.num_rays_per_batch = (self.config.num_rays_per_batch // (self.patch_size**2)) * ( + self.patch_size**2 + ) + super().__init__(self.config, **kwargs) def set_num_rays_per_batch(self, num_rays_per_batch: int): """Set the number of rays to sample per batch. Overridden to deal with patch-based sampling. @@ -302,20 +345,29 @@ class PatchPixelSampler(PixelSampler): return indices +@dataclass +class PairPixelSamplerConfig(PixelSamplerConfig): + """Config dataclass for PatchPixelSampler.""" + + _target: Type = field(default_factory=lambda: PairPixelSampler) + """Target class to instantiate.""" + radius: int = 2 + """max distance between pairs of pixels.""" + + class PairPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods """Samples pair of pixels from 'image_batch's. Samples pairs of pixels from from the images randomly within a 'radius' distance apart. Useful for pair-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 - radius: max distance between pairs of pixels + config: the PairPixelSamplerConfig used to instantiate class """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: - self.radius = kwargs["radius"] - self.rays_to_sample = num_rays_per_batch // 2 - super().__init__(num_rays_per_batch, keep_full_image, **kwargs) + def __init__(self, config: PairPixelSamplerConfig, **kwargs) -> None: + self.config = config + self.radius = self.config.radius + self.rays_to_sample = self.config.num_rays_per_batch // 2 + super().__init__(self.config, **kwargs) # overrides base method def sample_method( # pylint: disable=no-self-use