Data#

Pixel Samplers#

Code for sampling pixels.

class nerfstudio.data.pixel_samplers.EquirectangularPixelSampler(num_rays_per_batch: int, keep_full_image: bool = False, **kwargs)[source]#

Bases: PixelSampler

Samples ‘pixel_batch’s from ‘image_batch’s. Assumes images are equirectangular and the sampling is done uniformly on the sphere.

Parameters:

  • 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

sample_method(batch_size: int, num_images: int, image_height: int, image_width: int, mask: Optional[Tensor] = None, device: Union[device, str] = 'cpu') → Int[Tensor, 'batch_size 3'][source]#

Naive pixel sampler, uniformly samples across all possible pixels of all possible images.

Parameters:

  • 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.

class nerfstudio.data.pixel_samplers.PatchPixelSampler(num_rays_per_batch: int, keep_full_image: bool = False, **kwargs)[source]#

Bases: PixelSampler

Samples ‘pixel_batch’s from ‘image_batch’s. Samples square patches from the images randomly. Useful for patch-based losses.

Parameters:

  • 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

  • correctly. (config in order for samples to be reshaped into patches) –

sample_method(batch_size: int, num_images: int, image_height: int, image_width: int, mask: Optional[Tensor] = None, device: Union[device, str] = 'cpu') → Int[Tensor, 'batch_size 3'][source]#

Naive pixel sampler, uniformly samples across all possible pixels of all possible images.

Parameters:

  • 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.

set_num_rays_per_batch(num_rays_per_batch: int)[source]#

Set the number of rays to sample per batch. Overridden to deal with patch-based sampling.

Parameters:

num_rays_per_batch – number of rays to sample per batch

class nerfstudio.data.pixel_samplers.PixelSampler(num_rays_per_batch: int, keep_full_image: bool = False, **kwargs)[source]#

Bases: object

Samples ‘pixel_batch’s from ‘image_batch’s.

Parameters:

  • 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

collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False)[source]#

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.

Parameters:

  • 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

collate_image_dataset_batch_list(batch: Dict, num_rays_per_batch: int, keep_full_image: bool = False)[source]#

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.

Parameters:

  • 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

sample(image_batch: Dict)[source]#

Sample an image batch and return a pixel batch.

Parameters:

image_batch – batch of images to sample from

sample_method(batch_size: int, num_images: int, image_height: int, image_width: int, mask: Optional[Tensor] = None, device: Union[device, str] = 'cpu') → Int[Tensor, 'batch_size 3'][source]#

Naive pixel sampler, uniformly samples across all possible pixels of all possible images.

Parameters:

  • 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.

set_num_rays_per_batch(num_rays_per_batch: int)[source]#

Set the number of rays to sample per batch.

Parameters:

num_rays_per_batch – number of rays to sample per batch

Scene Box#

Dataset input structures.

class nerfstudio.data.scene_box.SceneBox(aabb: Float[Tensor, '2 3'])[source]#

Bases: object

Data to represent the scene box.

aabb: Float[Tensor, '2 3']#

axis-aligned bounding box. aabb[0] is the minimum (x,y,z) point. aabb[1] is the maximum (x,y,z) point.

Type:

aabb

static from_camera_poses(poses: Float[Tensor, '*batch 3 4'], scale_factor: float) → SceneBox[source]#

Returns the instance of SceneBox that fully envelopes a set of poses

Parameters:

  • poses – tensor of camera pose matrices

  • scale_factor – How much to scale the camera origins by.

get_center()[source]#

Returns the center of the box.

get_centered_and_scaled_scene_box(scale_factor: Union[float, Tensor] = 1.0)[source]#

Returns a new box that has been shifted and rescaled to be centered about the origin.

Parameters:

scale_factor – How much to scale the camera origins by.

get_diagonal_length()[source]#

Returns the longest diagonal length.

static get_normalized_positions(positions: Float[Tensor, '*batch 3'], aabb: Float[Tensor, '2 3'])[source]#

Return normalized positions in range [0, 1] based on the aabb axis-aligned bounding box.

Parameters:

  • positions – the xyz positions

  • aabb – the axis-aligned bounding box