pytorch3d.implicitron.models.implicit_function.voxel_grid_implicit_function
voxel_grid_implicit_function
- class pytorch3d.implicitron.models.implicit_function.voxel_grid_implicit_function.VoxelGridImplicitFunction(*args, **kwargs)[source]
Bases:
ImplicitFunctionBase
,Module
This implicit function consists of two streams, one for the density calculation and one for the color calculation. Each of these streams has three main parts:
- Voxel grids:
They take the (x, y, z) position and return the embedding of that point. These components are replaceable, you can make your own or choose one of several options.
- Harmonic embeddings:
Convert each feature into series of ‘harmonic features’, feature is passed through sine and cosine functions. Input is of shape [minibatch, …, D] output [minibatch, …, (n_harmonic_functions * 2 + int(append_input)) * D]. Appends input by default. If you want it to behave like identity, put n_harmonic_functions=0 and append_input=True.
- Decoding functions:
The decoder is an instance of the DecoderFunctionBase and converts the embedding of a spatial location to density/color. Examples are Identity which returns its input and the MLP which uses fully connected nerual network to transform the input. These components are replaceable, you can make your own or choose from several options.
- Calculating density is done in three steps:
Evaluating the voxel grid on points
Embedding the outputs with harmonic embedding
Passing through the Density decoder
- To calculate the color we need the embedding and the viewing direction, it has five steps:
Transforming the viewing direction with camera
Evaluating the voxel grid on points
Embedding the outputs with harmonic embedding
Embedding the normalized direction with harmonic embedding
Passing everything through the Color decoder
If using the Implicitron configuration system the input_dim to the decoding functions will be set to the output_dim of the Harmonic embeddings.
A speed up comes from using the scaffold, a low resolution voxel grid. The scaffold is referenced as “binary occupancy grid mask” in TensoRF paper and “AlphaMask” in official TensoRF implementation. The scaffold is used in:
- filtering points in empty space
- controlled by scaffold_filter_points boolean. If set to True, points for which
scaffold predicts that are in empty space will return 0 density and (0, 0, 0) color.
- calculating the bounding box of an object and cropping the voxel grids
controlled by volume_cropping_epochs.
- at those epochs the implicit function will find the bounding box of an object
inside it and crop density and color grids. Cropping of the voxel grids means preserving only voxel values that are inside the bounding box and changing the resolution to match the original, while preserving the new cropped location in world coordinates.
The scaffold has to exist before attempting filtering and cropping, and is created on scaffold_calculating_epochs. Each voxel in the scaffold is labeled as having density 1 if the point in the center of it evaluates to greater than scaffold_empty_space_threshold. 3D max pooling is performed on the densities of the points in 3D. Scaffold features are off by default.
- Members:
voxel_grid_density (VoxelGridBase): voxel grid to use for density estimation voxel_grid_color (VoxelGridBase): voxel grid to use for color estimation
- harmonic_embedder_xyz_density (HarmonicEmbedder): Function to transform the outputs of
the voxel_grid_density
- harmonic_embedder_xyz_color (HarmonicEmbedder): Function to transform the outputs of
the voxel_grid_color for density
- harmonic_embedder_dir_color (HarmonicEmbedder): Function to transform the outputs of
the voxel_grid_color for color
decoder_density (DecoderFunctionBase): decoder function to use for density estimation color_density (DecoderFunctionBase): decoder function to use for color estimation
- use_multiple_streams (bool): if you want the density and color calculations to run on
different cuda streams set this to True. Default True.
- xyz_ray_dir_in_camera_coords (bool): This is true if the directions are given in
camera coordinates. Default False.
- voxel_grid_scaffold (VoxelGridModule): which holds the scaffold. Extents and
translation of it are set to those of voxel_grid_density.
- scaffold_calculating_epochs (Tuple[int, …]): at which epochs to recalculate the
scaffold. (The scaffold will be created automatically at the beginning of the calculation.)
- scaffold_resolution (Tuple[int, int, int]): (width, height, depth) of the underlying
voxel grid which stores scaffold
- scaffold_empty_space_threshold (float): if self._get_density evaluates to less than
this it will be considered as empty space and the scaffold at that point would evaluate as empty space.
- scaffold_occupancy_chunk_size (str or int): Number of xy scaffold planes to calculate
at the same time. To calculate the scaffold we need to query _get_density() at every voxel, this calculation can be split into scaffold depth number of xy plane calculations if you want the lowest memory usage, one calculation to calculate the whole scaffold, but with higher memory footprint or any other number of planes. Setting to a non-positive number calculates all planes at the same time. Defaults to ‘-1’ (=calculating all planes).
- scaffold_max_pool_kernel_size (int): Size of the pooling region to use when
calculating the scaffold. Defaults to 3.
- scaffold_filter_points (bool): If set to True the points will be filtered using
self.voxel_grid_scaffold. Filtered points will be predicted as having 0 density and (0, 0, 0) color. The points which were not evaluated as empty space will be passed through the steps outlined above.
- volume_cropping_epochs: on which epochs to crop the voxel grids to fit the object’s
bounding box. Scaffold has to be calculated before cropping.
- voxel_grid_density: VoxelGridModule
- voxel_grid_color: VoxelGridModule
- harmonic_embedder_xyz_density_args: DictConfig = Field(name=None,type=None,default=<dataclasses._MISSING_TYPE object>,default_factory=<function get_default_args_field.<locals>.create>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=<dataclasses._MISSING_TYPE object>,_field_type=None)
- harmonic_embedder_xyz_color_args: DictConfig = Field(name=None,type=None,default=<dataclasses._MISSING_TYPE object>,default_factory=<function get_default_args_field.<locals>.create>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=<dataclasses._MISSING_TYPE object>,_field_type=None)
- harmonic_embedder_dir_color_args: DictConfig = Field(name=None,type=None,default=<dataclasses._MISSING_TYPE object>,default_factory=<function get_default_args_field.<locals>.create>,init=True,repr=True,hash=None,compare=True,metadata=mappingproxy({}),kw_only=<dataclasses._MISSING_TYPE object>,_field_type=None)
- decoder_density_class_type: str = 'MLPDecoder'
- decoder_density: DecoderFunctionBase
- decoder_color_class_type: str = 'MLPDecoder'
- decoder_color: DecoderFunctionBase
- use_multiple_streams: bool = True
- xyz_ray_dir_in_camera_coords: bool = False
- scaffold_calculating_epochs: Tuple[int, ...] = ()
- scaffold_resolution: Tuple[int, int, int] = (128, 128, 128)
- scaffold_empty_space_threshold: float = 0.001
- scaffold_occupancy_chunk_size: int = -1
- scaffold_max_pool_kernel_size: int = 3
- scaffold_filter_points: bool = True
- volume_cropping_epochs: Tuple[int, ...] = ()
- forward(ray_bundle: ImplicitronRayBundle, fun_viewpool=None, camera: CamerasBase | None = None, global_code=None, **kwargs) Tuple[Tensor, Tensor, Dict] [source]
The forward function accepts the parametrizations of 3D points sampled along projection rays. The forward pass is responsible for attaching a 3D vector and a 1D scalar representing the point’s RGB color and opacity respectively.
- Parameters:
ray_bundle –
An ImplicitronRayBundle object containing the following variables: origins: A tensor of shape (minibatch, …, 3) denoting the
origins of the sampling rays in world coords.
- directions: A tensor of shape (minibatch, …, 3)
containing the direction vectors of sampling rays in world coords.
- lengths: A tensor of shape (minibatch, …, num_points_per_ray)
containing the lengths at which the rays are sampled.
fun_viewpool –
- an optional callback with the signature
fun_fiewpool(points) -> pooled_features
where points is a [N_TGT x N x 3] tensor of world coords, and pooled_features is a [N_TGT x … x N_SRC x latent_dim] tensor of the features pooled from the context images.
camera – A camera model which will be used to transform the viewing directions
- Returns:
rays_densities –
- A tensor of shape (minibatch, …, num_points_per_ray, 1)
denoting the opacitiy of each ray point.
- rays_colors: A tensor of shape (minibatch, …, num_points_per_ray, 3)
denoting the color of each ray point.
- static allows_multiple_passes() bool [source]
Returns True as this implicit function allows multiple passes. Overridden from ImplicitFunctionBase.
- subscribe_to_epochs() Tuple[Tuple[int, ...], Callable[[int], bool]] [source]
Method which expresses interest in subscribing to optimization epoch updates. This implicit function subscribes to epochs to calculate the scaffold and to crop voxel grids, so this method combines wanted epochs and wraps their callbacks.
- Returns:
- list of epochs on which to call a callable and callable to be called on
particular epoch. The callable returns True if parameter change has happened else False and it must be supplied with one argument, epoch.
- create_decoder_density_impl(type_, args: DictConfig) None [source]
Decoding functions come after harmonic embedding and voxel grid. In order to not calculate the input dimension of the decoder in the config file this function calculates the required input dimension and sets the input dimension of the decoding function to this value.
- create_decoder_color_impl(type_, args: DictConfig) None [source]
Decoding functions come after harmonic embedding and voxel grid. In order to not calculate the input dimension of the decoder in the config file this function calculates the required input dimension and sets the input dimension of the decoding function to this value.