From 4f2157cd5849d2565b1773df6f19de6fa4f848ab Mon Sep 17 00:00:00 2001 From: hturki Date: Thu, 20 Jul 2023 12:48:25 -0400 Subject: [PATCH] support packed samples in semantic renderer (#2247) --- nerfstudio/model_components/renderers.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index ae50b0e12..24090872e 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -410,10 +410,17 @@ class SemanticRenderer(nn.Module): cls, semantics: Float[Tensor, "*bs num_samples num_classes"], weights: Float[Tensor, "*bs num_samples 1"], + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, + num_rays: Optional[int] = None, ) -> Float[Tensor, "*bs num_classes"]: """Calculate semantics along the ray.""" - sem = torch.sum(weights * semantics, dim=-2) - return sem + if ray_indices is not None and num_rays is not None: + # Necessary for packed samples from volumetric ray sampler + return nerfacc.accumulate_along_rays( + weights[..., 0], values=semantics, ray_indices=ray_indices, n_rays=num_rays + ) + else: + return torch.sum(weights * semantics, dim=-2) class NormalsRenderer(nn.Module):