Reference¶
N^3-Tree¶
- class N3Tree(N=2, data_dim=None, depth_limit=10, init_reserve=1, init_refine=0, geom_resize_fact=1.0, radius=0.5, center=[0.5, 0.5, 0.5], data_format='RGBA', extra_data=None, device='cpu', dtype=torch.float32, map_location=None)¶
Bases:
torch.nn.modules.module.Module
PyTorch \(N^3\)-tree library with CUDA acceleration. By \(N^3\)-tree we mean a 3D tree with branching factor N at each interior node, where \(N=2\) is the familiar octree.
Warning
nn.Parameters can change size, which makes current optimizers invalid. If any
refine(): or :code:`shrink_to_fit()
call returns True, orexpand(), shrink()
is used, please re-make any optimizers- __init__(N=2, data_dim=None, depth_limit=10, init_reserve=1, init_refine=0, geom_resize_fact=1.0, radius=0.5, center=[0.5, 0.5, 0.5], data_format='RGBA', extra_data=None, device='cpu', dtype=torch.float32, map_location=None)¶
Construct N^3 Tree
- Parameters
N – int branching factor N
data_dim – int size of data stored at each leaf (NEW in 0.2.28: optional if data_format other than RGBA is given). If data_format = “RGBA” or empty, this defaults to 4.
depth_limit – int maximum depth of tree to stop branching/refining Note that the root is at depth -1. Size
N^[-10]
leaves (1/1024 for octree) for example are depth 9.max_depth
applies to the same depth values.init_reserve – int amount of nodes to reserve initially
init_refine – int number of times to refine entire tree initially inital resolution will be
[N^(init_refine + 1)]^3
. initial max_depth will be init_refine.geom_resize_fact – float geometric resizing factor
radius – float or list, 1/2 side length of cube (possibly in each dim)
center – list center of space
data_format – a string to indicate the data format.
RGBA | SH# | SG# | ASG#
extra_data – extra data to include with tree
device – str device to put data
dtype – str tree data type, torch.float32 (default) | torch.float64
map_location – str DEPRECATED old name for device (will override device and warn)
- set(indices, values, cuda=True)¶
Set tree values,
- Parameters
indices – torch.Tensor
(Q, 3)
values – torch.Tensor
(Q, K)
cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version.
- forward(indices, cuda=True, want_node_ids=False, world=True)¶
Get tree values. Differentiable.
- Parameters
indices –
(Q, 3)
the pointscuda – whether to use CUDA kernel if available. If false, uses only PyTorch version.
want_node_ids – if true, returns node ID for each query.
world – use world space instead of
[0,1]^3
, default True
- Returns
(Q, data_dim), [(Q)]
- snap(indices)¶
Snap indices to lowest corner of corresponding leaf voxel
- Parameters
indices –
(B, 3)
indices to snap- Returns
(B, 3)
- partial(data_sel=None, data_format=None, dtype=None, device=None)¶
Get partial tree with some of the data dimensions (channels) E.g.
tree.partial(-1)
to get tree with data_dim 1 of last channel only- Parameters
data_sel – data channel selector, default is all channels
data_format – data format for new tree, default is current format
dtype – new data type, torch.float32 | torch.float64
device – where to put result tree
- Returns
partial N3Tree (copy)
- expand(data_format, data_dim=None, remap=None)¶
Modify the size of the data stored at the octree leaves.
- Parameters
data_format – new data format,
RGBA | SH# | SG# | ASG#
data_dim – data dimension; inferred from data_format by default only needed if data_format is RGBA.
remap – mapping of old data to new data. For each leaf, we will do
new_data[remap] = old_data
if the new data_dim is equal or larger, or new_data = old_data[remap] else. By default, this will be inferred automatically (maps basis functions in the correct way).
Warning
Will change the nn.Parameter size (data), breaking optimizer! Please re-create the optimizer
- shrink(data_format, data_dim=None, remap=None)¶
Modify the size of the data stored at the octree leaves. (Alias of expand, because it can actually both expand and shrink)
- Parameters
data_format – new data format,
RGBA | SH# | SG# | ASG#
data_dim – data dimension; inferred from data_format by default only needed if data_format is RGBA.
remap – mapping of old data to new data. For each leaf, we will do
new_data = old_data[remap]
if the new data_dim gets smaller, ornew_data[remap] = old_data
else. By default, this will be inferred automatically (maps basis functions in the correct way).
Warning
Will change the nn.Parameter size (data), breaking optimizer! Please re-create the optimizer
- clone(device=None)¶
Deep copy the tree
- Parameters
device – device of output tree (could e.g. copy cuda tree to cpu)
- merge(frontier_sel=None, op=<built-in method mean of type object>)¶
Merge leaves into selected ‘frontier’ nodes (i.e., nodes for which all children are leaves). Use
shrink_to_fit()
to recover memory freed.- Parameters
frontier_sel – selector (int, mask, list of indices etc) for frontier nodes. In same order as
reduce_frontier()
. Default all nodes. Typical use: usereduce_frontier(...)
to determine conditions for merge, then pass bool mask (of lengthn_frontier
) or indices tomerge()
.op – reduction to combine child leaves into node. E.g. torch.max, torch.mean. Should take a positional argument
x
(B, N, data_dim)
and a named parameterdim
(always 1), and return a matrix of(B, data_dim)
. If a tuple is returned, uses first result.
- property n_frontier¶
Number of frontier nodes
- property frontier_depth¶
Depth of frontier nodes, size
(n_frontier)
- reduce_frontier(op=<built-in method mean of type object>, dim=None, grad=False)¶
Reduce child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves).
- Parameters
op – reduction to combine child leaves into node. E.g. torch.max, torch.mean. Should take a positional argument
x
(B, N, in_dim <= data_dim)
and a named parameterdim
(always 1), and return a matrix of(B, your_out_dim)
.dim – dimension(s) of data to return, e.g.
-1
returns last data dimension for all ‘frontier’ nodesgrad – if True, returns a tensor differentiable wrt tree data. Default False.
- Example
def mean_and_var_func(x, dim=1): # (n_frontier, tree.N^3, data_dim) -> (n_frontier, 2 * data_dim) # Outputs mean and variance over children of each frontier node return torch.cat([torch.mean(x, dim=1), torch.var(x, dim=1)], dim=-1) # returns: (n_frontier, 2 * data_dim) tree.reduce_frontier(mean_and_var_func)
- Returns
reduced tensor
- max_frontier(dim=None, grad=False)¶
Takes max over child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves). This is simply reduce_frontier with torch.max operation, taking the returned values and discarding the argmax part.
- Parameters
dim – dimension(s) of data to return, e.g.
-1
returns last data dimension for all ‘frontier’ nodesgrad – if True, returns a tensor differentiable wrt tree data. Default False.
- Returns
reduced tensor
- diam_frontier(dim=None, grad=False, scale=1.0)¶
Takes diameter over child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves).
- Parameters
dim – dimension(s) of data to return, e.g.
-1
returns last data dimension for all ‘frontier’ nodesgrad – if True, returns a tensor differentiable wrt tree data. Default False.
- Returns
reduced tensor
- check_integrity()¶
Do some checks to verify the tree’s structural integrity, mostly for debugging. Errors with message if check fails; does nothing else.
- refine(repeats=1, sel=None)¶
Refine each selected leaf node, respecting depth_limit.
- Parameters
repeats – int number of times to repeat refinement
sel –
(N, 4)
node selector. Default selects all leaves.
- Returns
True iff N3Tree.data parameter was resized, requiring optimizer reinitialization if you’re using an optimizer
Warning
The parameter
tree.data
can change due to refinement. If any refine() call returns True, please re-make any optimizers usingtree.params()
.Warning
The selector
sel
is assumed to contain unique leaf indices. If there are duplicates memory will be wasted. We do not dedup here for efficiency reasons.
- shrink_to_fit()¶
Shrink data & buffers to tightly needed fit tree data, possibly dealing with fragmentation caused by merging. This is called by the
save()
function by default, unlessshrink=False
is specified there.Warning
Will change the nn.Parameter size (data), breaking optimizer!
- property max_depth¶
Maximum tree depth - 1
- accumulate_weights(op: str = 'sum')¶
Begin weight accumulation.
- Parameters
op – reduction to apply weight in each voxel, sum | max
Warning
Weight accumulator has not been validated and may have bugs
- Example
with tree.accumulate_weights() as accum: ... # (n_leaves) in same order as values etc. accum = accum()
- save(path, shrink=True, compress=True)¶
Save to from npz file
- Parameters
path – npz path
shrink – if True (default), applies shrink_to_fit before saving
compress – whether to compress the npz; may be slow
- classmethod load(path, device='cpu', dtype=torch.float32, map_location=None)¶
Load from npz file
- Parameters
path – npz path
device – str device to put data
dtype – str torch.float32 (default) | torch.float64
map_location – str DEPRECATED old name for device
- world2tree(indices)¶
Scale world points to tree (\([0,1]^3\))
- tree2world(indices)¶
Scale tree points (\([0,1]^3\)) to world accoording to center/radius
N^3-Tree View¶
- class N3TreeView(tree, key)¶
Bases:
object
- __init__(tree, key)¶
Initialize self. See help(type(self)) for accurate signature.
- refine(repeats=1)¶
Refine selected leaves using tree.refine
- property values¶
Values of the selected leaves (autograd enabled)
- Returns
(n_leaves, data_dim) float32 note this is 2D even if key is int
- property values_nograd¶
Values of the selected leaves (no autograd)
- Returns
(n_leaves, data_dim) float32 note this is 2D even if key is int
- property depths¶
Get a list of selected leaf depths in tree, in same order as values, corners. Root is at depth -1. Any children of root will have depth 0, etc.
- Returns
(n_leaves) int32
- property lengths¶
Get a list of selected leaf side lengths in tree (world dimensions), in same order as values, corners, depths
- Returns
(n_leaves, 3) float
- property lengths_local¶
Get a list of selected leaf side lengths in tree (local index \([0, 1]^3\)), in same order as values, corners, depths
- Returns
(n_leaves) float
- property corners¶
Get a list of selected leaf lower corners in tree (world coordinates), in same order as values, lengths, depths
- Returns
(n_leaves, 3) float
- property corners_local¶
Get a list of selected leaf lower corners in tree (local index \([0, 1]^3\)), in same order as values, lengths, depths
- Returns
(n_leaves, 3) float
- sample(n_samples, device=None)¶
Sample n_samples uniform points in each selected leaf (world coordinates)
- Parameters
n_samples – samples for each leaf
device – device to output random samples in
- Returns
(n_leaves, n_samples, 3) float
- sample_local(n_samples)¶
Sample n_samples uniform points in each selected leaf (local index \([0, 1]^3\))
- Returns
(n_leaves, n_samples, 3) float
- aux(arr)¶
Index an auxiliary tree data array of size (capacity, N, N, N, Any) using this view
- normal_(mean=0.0, std=1.0)¶
Set all values to random normal FIXME: inefficient
- Parameters
mean – normal mean
std – normal std
- uniform_(min=0.0, max=1.0)¶
Set all values to random uniform FIXME: inefficient
- Parameters
min – interval min
max – interval max
- clamp_(min=None, max=None)¶
Clamp. FIXME: inefficient
- Parameters
min – clamp min value, None=disable
max – clamp max value, None=disable
- relu_()¶
Apply relu to all elements. FIXME: inefficient
- sigmoid_()¶
Apply sigmoid to all elements. FIXME: inefficient
- nan_to_num_(inf_val=20000.0)¶
Convert nans to 0.0 and infs to inf_val FIXME: inefficient
Differentiable Volume Renderer¶
- class VolumeRenderer(tree, step_size: float = 0.001, background_brightness: float = 1.0, ndc: Optional[svox.renderer.NDCConfig] = None, min_comp: int = 0, max_comp: int = - 1, density_softplus: bool = False, rgb_padding: float = 0.0)¶
Bases:
torch.nn.modules.module.Module
Volume renderer
- __init__(tree, step_size: float = 0.001, background_brightness: float = 1.0, ndc: Optional[svox.renderer.NDCConfig] = None, min_comp: int = 0, max_comp: int = - 1, density_softplus: bool = False, rgb_padding: float = 0.0)¶
Construct volume renderer associated with given N^3 tree. You can construct multiple renderer instances for the same tree; the renderer class only stores configurations and no persistent tensors.
The renderer traces rays with origins/dirs within the octree boundaries, detection ray-voxel intersections. The color and density within each voxel is assumed constant, and no interpolation is performed.
For each intersection point, it queries the tree, assuming the last data dimension is density (sigma) and the rest of the dimensions are color, formatted according to tree.data_format. It then applies SH/SG/ASG basis functions, if any, according to viewdirs. Sigmoid will be applied to these colors to normalize them, and optionally a shifted softplus is applied to the density.
- Parameters
tree – N3Tree instance for rendering
step_size – float step size eps, added to each voxel aabb intersection step
background_brightness – float background brightness, 1.0 = white
ndc – NDCConfig, NDC coordinate configuration, namedtuple(width, height, focal). None = no NDC, use usual coordinates
min_comp – minimum SH/SG component to render.
max_comp – maximum SH/SG component to render, -1=last. Set
min_comp = max_comp
to render a particular component. Default means all. Tip: If you setmin_comp > max_comp
, the renderer will render all colors as 0.5 luminosity gray. This is still differentiable and be used to implement a mask loss.density_softplus – if true, applies \(\log(1 + \exp(sigma - 1))\). Mind the shift -1! (from mip-NeRF). Please note softplus will NOT be compatible with volrend, please pre-apply it .
rgb_padding – to avoid oversaturating the sigmoid, applies
* (1 + 2 * rgb_padding) - rgb_padding
to colors after sigmoid (from mip-NeRF). Please note the padding will NOT be compatible with volrend, although most likely the effect is very small. 0.001 is a reasonable value to try.
- forward(rays: svox.renderer.Rays, cuda=True, fast=False)¶
Render a batch of rays. Differentiable.
- Parameters
rays – namedtuple
svox.Rays
of origins(B, 3)
, dirs(B, 3):, viewdirs :code:`(B, 3)
cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version. Only True supported right now
fast – if True, enables faster evaluation, potentially leading to some loss of accuracy.
- Returns
(B, rgb_dim)
. Where rgb_dim istree.data_dim - 1
ifdata_format.format == DataFormat.RGBA
or(tree.data_dim - 1) / tree.data_format.basis_dim
else.
- render_persp(c2w, width=800, height=800, fx=1111.111, fy=None, cuda=True, fast=False)¶
Render a perspective image. Differentiable.
- Parameters
c2w – torch.Tensor (3, 4) or (4, 4) camera pose matrix (c2w)
width – int output image width
height – int output image height
fx – float output image focal length (x)
fy – float output image focal length (y), if not specified uses fx
cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version. Only True supported right now
fast – if True, enables faster evaluation, potentially leading to some loss of accuracy.
- Returns
(height, width, rgb_dim)
where rgb_dim istree.data_dim - 1
ifdata_format.format == DataFormat.RGBA
or(tree.data_dim - 1) / tree.data_format.basis_dim
else.
- se_grad(rays: svox.renderer.Rays, colors)¶
Returns rendered color + gradient and Hessian diagonal of the total squared error: \(\frac{1}{2} \sum_{r \in \mathcal{R}} (\hat{C}(r) - C(r))^2\) where \(\hat{C}(r)\) is computed from the ray and \(C(r)\) comes from the provided tensor
colors
. This is the arbitrary ray-batch version ofse_grad
. This is useful for diagonal NNLS methods for scaling step sizes. Note currently the Hessian is actually the squared norm of Jacobian rows as in Gauss-Newton algorithms.The tree’s rendered output dimension (rgb_dim) cannot be greater than 4 (this is almost always true, don’t need to worry).
- Parameters
rays – namedtuple
svox.Rays
of origins(B, 3)
, dirs(B, 3):, viewdirs :code:`(B, 3)
colors – torch.Tensor
(B, 3)
reference colors
- Returns
colors (B, rgb_dim), grad (shape of tree.data), diag_hessian (shape of tree.data)
- se_grad_persp(c2w, colors, width=800, height=800, fx=1111.111, fy=None)¶
Returns rendered color + gradient and Hessian diagonal of the total squared error: \(\frac{1}{2} \sum_{r \in \mathcal{R}} (\hat{C}(r) - C(r))^2\) where \(\hat{C}(r)\) is computed from the ray and \(C(r)\) comes from the provided tensor
colors
. This is the image-batch version ofse_grad
. This is useful for diagonal NNLS methods for scaling step sizes. Note currently the Hessian is actually the squared norm of Jacobian rows as in Gauss-Newton algorithms.The tree’s rendered output dimension (rgb_dim) cannot be greater than 4 (this is almost always true, don’t need to worry).
- Parameters
c2w – torch.Tensor (3, 4) or (4, 4) camera pose matrix (c2w)
colors – torch.Tensor
(H, W, 3)
reference colorswidth – int output image width
height – int output image height
fx – float output image focal length (x)
fy – float output image focal length (y), if not specified uses fx
- Returns
colors (H, W, rgb_dim), grad (shape of tree.data), diag_hessian (shape of tree.data)
- static persp_rays(c2w, width=800, height=800, fx=1111.111, fy=None)¶
Generate perspective camera rays in row major order, then usable for renderer’s forward method. NDC is not supported currently.
- Parameters
c2w – torch.Tensor (3, 4) or (4, 4) camera pose matrix (c2w)
width – int output image width
height – int output image height
fx – float output image focal length (x)
fy – float output image focal length (y), if not specified uses fx
- Returns
rays namedtuple svox.Rays of origins
(H*W, 3)
, dirs(H*W, 3):, viewdirs :code:`(H*W, 3)
, where H = W.