diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index eb6c5f9b6..9f4168517 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -151,6 +151,10 @@ def generate_point_cloud( CONSOLE.print(f"Please set --normal_output_name to one of: {outputs.keys()}", justify="center") sys.exit(1) normal = outputs[normal_output_name] + assert ( + torch.min(normal) >= 0.0 and torch.max(normal) <= 1.0 + ), "Normal values from method output must be in [0, 1]" + normal = (normal * 2.0) - 1.0 point = ray_bundle.origins + ray_bundle.directions * depth if use_bounding_box: diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index e7ee3d14c..e10b3cad2 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -50,6 +50,7 @@ from nerfstudio.model_components.renderers import ( RGBRenderer, ) from nerfstudio.model_components.scene_colliders import NearFarCollider +from nerfstudio.model_components.shaders import NormalsShader from nerfstudio.models.base_model import Model from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig @@ -162,6 +163,9 @@ class NerfplayerNerfactoModel(NerfactoModel): self.renderer_depth = DepthRenderer(method="expected") # for depth loss self.renderer_normals = NormalsRenderer() + # shaders + self.normals_shader = NormalsShader() + # losses self.rgb_loss = MSELoss() @@ -192,8 +196,12 @@ class NerfplayerNerfactoModel(NerfactoModel): } 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) + outputs["normals"] = self.normals_shader( + self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights) + ) + outputs["pred_normals"] = self.normals_shader( + self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights) + ) # These use a lot of GPU memory, so we avoid storing them for eval. if self.training: