diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 07184b47c..8f0878568 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -17,120 +17,10 @@ Code for sampling pixels. """ import random -from typing import Dict +from typing import Dict, Optional, Union import torch - - -def collate_image_dataset_batch(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. - Returns a collated batch which is input to the Graph. - It will sample only within the valid 'mask' if it's specified. - - Args: - batch: batch of images to sample from - 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 - """ - device = batch["image"].device - num_images, image_height, image_width, _ = batch["image"].shape - - # only sample within the mask, if the mask is in the batch - if "mask" in batch: - nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) - chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) - indices = nonzero_indices[chosen_indices] - else: - indices = torch.floor( - torch.rand((num_rays_per_batch, 3), device=device) - * torch.tensor([num_images, image_height, image_width], device=device) - ).long() - - c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None} - - assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape - - # Needed to correct the random indices to their actual camera idx locations. - indices[:, 0] = batch["image_idx"][c] - collated_batch["indices"] = indices # with the abs camera indices - - if keep_full_image: - collated_batch["full_image"] = batch["image"] - - return collated_batch - - -def collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): - """ - Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside - a list. - - We will use this with the intent of DEPRECIATING it as soon as we find a viable alternative. - The intention will be to replace this with a more efficient implementation that doesn't require a for loop, but - since pytorch's ragged tensors are still in beta (this would allow for some vectorization), this will do - - Args: - batch: batch of images to sample from - 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 - """ - - device = batch["image"][0].device - num_images = len(batch["image"]) - - # only sample within the mask, if the mask is in the batch - all_indices = [] - all_images = [] - - if "mask" in batch: - num_rays_in_batch = num_rays_per_batch // num_images - for i in range(num_images): - if i == num_images - 1: - num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch - nonzero_indices = torch.nonzero(batch["mask"][i][..., 0], as_tuple=False) - chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_in_batch) - indices = nonzero_indices[chosen_indices] - indices = torch.cat([torch.full((num_rays_in_batch, 1), i, device=device), indices], dim=-1) - all_indices.append(indices) - all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) - - else: - num_rays_in_batch = num_rays_per_batch // num_images - for i in range(num_images): - image_height, image_width, _ = batch["image"][i].shape - if i == num_images - 1: - num_rays_in_batch = num_rays_per_batch - (num_images - 1) * num_rays_in_batch - indices = torch.floor( - torch.rand((num_rays_in_batch, 3), device=device) - * torch.tensor([1, image_height, image_width], device=device) - ).long() - indices[:, 0] = i - all_indices.append(indices) - all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) - - indices = torch.cat(all_indices, dim=0) - - c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - collated_batch = { - key: value[c, y, x] - for key, value in batch.items() - if key != "image_idx" and key != "image" and key != "mask" and value is not None - } - - collated_batch["image"] = torch.cat(all_images, dim=0) - - assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape - - # Needed to correct the random indices to their actual camera idx locations. - indices[:, 0] = batch["image_idx"][c] - collated_batch["indices"] = indices # with the abs camera indices - - if keep_full_image: - collated_batch["full_image"] = batch["image"] - - return collated_batch +from torchtyping import TensorType class PixelSampler: # pylint: disable=too-few-public-methods @@ -141,7 +31,8 @@ class PixelSampler: # pylint: disable=too-few-public-methods keep_full_image: whether or not to include a reference to the full image in returned batch """ - def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False) -> None: + def __init__(self, num_rays_per_batch: int, keep_full_image: bool = False, **kwargs) -> None: + self.kwargs = kwargs self.num_rays_per_batch = num_rays_per_batch self.keep_full_image = keep_full_image @@ -153,6 +44,144 @@ class PixelSampler: # pylint: disable=too-few-public-methods """ self.num_rays_per_batch = num_rays_per_batch + 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]: + """ + Naive pixel sampler, uniformly samples across all possible pixels of all possible images. + + Args: + batch_size: number of samples in a batch + num_images: number of images to sample over + mask: mask of possible pixels in an image to sample from. + """ + if mask: + nonzero_indices = torch.nonzero(mask[..., 0], as_tuple=False) + chosen_indices = random.sample(range(len(nonzero_indices)), k=batch_size) + indices = nonzero_indices[chosen_indices] + else: + indices = torch.floor( + torch.rand((batch_size, 3), device=device) + * 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. + Returns a collated batch which is input to the Graph. + It will sample only within the valid 'mask' if it's specified. + + Args: + batch: batch of images to sample from + 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 + """ + + device = batch["image"].device + 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 + ) + 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)) + collated_batch = { + key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None + } + + assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape + + # Needed to correct the random indices to their actual camera idx locations. + indices[:, 0] = batch["image_idx"][c] + collated_batch["indices"] = indices # with the abs camera indices + + if keep_full_image: + collated_batch["full_image"] = batch["image"] + + return collated_batch + + def collate_image_dataset_batch_list(self, batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): + """ + Does the same as collate_image_dataset_batch, except it will operate over a list of images / masks inside + a list. + + We will use this with the intent of DEPRECIATING it as soon as we find a viable alternative. + The intention will be to replace this with a more efficient implementation that doesn't require a for loop, but + since pytorch's ragged tensors are still in beta (this would allow for some vectorization), this will do. + + Args: + batch: batch of images to sample from + 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 + """ + + device = batch["image"][0].device + num_images = len(batch["image"]) + + # only sample within the mask, if the mask is in the batch + all_indices = [] + all_images = [] + + if "mask" in batch: + num_rays_in_batch = num_rays_per_batch // num_images + for i in range(num_images): + + image_height, image_width, _ = batch["image"][i].shape + + 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, mask=batch["mask"][i], device=device + ) + indices[:, 0] = i + all_indices.append(indices) + all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) + + else: + num_rays_in_batch = num_rays_per_batch // num_images + for i in range(num_images): + image_height, image_width, _ = batch["image"][i].shape + 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) + indices[:, 0] = i + all_indices.append(indices) + all_images.append(batch["image"][i][indices[:, 1], indices[:, 2]]) + + indices = torch.cat(all_indices, dim=0) + + c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) + collated_batch = { + key: value[c, y, x] + for key, value in batch.items() + if key != "image_idx" and key != "image" and key != "mask" and value is not None + } + + collated_batch["image"] = torch.cat(all_images, dim=0) + + assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape + + # Needed to correct the random indices to their actual camera idx locations. + indices[:, 0] = batch["image_idx"][c] + collated_batch["indices"] = indices # with the abs camera indices + + if keep_full_image: + collated_batch["full_image"] = batch["image"] + + return collated_batch + def sample(self, image_batch: Dict): """Sample an image batch and return a pixel batch. @@ -161,11 +190,11 @@ class PixelSampler: # pylint: disable=too-few-public-methods """ if isinstance(image_batch["image"], list): image_batch = dict(image_batch.items()) # copy the dictionary so we don't modify the original - pixel_batch = collate_image_dataset_batch_list( + pixel_batch = self.collate_image_dataset_batch_list( image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image ) elif isinstance(image_batch["image"], torch.Tensor): - pixel_batch = collate_image_dataset_batch( + pixel_batch = self.collate_image_dataset_batch( image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image ) else: @@ -173,54 +202,6 @@ class PixelSampler: # pylint: disable=too-few-public-methods return pixel_batch -def collate_image_dataset_batch_equirectangular(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False): - """ - Operates on a batch of equirectangular images and samples pixels to use for - generating rays. Rays will be generated uniformly on the sphere. - Returns a collated batch which is input to the Graph. - It will sample only within the valid 'mask' if it's specified. - - Args: - batch: batch of images to sample from - 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 - """ - # TODO(kevinddchen): make more DRY - device = batch["image"].device - num_images, image_height, image_width, _ = batch["image"].shape - - # only sample within the mask, if the mask is in the batch - if "mask" in batch: - # TODO(kevinddchen): implement this - raise NotImplementedError("Masking not implemented for equirectangular images.") - - # 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(num_rays_per_batch, device=device) - phi_rand = torch.acos(1 - 2 * torch.rand(num_rays_per_batch, device=device)) / torch.pi - theta_rand = torch.rand(num_rays_per_batch, 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() - - c, y, x = (i.flatten() for i in torch.split(indices, 1, dim=-1)) - collated_batch = {key: value[c, y, x] for key, value in batch.items() if key != "image_idx" and value is not None} - - assert collated_batch["image"].shape == (num_rays_per_batch, 3), collated_batch["image"].shape - - # Needed to correct the random indices to their actual camera idx locations. - indices[:, 0] = batch["image_idx"][c] - collated_batch["indices"] = indices # with the abs camera indices - - if keep_full_image: - collated_batch["full_image"] = batch["image"] - - return collated_batch - - class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods """Samples 'pixel_batch's from 'image_batch's. Assumes images are equirectangular and the sampling is done uniformly on the sphere. @@ -231,9 +212,34 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ """ # overrides base method - def sample(self, image_batch: Dict): + 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]: - pixel_batch = collate_image_dataset_batch_equirectangular( - image_batch, self.num_rays_per_batch, keep_full_image=self.keep_full_image - ) - return pixel_batch + if mask: + # 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