add prefix_trans to all render functions

nerfacc/volrend.py

@@ -155,6 +155,7 @@ def render_transmittance_from_alpha(
     packed_info: Optional[Tensor] = None,
     ray_indices: Optional[Tensor] = None,
     n_rays: Optional[int] = None,
+    prefix_trans: Optional[Tensor] = None,
 ) -> Tensor:
     """Compute transmittance :math:`T_i` from alpha :math:`\\alpha_i`.
 
@@ -171,6 +172,7 @@ def render_transmittance_from_alpha(
             Useful for flattened input.
         ray_indices: Ray indices of the flattened samples. LongTensor with shape (all_samples).
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         The rendering transmittance with the same shape as `alphas`.
@@ -191,6 +193,8 @@ def render_transmittance_from_alpha(
         packed_info = pack_info(ray_indices, n_rays)
 
     trans = exclusive_prod(1 - alphas, packed_info)
+    if prefix_trans is not None:
+        trans *= prefix_trans
     return trans
 
 
@@ -223,6 +227,7 @@ def render_transmittance_from_density(
         ray_indices: Ray indices of the flattened samples. LongTensor with shape (all_samples).
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
         prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         The rendering transmittance and opacities, both with the same shape as `sigmas`.
@@ -257,6 +262,7 @@ def render_weight_from_alpha(
     packed_info: Optional[Tensor] = None,
     ray_indices: Optional[Tensor] = None,
     n_rays: Optional[int] = None,
+    prefix_trans: Optional[Tensor] = None,
 ) -> Tuple[Tensor, Tensor]:
     """Compute rendering weights :math:`w_i` from opacity :math:`\\alpha_i`.
 
@@ -273,6 +279,7 @@ def render_weight_from_alpha(
             Useful for flattened input.
         ray_indices: Ray indices of the flattened samples. LongTensor with shape (all_samples).
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         The rendering weights and transmittance, both with the same shape as `alphas`.
@@ -289,7 +296,7 @@ def render_weight_from_alpha(
 
     """
     trans = render_transmittance_from_alpha(
-        alphas, packed_info, ray_indices, n_rays
+        alphas, packed_info, ray_indices, n_rays, prefix_trans
     )
     weights = trans * alphas
     return weights, trans
@@ -302,6 +309,7 @@ def render_weight_from_density(
     packed_info: Optional[Tensor] = None,
     ray_indices: Optional[Tensor] = None,
     n_rays: Optional[int] = None,
+    prefix_trans: Optional[Tensor] = None,
 ) -> Tuple[Tensor, Tensor, Tensor]:
     """Compute rendering weights :math:`w_i` from density :math:`\\sigma_i` and interval :math:`\\delta_i`.
 
@@ -320,6 +328,7 @@ def render_weight_from_density(
             Useful for flattened input.
         ray_indices: Ray indices of the flattened samples. LongTensor with shape (all_samples).
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         The rendering weights, transmittance and opacities, both with the same shape as `sigmas`.
@@ -340,7 +349,7 @@ def render_weight_from_density(
 
     """
     trans, alphas = render_transmittance_from_density(
-        t_starts, t_ends, sigmas, packed_info, ray_indices, n_rays
+        t_starts, t_ends, sigmas, packed_info, ray_indices, n_rays, prefix_trans
     )
     weights = trans * alphas
     return weights, trans, alphas
@@ -354,6 +363,7 @@ def render_visibility_from_alpha(
     n_rays: Optional[int] = None,
     early_stop_eps: float = 1e-4,
     alpha_thre: float = 0.0,
+    prefix_trans: Optional[Tensor] = None,
 ) -> Tensor:
     """Compute visibility from opacity :math:`\\alpha_i`.
 
@@ -374,6 +384,7 @@ def render_visibility_from_alpha(
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
         early_stop_eps: The early stopping threshold on transmittance.
         alpha_thre: The threshold on opacity.
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         A boolean tensor indicating which samples are visible. Same shape as `alphas`.
@@ -392,7 +403,7 @@ def render_visibility_from_alpha(
 
     """
     trans = render_transmittance_from_alpha(
-        alphas, packed_info, ray_indices, n_rays
+        alphas, packed_info, ray_indices, n_rays, prefix_trans
     )
     vis = trans >= early_stop_eps
     if alpha_thre > 0:
@@ -410,6 +421,7 @@ def render_visibility_from_density(
     n_rays: Optional[int] = None,
     early_stop_eps: float = 1e-4,
     alpha_thre: float = 0.0,
+    prefix_trans: Optional[Tensor] = None,
 ) -> Tensor:
     """Compute visibility from density :math:`\\sigma_i` and interval :math:`\\delta_i`.
 
@@ -430,6 +442,7 @@ def render_visibility_from_density(
         n_rays: Number of rays. Only useful when `ray_indices` is provided.
         early_stop_eps: The early stopping threshold on transmittance.
         alpha_thre: The threshold on opacity.
+        prefix_trans: The pre-computed transmittance of the samples. Tensor with shape (all_samples,).
 
     Returns:
         A boolean tensor indicating which samples are visible. Same shape as `alphas`.
@@ -452,7 +465,7 @@ def render_visibility_from_density(
 
     """
     trans, alphas = render_transmittance_from_density(
-        t_starts, t_ends, sigmas, packed_info, ray_indices, n_rays
+        t_starts, t_ends, sigmas, packed_info, ray_indices, n_rays, prefix_trans
     )
     vis = trans >= early_stop_eps
     if alpha_thre > 0: