diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 1f95de95f..6a8e2a5f9 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -37,7 +37,7 @@ from torchtyping import TensorType from typing_extensions import Literal from nerfstudio.cameras.rays import RaySamples -from nerfstudio.utils.math import components_from_spherical_harmonics +from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize BACKGROUND_COLOR_OVERRIDE: Optional[TensorType[3]] = None @@ -328,9 +328,18 @@ class NormalsRenderer(nn.Module): cls, normals: TensorType["bs":..., "num_samples", 3], weights: TensorType["bs":..., "num_samples", 1], + normalize: bool = True, ) -> TensorType["bs":..., 3]: - """Calculate normals along the ray.""" + """Calculate normals along the ray. + + Args: + normals: Normals for each sample. + weights: Weights of each sample. + normalize: Normalize normals. + """ n = torch.sum(weights * normals, dim=-2) + if normalize: + n = safe_normalize(n) return n diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index c996e20b3..91b044884 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -250,8 +250,10 @@ class NerfactoModel(Model): } if self.config.predict_normals: - outputs["normals"] = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) - outputs["pred_normals"] = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + outputs["normals"] = (normals + 1) / 2 + outputs["pred_normals"] = (pred_normals + 1) / 2 # These use a lot of GPU memory, so we avoid storing them for eval. if self.training: