From 3377c5344b6c809a266eeaa499e535d985959c8f Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 2 Dec 2022 17:51:21 -0800 Subject: [PATCH] [WIP] Mesh export UI (#1076) * starting PR * export mesh ui functional * better descriptions --- nerfstudio/viewer/app/package.json | 2 +- .../SidePanel/ExportPanel/ExportPanel.jsx | 49 +++- .../SidePanel/ExportPanel/MeshSubPanel.jsx | 221 ++++++++++++++++++ .../ExportPanel/PointcloudSubPanel.jsx | 32 +-- scripts/exporter.py | 4 +- 5 files changed, 273 insertions(+), 35 deletions(-) create mode 100644 nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx diff --git a/nerfstudio/viewer/app/package.json b/nerfstudio/viewer/app/package.json index 400b5da17..b6183729e 100644 --- a/nerfstudio/viewer/app/package.json +++ b/nerfstudio/viewer/app/package.json @@ -1,7 +1,7 @@ { "name": "viewer", "homepage": ".", - "version": "22-11-28-0", + "version": "22-12-02-0", "private": true, "dependencies": { "@emotion/react": "^11.10.4", diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx index 9703e2c14..ac998daa2 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/ExportPanel.jsx @@ -1,11 +1,14 @@ /* eslint-disable react/jsx-props-no-spreading */ import * as React from 'react'; +import * as THREE from 'three'; +import { useSelector } from 'react-redux'; import { Box, Typography, Tab, Tabs } from '@mui/material'; import { LevaPanel, LevaStoreProvider, useCreateStore } from 'leva'; import BlurOnIcon from '@mui/icons-material/BlurOn'; import CategoryIcon from '@mui/icons-material/Category'; import PointcloudSubPanel from './PointcloudSubPanel'; +import MeshSubPanel from './MeshSubPanel'; import LevaTheme from '../../../themes/leva_theme.json'; interface TabPanelProps { @@ -41,6 +44,8 @@ function a11yProps(index: number) { }; } +const CLIPPING_BOX_NAME = 'Clipping Box'; + export default function ExportPanel(props) { // unpack relevant information const sceneTree = props.sceneTree; @@ -52,7 +57,32 @@ export default function ExportPanel(props) { setTypeValue(newValue); }; + // create the clipping box + const clippingEnabled = useSelector( + (state) => state.renderingState.clipping_enabled, + ); + const geometry = new THREE.BoxGeometry(1, 1, 1); + const material = new THREE.MeshBasicMaterial({ + color: 0xffd369, + transparent: true, + opacity: 0.3, + side: THREE.DoubleSide, + }); + const cube = new THREE.Mesh(geometry, material); + sceneTree.set_object_from_path([CLIPPING_BOX_NAME], cube); + cube.visible = showExportBox && clippingEnabled; + const update_box_center = (value) => { + const box = sceneTree.find_object_no_create([CLIPPING_BOX_NAME]); + box.position.set(value[0], value[1], value[2]); + }; + + const update_box_scale = (value) => { + const box = sceneTree.find_object_no_create([CLIPPING_BOX_NAME]); + box.scale.set(value[0], value[1], value[2]); + }; + const pointcloudStore = useCreateStore(); + const meshStore = useCreateStore(); return (
@@ -90,13 +120,26 @@ export default function ExportPanel(props) { /> - Coming Soon + + + +
); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx new file mode 100644 index 000000000..92e306be2 --- /dev/null +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/MeshSubPanel.jsx @@ -0,0 +1,221 @@ +import * as React from 'react'; +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; +} + +export default function MeshSubPanel(props) { + const update_box_center = props.update_box_center; + const update_box_scale = props.update_box_scale; + const store = useStoreContext(); + + const dispatch = useDispatch(); + + // redux store state + const config_base_dir = useSelector( + (state) => state.renderingState.config_base_dir, + ); + const clippingEnabled = useSelector( + (state) => state.renderingState.clipping_enabled, + ); + const clippingCenter = useSelector( + (state) => state.renderingState.clipping_center, + ); + const clippingScale = useSelector( + (state) => state.renderingState.clipping_box_scale, + ); + const output_options = useSelector( + (state) => state.renderingState.output_options, + ); + const normal_options = get_normal_outputs(output_options); + const mesh_method_options = ['poisson', 'tsdf']; + + const config_filename = `${config_base_dir}/config.yml`; + + const controlValues = useControls( + { + mesh_method_options: { + label: 'Mesh Method Options', + options: { + 'Poisson (best)': 'poisson', + TSDF: 'tsdf', + }, + value: mesh_method_options[0], + }, + normal_options: { + label: 'Normal Options', + options: [...new Set(normal_options)], + value: normal_options[0], + render: (get) => get('mesh_method_options') === 'poisson (best)', + }, + numFaces: { label: 'Number of Faces', value: 50000, min: 1 }, + textureResolution: { label: 'Texture Resolution', value: 2048, min: 1 }, + numPoints: { + label: 'Number of Points', + value: 1000000, + min: 1, + render: (get) => get('mesh_method_options') === 'poisson (best)', + }, + removeOutliers: { + label: 'Remove Outliers', + value: true, + render: (get) => get('mesh_method_options') === 'poisson (best)', + }, + useBoundingBox: { + label: 'Crop', + value: clippingEnabled, + onChange: (value) => { + dispatch({ + type: 'write', + path: 'renderingState/clipping_enabled', + data: value, + }); + }, + }, + center: { + label: '| Center', + value: clippingCenter, + render: (get) => get('useBoundingBox'), + onChange: (value) => { + update_box_center(value); + }, + onEditEnd: (value) => { + dispatch({ + type: 'write', + path: 'renderingState/clipping_center', + data: value, + }); + }, + }, + scale: { + label: '| Scale', + value: clippingScale, + min: 0.01, + render: (get) => get('useBoundingBox'), + onChange: (value) => { + update_box_scale(value); + }, + onEditEnd: (value) => { + dispatch({ + type: 'write', + path: 'renderingState/clipping_box_scale', + data: value, + }); + }, + }, + outputDir: { label: 'Output Directory', value: 'exports/mesh/' }, + }, + { store }, + ); + + const getPythonBool = (boolean) => { + return boolean ? 'True' : 'False'; + }; + + console.log(controlValues.outputDir); + console.log(getPythonBool(controlValues.removeOutliers)); + + const bbox_min = [ + clippingCenter[0] - clippingScale[0] / 2, + clippingCenter[1] - clippingScale[1] / 2, + clippingCenter[2] - clippingScale[2] / 2, + ]; + const bbox_max = [ + clippingCenter[0] + clippingScale[0] / 2, + clippingCenter[1] + clippingScale[1] / 2, + clippingCenter[2] + clippingScale[2] / 2, + ]; + + const mesh_method_choice = controlValues.mesh_method_options; + let cmd = ''; + if (mesh_method_choice === 'tsdf') { + cmd = + `ns-export ${mesh_method_choice}` + + ` --load-config ${config_filename}` + + ` --output-dir ${controlValues.outputDir}` + + ` --target-num-faces ${controlValues.numFaces}` + + ` --num-pixels-per-side ${controlValues.textureResolution}` + + ` --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]}`; + } else if (mesh_method_choice === 'poisson') { + cmd = + `ns-export ${mesh_method_choice}` + + ` --load-config ${config_filename}` + + ` --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}` + + ` --remove-outliers ${getPythonBool(controlValues.removeOutliers)}` + + ` --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]}`; + } + + const handleCopy = () => { + navigator.clipboard.writeText(cmd); + }; + + useEffect(() => { + update_box_center(clippingCenter); + update_box_scale(clippingScale); + }); + + let inner_html; + if ( + mesh_method_choice === 'poisson' && + controlValues.normal_options === 'none' + ) { + // check if normals aren't available when using poisson meshing + inner_html = ( +
+ You need to train a model with normals to use poisson meshing. Try tsdf + or retrain the model with normals. You can use nerfacto with
+ --pipeline.model.predict-normals True +
+ ); + } else { + // if valid, provide the command needed for creating the mesh + inner_html = ( + <> +
+ Run the following command in a terminal to export the mesh: +
+
{cmd}
+
+ +
+ + ); + } + + return
{inner_html}
; +} diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx index 8951c8051..480f3694c 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ExportPanel/PointcloudSubPanel.jsx @@ -1,5 +1,4 @@ import * as React from 'react'; -import * as THREE from 'three'; import { useControls, useStoreContext } from 'leva'; import { useEffect } from 'react'; import { useDispatch, useSelector } from 'react-redux'; @@ -7,11 +6,9 @@ import ContentCopyRoundedIcon from '@mui/icons-material/ContentCopyRounded'; import { Button } from '@mui/material'; -const PCD_CLIPPING_BOX_NAME = 'PCD Clipping Box'; - export default function PointcloudSubPanel(props) { - const sceneTree = props.sceneTree; - const showExportBox = props.showExportBox; + const update_box_center = props.update_box_center; + const update_box_scale = props.update_box_scale; const store = useStoreContext(); const dispatch = useDispatch(); @@ -32,27 +29,6 @@ export default function PointcloudSubPanel(props) { const config_filename = `${config_base_dir}/config.yml`; - const geometry = new THREE.BoxGeometry(1, 1, 1); - const material = new THREE.MeshBasicMaterial({ - color: 0xffd369, - transparent: true, - opacity: 0.7, - side: THREE.DoubleSide, - }); - const cube = new THREE.Mesh(geometry, material); - sceneTree.set_object_from_path([PCD_CLIPPING_BOX_NAME], cube); - cube.visible = showExportBox && clippingEnabled; - - const update_box_center = (value) => { - const box = sceneTree.find_object_no_create([PCD_CLIPPING_BOX_NAME]); - box.position.set(value[0], value[1], value[2]); - }; - - const update_box_scale = (value) => { - const box = sceneTree.find_object_no_create([PCD_CLIPPING_BOX_NAME]); - box.scale.set(value[0], value[1], value[2]); - }; - const controlValues = useControls( { numPoints: { label: 'Number of Points', value: 1000000, min: 1 }, @@ -62,8 +38,6 @@ export default function PointcloudSubPanel(props) { label: 'Crop', value: clippingEnabled, onChange: (value) => { - const box = sceneTree.find_object_no_create([PCD_CLIPPING_BOX_NAME]); - box.visible = value; dispatch({ type: 'write', path: 'renderingState/clipping_enabled', @@ -102,7 +76,7 @@ export default function PointcloudSubPanel(props) { }); }, }, - outputDir: { label: 'Output Directory', value: 'pcd/' }, + outputDir: { label: 'Output Directory', value: 'exports/pcd/' }, }, { store }, ); diff --git a/scripts/exporter.py b/scripts/exporter.py index a739b8814..658d5e191 100644 --- a/scripts/exporter.py +++ b/scripts/exporter.py @@ -58,7 +58,7 @@ class ExportPointCloud(Exporter): """Minimum of the bounding box, used if use_bounding_box is True.""" bounding_box_max: Tuple[float, float, float] = (1, 1, 1) """Minimum of the bounding box, used if use_bounding_box is True.""" - num_rays_per_batch: int = 65536 + num_rays_per_batch: int = 32768 """Number of rays to evaluate per batch. Decrease if you run out of memory.""" std_ratio: float = 10.0 """Threshold based on STD of the average distances across the point cloud to remove outliers.""" @@ -195,7 +195,7 @@ class ExportPoissonMesh(Exporter): """Minimum of the bounding box, used if use_bounding_box is True.""" bounding_box_max: Tuple[float, float, float] = (1, 1, 1) """Minimum of the bounding box, used if use_bounding_box is True.""" - num_rays_per_batch: int = 65536 + num_rays_per_batch: int = 32768 """Number of rays to evaluate per batch. Decrease if you run out of memory.""" texture_method: Literal["point_cloud", "nerf"] = "nerf" """Method to texture the mesh with. Either 'point_cloud' or 'nerf'."""