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):