From 83d0c58c0f71a358e18e02877e5af3b9498da35a Mon Sep 17 00:00:00 2001 From: LoYoT Date: Sat, 22 Apr 2023 01:01:00 +0800 Subject: [PATCH] add prefix_trans to all render functions --- nerfacc/volrend.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/nerfacc/volrend.py b/nerfacc/volrend.py index 2d567eba..a5fc7108 100644 --- a/nerfacc/volrend.py +++ b/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: