def forward(
self,
rays_densities: torch.Tensor,
rays_features: torch.Tensor,
eps: float = 1e-10,
**kwargs,
) -> torch.Tensor:
"""
Args:
rays_densities: Per-ray density values represented with a tensor
of shape `(..., n_points_per_ray, 1)` whose values range in [0, 1].
rays_features: Per-ray feature values represented with a tensor
of shape `(..., n_points_per_ray, feature_dim)`.
eps: A lower bound added to `rays_densities` before computing
the absorbtion function (cumprod of `1-rays_densities` along
each ray). This prevents the cumprod to yield exact 0
which would inhibit any gradient-based learning.
Returns:
features: A tensor of shape `(..., feature_dim)` containing
the rendered features for each ray.
weights: A tensor of shape `(..., n_points_per_ray)` containing
the ray-specific emission-absorbtion distribution.
Each ray distribution `(..., :)` is a valid probability
distribution, i.e. it contains non-negative values that integrate
to 1, such that `weights.sum(dim=-1)==1).all()` yields `True`.
"""
_check_raymarcher_inputs(
rays_densities,
rays_features,
None,
z_can_be_none=True,
features_can_be_none=False,
density_1d=True,
)
_check_density_bounds(rays_densities)
rays_densities = rays_densities[..., 0]
absorption = _shifted_cumprod(
(1.0 + eps) - rays_densities, shift=self.surface_thickness
)
weights = rays_densities * absorption
features = (weights[..., None] * rays_features).sum(dim=-2)
return features, weights
def forward(
self,
rays_densities: torch.Tensor,
rays_features: torch.Tensor,
aux: Dict[str, Any],
ray_lengths: torch.Tensor,
density_noise_std: float = 0.0,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Dict[
str, Any]]:
"""
Args:
rays_densities: Per-ray density values represented with a tensor
of shape `(..., n_points_per_ray, 1)`.
rays_features: Per-ray feature values represented with a tensor
of shape `(..., n_points_per_ray, feature_dim)`.
aux: a dictionary with extra information.
ray_lengths: Per-ray depth values represented with a tensor
of shape `(..., n_points_per_ray, feature_dim)`.
density_noise_std: the magnitude of the noise added to densities.
Returns:
features: A tensor of shape `(..., feature_dim)` containing
the rendered features for each ray.
depth: A tensor of shape `(..., 1)` containing estimated depth.
opacities: A tensor of shape `(..., 1)` containing rendered opacsities.
weights: A tensor of shape `(..., n_points_per_ray)` containing
the ray-specific non-negative opacity weights. In general, they
don't sum to 1 but do not overcome it, i.e.
`(weights.sum(dim=-1) <= 1.0).all()` holds.
"""
_check_raymarcher_inputs(
rays_densities,
rays_features,
ray_lengths,
z_can_be_none=True,
features_can_be_none=False,
density_1d=True,
)
deltas = torch.cat(
(
ray_lengths[..., 1:] - ray_lengths[..., :-1],
self.background_opacity *
torch.ones_like(ray_lengths[..., :1]),
),
dim=-1,
)
rays_densities = rays_densities[..., 0]
if density_noise_std > 0.0:
rays_densities = (
rays_densities +
torch.randn_like(rays_densities) * density_noise_std)
if self.density_relu:
rays_densities = torch.relu(rays_densities)
weighted_densities = deltas * rays_densities
capped_densities = self._capping_function(weighted_densities)
rays_opacities = self._capping_function(
torch.cumsum(weighted_densities, dim=-1))
opacities = rays_opacities[..., -1:]
absorption_shifted = (-rays_opacities + 1.0).roll(
self.surface_thickness, dims=-1)
absorption_shifted[..., :self.surface_thickness] = 1.0
weights = self._weight_function(capped_densities, absorption_shifted)
features = (weights[..., None] * rays_features).sum(dim=-2)
depth = (weights * ray_lengths)[..., None].sum(dim=-2)
alpha = opacities if self.blend_output else 1
if self._bg_color.shape[-1] not in [1, features.shape[-1]]:
raise ValueError("Wrong number of background color channels.")
features = alpha * features + (1 - opacities) * self._bg_color
return features, depth, opacities, weights, aux