[WIP] Mesh export UI (#1076)

* starting PR

* export mesh ui functional

* better descriptions

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",

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 (
     <div className="ExportPanel">
@@ -90,13 +120,26 @@ export default function ExportPanel(props) {
         />
         <LevaStoreProvider store={pointcloudStore}>
           <PointcloudSubPanel
-            sceneTree={sceneTree}
-            showExportBox={showExportBox}
+            update_box_center={update_box_center}
+            update_box_scale={update_box_scale}
           />
         </LevaStoreProvider>
       </TabPanel>
       <TabPanel value={type_value} index={1}>
-        Coming Soon
+        <LevaPanel
+          store={meshStore}
+          className="Leva-panel"
+          theme={LevaTheme}
+          titleBar={false}
+          fill
+          flat
+        />
+        <LevaStoreProvider store={meshStore}>
+          <MeshSubPanel
+            update_box_center={update_box_center}
+            update_box_scale={update_box_scale}
+          />
+        </LevaStoreProvider>
       </TabPanel>
     </div>
   );

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 = (
+      <div className="ExportModal-text">
+        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 <br/>
+        <b>--pipeline.model.predict-normals True</b>
+      </div>
+    );
+  } else {
+    // if valid, provide the command needed for creating the mesh
+    inner_html = (
+      <>
+        <div className="ExportModal-text">
+          Run the following command in a terminal to export the mesh:
+        </div>
+        <div className="ExportModal-code">{cmd}</div>
+        <div className="ExportModal-button" style={{ textAlign: 'center' }}>
+          <Button
+            variant="outlined"
+            size="small"
+            startIcon={<ContentCopyRoundedIcon />}
+            onClick={handleCopy}
+          >
+            Copy Command
+          </Button>
+        </div>
+      </>
+    );
+  }
+
+  return <div className="MeshPanel">{inner_html}</div>;
+}

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 },
   );

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'."""