|
|
|
"""Contains the function to march rays (integration).""" |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
|
|
__all__ = ['Integrator'] |
|
|
|
|
|
class Integrator(torch.nn.Module): |
|
"""Defines the class to help march rays, i.e. do integral along each ray. |
|
|
|
The ray marcher takes the raw output of the implicit representation |
|
(including colors(i.e. rgbs) and densities(i.e. sigmas)) and uses the |
|
volume rendering equation to produce composited colors and depths. |
|
""" |
|
|
|
def __init__(self): |
|
super().__init__() |
|
|
|
def integration(self, rgbs, sigmas, depths, rendering_options): |
|
"""Integrate the values along the ray. |
|
|
|
`N` denotes batch size. |
|
`R` denotes the number of rays, equals `H * W`. |
|
`K` denotes the number of points on each ray. |
|
|
|
Args: |
|
rgbs (torch.tensor): colors' value of each point in the fields, with |
|
shape [N, R, K, 3]. |
|
sigmas (torch.tensor): densities' value of each point in the fields, |
|
with shape [N, R, K, 1]. |
|
depths (torch.tensor): depths' value of each point in the fields, |
|
with shape [N, R, K, 1]. |
|
rendering_options (dict): Additional keyword arguments of rendering |
|
option. |
|
|
|
Returns: |
|
A dictionary, containing |
|
- `composite_rgb`: camera radius w.r.t. the world coordinate |
|
system, with shape [N, R, 3]. |
|
- `composite_depth`: camera polar w.r.t. the world coordinate |
|
system, with shape [N, R, 1]. |
|
- `weights`: importance weights of each point in the field, |
|
with shape [N, R, K, 1]. |
|
""" |
|
num_dims = rgbs.ndim |
|
assert num_dims == 4 |
|
assert sigmas.ndim == num_dims and depths.ndim == num_dims |
|
|
|
N, R, K = rgbs.shape[:3] |
|
|
|
|
|
deltas = depths[:, :, 1:] - depths[:, :, :-1] |
|
if rendering_options.get('use_max_depth', False): |
|
max_depth = rendering_options.get('max_depth', None) |
|
if max_depth is not None: |
|
delta_inf = max_depth - deltas[:, :, -1:] |
|
else: |
|
delta_inf = 1e10 * torch.ones_like(deltas[:, :, :1]) |
|
deltas = torch.cat([deltas, delta_inf], -2) |
|
if rendering_options.get('no_dist', False): |
|
deltas[:] = 1 |
|
|
|
use_mid_point = rendering_options.get('use_mid_point', True) |
|
if use_mid_point: |
|
rgbs = (rgbs[:, :, :-1] + rgbs[:, :, 1:]) / 2 |
|
sigmas = (sigmas[:, :, :-1] + sigmas[:, :, 1:]) / 2 |
|
depths = (depths[:, :, :-1] + depths[:, :, 1:]) / 2 |
|
|
|
clamp_mode = rendering_options.get('clamp_mode', 'mipnerf') |
|
if clamp_mode == 'softplus': |
|
sigmas = F.softplus(sigmas) |
|
elif clamp_mode == 'relu': |
|
sigmas = F.relu(sigmas) |
|
elif clamp_mode == 'mipnerf': |
|
sigmas = F.softplus(sigmas - 1) |
|
else: |
|
raise ValueError(f'Invalid clamping mode: `{clamp_mode}`!\n') |
|
|
|
alphas = 1 - torch.exp(- deltas * sigmas) |
|
alphas_shifted = torch.cat( |
|
[torch.ones_like(alphas[:, :, :1]), 1 - alphas + 1e-10], -2) |
|
weights = alphas * torch.cumprod(alphas_shifted, -2)[:, :, :-1] |
|
weights_sum = weights.sum(2) |
|
if rendering_options.get('last_back', False): |
|
weights[:, :, -1] = weights[:, :, -1] + (1 - weights_sum) |
|
|
|
composite_rgb = torch.sum(weights * rgbs, -2) |
|
composite_depth = torch.sum(weights * depths, -2) |
|
|
|
if rendering_options.get('normalize_rgb', False): |
|
composite_rgb = composite_rgb / weights_sum |
|
if rendering_options.get('normalize_depth', True): |
|
composite_depth = composite_depth / weights_sum |
|
if rendering_options.get('clip_depth', True): |
|
composite_depth = torch.nan_to_num(composite_depth, float('inf')) |
|
composite_depth = torch.clip(composite_depth, torch.min(depths), |
|
torch.max(depths)) |
|
|
|
if rendering_options.get('white_back', False): |
|
composite_rgb = composite_rgb + 1 - weights_sum |
|
|
|
composite_rgb = composite_rgb * 2 - 1 |
|
|
|
results = { |
|
'composite_rgb': composite_rgb, |
|
'composite_depth': composite_depth, |
|
'weights': weights |
|
} |
|
|
|
return results |
|
|
|
def forward(self, rgbs, sigmas, depths, rendering_options): |
|
results = self.integration(rgbs, sigmas, depths, rendering_options) |
|
return results |
|
|
