Estimators¶

class nerfacc.OccGridEstimator(roi_aabb, resolution=128, levels=1, **kwargs)¶

Occupancy grid transmittance estimator for spatial skipping.

References: “Instant Neural Graphics Primitives.”

Parameters:

roi_aabb (Union[List[int], Tensor]) – The axis-aligned bounding box of the region of interest. Useful for mapping the 3D space to the grid.
resolution (Union[int, List[int], Tensor]) – The resolution of the grid. If an integer is given, the grid is assumed to be a cube. Otherwise, a list or a tensor of shape (3,) is expected. Default: 128.
levels (int) – The number of levels of the grid. Default: 1.

mark_invisible_cells(K, c2w, width, height, near_plane=0.0, chunk=32768)¶

Mark the cells that aren’t covered by the cameras with density -1. Should only be executed once before training starts.

Parameters:

K (Tensor) – Camera intrinsics of shape (N, 3, 3) or (1, 3, 3).
c2w (Tensor) – Camera to world poses of shape (N, 3, 4) or (N, 4, 4).
width (int) – Image width in pixels
height (int) – Image height in pixels
near_plane (float) – Near plane distance
chunk (int) – The chunk size to split the cells (to avoid OOM)

Return type:

None

sampling(rays_o, rays_d, sigma_fn=None, alpha_fn=None, near_plane=0.0, far_plane=10000000000.0, t_min=None, t_max=None, render_step_size=0.001, early_stop_eps=0.0001, alpha_thre=0.0, stratified=False, cone_angle=0.0)¶

Sampling with spatial skipping.

Note

This function is not differentiable to any inputs.

Parameters:

rays_o (Tensor) – Ray origins of shape (n_rays, 3).
rays_d (Tensor) – Normalized ray directions of shape (n_rays, 3).
sigma_fn (Optional[Callable]) – Optional. If provided, the marching will skip the invisible space by evaluating the density along the ray with sigma_fn. It should be a function that takes in samples {t_starts (N,), t_ends (N,), ray indices (N,)} and returns the post-activation density values (N,). You should only provide either sigma_fn or alpha_fn.
alpha_fn (Optional[Callable]) – Optional. If provided, the marching will skip the invisible space by evaluating the density along the ray with alpha_fn. It should be a function that takes in samples {t_starts (N,), t_ends (N,), ray indices (N,)} and returns the post-activation opacity values (N,). You should only provide either sigma_fn or alpha_fn.
near_plane (float) – Optional. Near plane distance. Default: 0.0.
far_plane (float) – Optional. Far plane distance. Default: 1e10.
t_min (Optional[Tensor]) – Optional. Per-ray minimum distance. Tensor with shape (n_rays). If profided, the marching will start from maximum of t_min and near_plane.
t_max (Optional[Tensor]) – Optional. Per-ray maximum distance. Tensor with shape (n_rays). If profided, the marching will stop by minimum of t_max and far_plane.
render_step_size (float) – Step size for marching. Default: 1e-3.
early_stop_eps (float) – Early stop threshold for skipping invisible space. Default: 1e-4.
alpha_thre (float) – Alpha threshold for skipping empty space. Default: 0.0.
stratified (bool) – Whether to use stratified sampling. Default: False.
cone_angle (float) – Cone angle for linearly-increased step size. 0. means constant step size. Default: 0.0.

Returns:

ray_indices: Ray index of each sample. IntTensor with shape (n_samples).
t_starts: Per-sample start distance. Tensor with shape (n_samples,).
t_ends: Per-sample end distance. Tensor with shape (n_samples,).

Return type:

A tuple of {LongTensor, Tensor, Tensor}

Examples:

>>> ray_indices, t_starts, t_ends = grid.sampling(
>>>     rays_o, rays_d, render_step_size=1e-3)
>>> t_mid = (t_starts + t_ends) / 2.0
>>> sample_locs = rays_o[ray_indices] + t_mid * rays_d[ray_indices]

update_every_n_steps(step, occ_eval_fn, occ_thre=0.01, ema_decay=0.95, warmup_steps=256, n=16)¶

Update the estimator every n steps during training.

Parameters:

step (int) – Current training step.
occ_eval_fn (Callable) – A function that takes in sample locations \((N, 3)\) and returns the occupancy values \((N, 1)\) at those locations.
occ_thre (float) – Threshold used to binarize the occupancy grid. Default: 1e-2.
ema_decay (float) – The decay rate for EMA updates. Default: 0.95.
warmup_steps (int) – Sample all cells during the warmup stage. After the warmup stage we change the sampling strategy to 1/4 uniformly sampled cells together with 1/4 occupied cells. Default: 256.
n (int) – Update the grid every n steps. Default: 16.

Return type:

None

class nerfacc.PropNetEstimator(optimizer=None, scheduler=None)¶

Proposal network transmittance estimator.

References: “Mip-NeRF 360: Unbounded Anti-Aliased Neural Radiance Fields.”

Parameters:

optimizer (Optional[Optimizer]) – The optimizer to use for the proposal networks.
scheduler (Optional[_LRScheduler]) – The learning rate scheduler to use for the proposal networks.

compute_loss(trans, loss_scaler=1.0)¶

Compute the loss for the proposal networks.

Parameters:

trans (Tensor) – The transmittance of all samples. Shape (n_rays, num_samples).
loss_scaler (float) – The loss scaler. Default to 1.0.

Returns:

The loss for the proposal networks.

Return type:

Tensor

sampling(prop_sigma_fns, prop_samples, num_samples, n_rays, near_plane, far_plane, sampling_type='lindisp', stratified=False, requires_grad=False)¶

Sampling with CDFs from proposal networks.

Note

When requires_grad is True, the gradients are allowed to flow through the proposal networks, and the outputs of the proposal networks are cached to update them later when calling update_every_n_steps()

Parameters:

prop_sigma_fns (List[Callable]) – Proposal network evaluate functions. It should be a list of functions that take in samples {t_starts (n_rays, n_samples), t_ends (n_rays, n_samples)} and returns the post-activation densities (n_rays, n_samples).
prop_samples (List[int]) – Number of samples to draw from each proposal network. Should be the same length as prop_sigma_fns.
num_samples (int) – Number of samples to draw in the end.
n_rays (int) – Number of rays.
near_plane (float) – Near plane.
far_plane (float) – Far plane.
sampling_type (Literal['uniform', 'lindisp']) – Sampling type. Either “uniform” or “lindisp”. Default to “lindisp”.
stratified (bool) – Whether to use stratified sampling. Default to False.
requires_grad (bool) – Whether to allow gradients to flow through the proposal networks. Default to False.

Returns:

t_starts: The starts of the samples. Shape (n_rays, num_samples).
t_ends: The ends of the samples. Shape (n_rays, num_samples).

Return type:

A tuple of {Tensor, Tensor}

update_every_n_steps(trans, requires_grad=False, loss_scaler=1.0)¶

Update the estimator every n steps during training.

Parameters:

trans (Tensor) – The transmittance of all samples. Shape (n_rays, num_samples).
requires_grad (bool) – Whether to allow gradients to flow through the proposal networks. Default to False.
loss_scaler (float) – The loss scaler to use. Default to 1.0.

Returns:

The loss of the proposal networks for logging (a float scalar).

Return type:

float