Add option for model normals on ply export (#1965)

* Update point cloud exporter to accept the same arguments as the poisson mesh export wrt normals + hoist validate pipeline up so both exporters can use it

* Update viewer to support both normal methods for point cloud and poisson export

* Lint fixes

* Undo eslint comments

* Add types and documentation for args in validate pipeline function

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>

nerfstudio/scripts/exporter.py

@@ -59,6 +59,36 @@ class Exporter:
     """Path to the output directory."""
 
 
+def validate_pipeline(normal_method: str, normal_output_name: str, pipeline: Pipeline) -> None:
+    """Check that the pipeline is valid for this exporter.
+
+    Args:
+        normal_method: Method to estimate normals with. Either "open3d" or "model_output".
+        normal_output_name: Name of the normal output.
+        pipeline: Pipeline to evaluate with.
+    """
+    if normal_method == "model_output":
+        CONSOLE.print("Checking that the pipeline has a normal output.")
+        origins = torch.zeros((1, 3), device=pipeline.device)
+        directions = torch.ones_like(origins)
+        pixel_area = torch.ones_like(origins[..., :1])
+        camera_indices = torch.zeros_like(origins[..., :1])
+        ray_bundle = RayBundle(
+            origins=origins, directions=directions, pixel_area=pixel_area, camera_indices=camera_indices
+        )
+        outputs = pipeline.model(ray_bundle)
+        if normal_output_name not in outputs:
+            CONSOLE.print(f"[bold yellow]Warning: Normal output '{normal_output_name}' not found in pipeline outputs.")
+            CONSOLE.print(f"Available outputs: {list(outputs.keys())}")
+            CONSOLE.print(
+                "[bold yellow]Warning: Please train a model with normals "
+                "(e.g., nerfacto with predicted normals turned on)."
+            )
+            CONSOLE.print("[bold yellow]Warning: Or change --normal-method")
+            CONSOLE.print("[bold yellow]Exiting early.")
+            sys.exit(1)
+
+
 @dataclass
 class ExportPointCloud(Exporter):
     """Export NeRF as a point cloud."""
@@ -67,8 +97,10 @@ class ExportPointCloud(Exporter):
     """Number of points to generate. May result in less if outlier removal is used."""
     remove_outliers: bool = True
     """Remove outliers from the point cloud."""
-    estimate_normals: bool = False
-    """Estimate normals for the point cloud."""
+    normal_method: Literal["open3d", "model_output"] = "model_output"
+    """Method to estimate normals with."""
+    normal_output_name: str = "normals"
+    """Name of the normal output."""
     depth_output_name: str = "depth"
     """Name of the depth output."""
     rgb_output_name: str = "rgb"
@@ -92,19 +124,24 @@ class ExportPointCloud(Exporter):
 
         _, pipeline, _, _ = eval_setup(self.load_config)
 
+        validate_pipeline(self.normal_method, self.normal_output_name, pipeline)
+
         # Increase the batchsize to speed up the evaluation.
         assert isinstance(pipeline.datamanager, VanillaDataManager)
         assert pipeline.datamanager.train_pixel_sampler is not None
         pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch
 
+        # Whether the normals should be estimated based on the point cloud.
+        estimate_normals = self.normal_method == "open3d"
+
         pcd = generate_point_cloud(
             pipeline=pipeline,
             num_points=self.num_points,
             remove_outliers=self.remove_outliers,
-            estimate_normals=self.estimate_normals,
+            estimate_normals=estimate_normals,
             rgb_output_name=self.rgb_output_name,
             depth_output_name=self.depth_output_name,
-            normal_output_name=None,
+            normal_output_name=self.normal_output_name if self.normal_method == "model_output" else None,
             use_bounding_box=self.use_bounding_box,
             bounding_box_min=self.bounding_box_min,
             bounding_box_max=self.bounding_box_max,
@@ -237,31 +274,6 @@ class ExportPoissonMesh(Exporter):
     std_ratio: float = 10.0
     """Threshold based on STD of the average distances across the point cloud to remove outliers."""
 
-    def validate_pipeline(self, pipeline: Pipeline) -> None:
-        """Check that the pipeline is valid for this exporter."""
-        if self.normal_method == "model_output":
-            CONSOLE.print("Checking that the pipeline has a normal output.")
-            origins = torch.zeros((1, 3), device=pipeline.device)
-            directions = torch.ones_like(origins)
-            pixel_area = torch.ones_like(origins[..., :1])
-            camera_indices = torch.zeros_like(origins[..., :1])
-            ray_bundle = RayBundle(
-                origins=origins, directions=directions, pixel_area=pixel_area, camera_indices=camera_indices
-            )
-            outputs = pipeline.model(ray_bundle)
-            if self.normal_output_name not in outputs:
-                CONSOLE.print(
-                    f"[bold yellow]Warning: Normal output '{self.normal_output_name}' not found in pipeline outputs."
-                )
-                CONSOLE.print(f"Available outputs: {list(outputs.keys())}")
-                CONSOLE.print(
-                    "[bold yellow]Warning: Please train a model with normals "
-                    "(e.g., nerfacto with predicted normals turned on)."
-                )
-                CONSOLE.print("[bold yellow]Warning: Or change --normal-method")
-                CONSOLE.print("[bold yellow]Exiting early.")
-                sys.exit(1)
-
     def main(self) -> None:
         """Export mesh"""
 
@@ -269,7 +281,8 @@ class ExportPoissonMesh(Exporter):
             self.output_dir.mkdir(parents=True)
 
         _, pipeline, _, _ = eval_setup(self.load_config)
-        self.validate_pipeline(pipeline)
+
+        validate_pipeline(self.normal_method, self.normal_output_name, pipeline)
 
         # Increase the batchsize to speed up the evaluation.
         assert isinstance(pipeline.datamanager, VanillaDataManager)

nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx

@@ -3,26 +3,8 @@ import { useControls, useStoreContext } from 'leva';
 import { useEffect } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 import ContentCopyRoundedIcon from '@mui/icons-material/ContentCopyRounded';
-
 import { Button } from '@mui/material';
-
-function get_normal_outputs(output_options) {
-  // Get a list of normal outputs from the Model
-  let normal_options = [];
-  if (output_options) {
-    // check which outputs have normals
-    for (let i = 0; i < output_options.length; i += 1) {
-      const output_name = output_options[i];
-      if (output_name.includes('normals')) {
-        normal_options.push(output_name);
-      }
-    }
-  }
-  if (normal_options.length === 0) {
-    normal_options = ['none'];
-  }
-  return normal_options;
-}
+import { get_normal_methods } from '../../../utils';
 
 export default function MeshSubPanel(props) {
   const update_box_center = props.update_box_center;
@@ -47,7 +29,7 @@ export default function MeshSubPanel(props) {
   const output_options = useSelector(
     (state) => state.renderingState.output_options,
   );
-  const normal_options = get_normal_outputs(output_options);
+  const normal_options = get_normal_methods(output_options);
   const mesh_method_options = ['poisson', 'tsdf'];
 
   const config_filename = `${config_base_dir}/config.yml`;
@@ -64,8 +46,11 @@ export default function MeshSubPanel(props) {
       },
       normal_options: {
         label: 'Normal Options',
-        options: [...new Set(normal_options)],
-        value: normal_options[0],
+        options: [...normal_options],
+        value:
+          normal_options.values().next().value !== 'none'
+            ? normal_options.values().next().value
+            : 'open3d',
         render: (get) => get('mesh_method_options') === 'poisson',
       },
       numFaces: { label: 'Number of Faces', value: 50000, min: 1 },
@@ -162,8 +147,13 @@ export default function MeshSubPanel(props) {
       ` --output-dir ${controlValues.outputDir}` +
       ` --target-num-faces ${controlValues.numFaces}` +
       ` --num-pixels-per-side ${controlValues.textureResolution}` +
-      ` --normal-output-name ${controlValues.normal_options}` +
-      ` --num-points ${controlValues.numPoints}` +
+      ` --normal-method ${
+        controlValues.normal_options === 'open3d' ? 'open3d' : 'model_output'
+      }${
+        controlValues.normal_options !== 'open3d'
+          ? ` --normal-output-name ${controlValues.normal_options}`
+          : ''
+      } --num-points ${controlValues.numPoints}` +
       ` --remove-outliers ${getPythonBool(controlValues.removeOutliers)}` +
       ` --use-bounding-box ${getPythonBool(clippingEnabled)}` +
       ` --bounding-box-min ${bbox_min[0]} ${bbox_min[1]} ${bbox_min[2]}` +

nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx

@@ -3,8 +3,8 @@ import { useControls, useStoreContext } from 'leva';
 import { useEffect } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 import ContentCopyRoundedIcon from '@mui/icons-material/ContentCopyRounded';
-
 import { Button } from '@mui/material';
+import { get_normal_methods } from '../../../utils';
 
 export default function PointcloudSubPanel(props) {
   const update_box_center = props.update_box_center;
@@ -26,14 +26,25 @@ export default function PointcloudSubPanel(props) {
   const clippingScale = useSelector(
     (state) => state.renderingState.clipping_box_scale,
   );
-
+  const output_options = useSelector(
+    (state) => state.renderingState.output_options,
+  );
   const config_filename = `${config_base_dir}/config.yml`;
 
+  const normal_options = get_normal_methods(output_options);
+
   const controlValues = useControls(
     {
       numPoints: { label: 'Number of Points', value: 1000000, min: 1 },
       removeOutliers: { label: 'Remove Outliers', value: true },
-      estimateNormals: { label: 'Estimate Normals', value: false },
+      normal_options: {
+        label: 'Normal Options',
+        options: [...normal_options],
+        value:
+          normal_options.values().next().value !== 'none'
+            ? normal_options.values().next().value
+            : 'open3d',
+      },
       useBoundingBox: {
         label: 'Crop',
         value: clippingEnabled,
@@ -102,8 +113,14 @@ export default function PointcloudSubPanel(props) {
     ` --output-dir ${controlValues.outputDir}` +
     ` --num-points ${controlValues.numPoints}` +
     ` --remove-outliers ${getPythonBool(controlValues.removeOutliers)}` +
-    ` --estimate-normals ${getPythonBool(controlValues.estimateNormals)}` +
-    ` --use-bounding-box ${getPythonBool(clippingEnabled)}` +
+    ` --normal-method ${
+      controlValues.normal_options === 'open3d' ? 'open3d' : 'model_output'
+    }${
+      controlValues.normal_options !== 'open3d' &&
+      controlValues.normal_options !== 'none'
+        ? ` --normal-output-name ${controlValues.normal_options}`
+        : ''
+    } --use-bounding-box ${getPythonBool(clippingEnabled)}` +
     ` --bounding-box-min ${bbox_min[0]} ${bbox_min[1]} ${bbox_min[2]}` +
     ` --bounding-box-max ${bbox_max[0]} ${bbox_max[1]} ${bbox_max[2]}`;
 

nerfstudio/viewer/app/src/utils.js

@@ -1,3 +1,25 @@
 export function split_path(path_str) {
   return path_str.split('/').filter((x) => x.length > 0);
 }
+
+export function get_normal_outputs(output_options) {
+  // Get a list of normal outputs from the Model
+  let normal_options = [];
+  if (output_options) {
+    // check which outputs have normals
+    for (let i = 0; i < output_options.length; i += 1) {
+      const output_name = output_options[i];
+      if (output_name.includes('normals')) {
+        normal_options.push(output_name);
+      }
+    }
+  }
+  if (normal_options.length === 0) {
+    normal_options = ['none'];
+  }
+  return normal_options;
+}
+
+export function get_normal_methods(output_options) {
+  return new Set([...get_normal_outputs(output_options), 'open3d']);
+}