Source code for pytorch3d.implicitron.models.renderer.multipass_ea

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# pyre-unsafe

from typing import List

import torch
from pytorch3d.implicitron.models.renderer.base import ImplicitronRayBundle
from import registry, run_auto_creation

from .base import BaseRenderer, EvaluationMode, ImplicitFunctionWrapper, RendererOutput
from .ray_point_refiner import RayPointRefiner
from .raymarcher import RaymarcherBase

[docs] @registry.register class MultiPassEmissionAbsorptionRenderer( # pyre-ignore: 13 BaseRenderer, torch.nn.Module ): """ Implements the multi-pass rendering function, in particular, with emission-absorption ray marching used in NeRF [1]. First, it evaluates opacity-based ray-point weights and then optionally (in case more implicit functions are given) resamples points using importance sampling and evaluates new weights. During each ray marching pass, features, depth map, and masks are integrated: Let o_i be the opacity estimated by the implicit function, and d_i be the offset between points `i` and `i+1` along the respective ray. Ray marching is performed using the following equations:: ray_opacity_n = cap_fn(sum_i=1^n cap_fn(d_i * o_i)), weight_n = weight_fn(cap_fn(d_i * o_i), 1 - ray_opacity_{n-1}), and the final rendered quantities are computed by a dot-product of ray values with the weights, e.g. `features = sum_n(weight_n * ray_features_n)`. By default, for the EA raymarcher from [1] ( activated with `self.raymarcher_class_type="EmissionAbsorptionRaymarcher"` ):: cap_fn(x) = 1 - exp(-x), weight_fn(x) = w * x. Note that the latter can altered by changing `self.raymarcher_class_type`, e.g. to "CumsumRaymarcher" which implements the cumulative-sum raymarcher from NeuralVolumes [2]. Settings: n_pts_per_ray_fine_training: The number of points sampled per ray for the fine rendering pass during training. n_pts_per_ray_fine_evaluation: The number of points sampled per ray for the fine rendering pass during evaluation. stratified_sampling_coarse_training: Enable/disable stratified sampling in the refiner during training. Only matters if there are multiple implicit functions (i.e. in GenericModel if num_passes>1). stratified_sampling_coarse_evaluation: Enable/disable stratified sampling in the refiner during evaluation. Only matters if there are multiple implicit functions (i.e. in GenericModel if num_passes>1). append_coarse_samples_to_fine: Add the fine ray points to the coarse points after sampling. density_noise_std_train: Standard deviation of the noise added to the opacity field. return_weights: Enables returning the rendering weights of the EA raymarcher. Setting to `True` can lead to a prohibitivelly large memory consumption. blurpool_weights: Use blurpool defined in [3], on the input weights of each implicit_function except the first (implicit_functions[0]). sample_pdf_eps: Padding applied to the weights (alpha in equation 18 of [3]). raymarcher_class_type: The type of self.raymarcher corresponding to a child of `RaymarcherBase` in the registry. raymarcher: The raymarcher object used to convert per-point features and opacities to a feature render. References: [1] Mildenhall, Ben, et al. "Nerf: Representing Scenes as Neural Radiance Fields for View Synthesis." ECCV 2020. [2] Lombardi, Stephen, et al. "Neural Volumes: Learning Dynamic Renderable Volumes from Images." SIGGRAPH 2019. [3] Jonathan T. Barron, et al. "Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields." ICCV 2021. """ raymarcher_class_type: str = "EmissionAbsorptionRaymarcher" raymarcher: RaymarcherBase n_pts_per_ray_fine_training: int = 64 n_pts_per_ray_fine_evaluation: int = 64 stratified_sampling_coarse_training: bool = True stratified_sampling_coarse_evaluation: bool = False append_coarse_samples_to_fine: bool = True density_noise_std_train: float = 0.0 return_weights: bool = False blurpool_weights: bool = False sample_pdf_eps: float = 1e-5 def __post_init__(self): self._refiners = { EvaluationMode.TRAINING: RayPointRefiner( n_pts_per_ray=self.n_pts_per_ray_fine_training, random_sampling=self.stratified_sampling_coarse_training, add_input_samples=self.append_coarse_samples_to_fine, blurpool_weights=self.blurpool_weights, sample_pdf_eps=self.sample_pdf_eps, ), EvaluationMode.EVALUATION: RayPointRefiner( n_pts_per_ray=self.n_pts_per_ray_fine_evaluation, random_sampling=self.stratified_sampling_coarse_evaluation, add_input_samples=self.append_coarse_samples_to_fine, blurpool_weights=self.blurpool_weights, sample_pdf_eps=self.sample_pdf_eps, ), } run_auto_creation(self)
[docs] def forward( self, ray_bundle: ImplicitronRayBundle, implicit_functions: List[ImplicitFunctionWrapper], evaluation_mode: EvaluationMode = EvaluationMode.EVALUATION, **kwargs, ) -> RendererOutput: """ Args: ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the sampled rendering rays. implicit_functions: List of ImplicitFunctionWrappers which define the implicit functions to be used sequentially in the raymarching step. The output of raymarching with implicit_functions[n-1] is refined, and then used as input for raymarching with implicit_functions[n]. evaluation_mode: one of EvaluationMode.TRAINING or EvaluationMode.EVALUATION which determines the settings used for rendering Returns: instance of RendererOutput """ if not implicit_functions: raise ValueError("EA renderer expects implicit functions") return self._run_raymarcher( ray_bundle, implicit_functions, None, evaluation_mode, )
def _run_raymarcher( self, ray_bundle, implicit_functions, prev_stage, evaluation_mode ): density_noise_std = ( self.density_noise_std_train if evaluation_mode == EvaluationMode.TRAINING else 0.0 ) ray_deltas = ( None if ray_bundle.bins is None else torch.diff(ray_bundle.bins, dim=-1) ) output = self.raymarcher( *implicit_functions[0](ray_bundle=ray_bundle), ray_lengths=ray_bundle.lengths, ray_deltas=ray_deltas, density_noise_std=density_noise_std, ) output.prev_stage = prev_stage weights = output.weights if not self.return_weights: output.weights = None # we may need to make a recursive call if len(implicit_functions) > 1: fine_ray_bundle = self._refiners[evaluation_mode](ray_bundle, weights) output = self._run_raymarcher( fine_ray_bundle, implicit_functions[1:], output, evaluation_mode, ) return output