tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

PixelSampler refactor (#1163)

Browse Source 
* moved collate methods as class methods, so sampling is only thing changed.

* disable pylint error, expected to be overrided.

* update.

* more pylint.

* final lint.

* cleanup, mask fix.
This commit is contained in:
Alexander Kristoffersen 2022-12-23 00:49:28 -08:00 committed by GitHub
parent 06642a5871
commit fd33a1e950
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 174 additions and 168 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

342
nerfstudio/data/pixel_samplers.py
View File

@ -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