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