diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index d8c8a0ac9..153ab4e4b 100644 --- a/nerfstudio/scripts/exporter.py +++ b/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) diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx index c322ce971..58c42a263 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx +++ b/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]}` + diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx index 9a8f7177e..76e99027d 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx +++ b/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]}`; diff --git a/nerfstudio/viewer/app/src/utils.js b/nerfstudio/viewer/app/src/utils.js index b63d54361..9422a9988 100644 --- a/nerfstudio/viewer/app/src/utils.js +++ b/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']); +}