Load path dropdown (nerfstudio-project#1290)

* started load path options * deleted random comment * got the camera paths to all send but it won't render right * finished first draft of new load path button * removed unnecessary prints * fixed minor name warning * clean up and ui change * Clean up UI Co-authored-by: Matthew Tancik <[email protected]> Co-authored-by: Matt Tancik <[email protected]>
DominikVincent · Jan 27, 2023 · dcadbbc · dcadbbc
1 parent b5031dd
commit dcadbbc
Show file tree

Hide file tree

Showing 9 changed files with 228 additions and 17 deletions.
diff --git a/nerfstudio/viewer/app/src/index.scss b/nerfstudio/viewer/app/src/index.scss
@@ -242,6 +242,25 @@ $sidebar-width: 100% !default;
   margin-top: 20px;
 }
 
+// ------------------------------------
+// Load Path Modal styling
+
+.LoadPathModal-box {
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  width: 500px; // 400
+  background-color: $base-color2;
+  border: 2px solid #000;
+  padding: 32px;
+  margin-top: 20px;
+}
+
+.LoadPathModal-upload_button {
+  text-align: center;
+}
+
 // ------------------------------------
 // Camera Pane
 

diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx
@@ -47,7 +47,6 @@ export function RenderControls() {
   const crop_enabled = useSelector(
     (state) => state.renderingState.crop_enabled,
   );
-  console.log('crop_enabled_config panel', crop_enabled);
 
   const crop_bg_color = useSelector(
     (state) => state.renderingState.crop_bg_color,
@@ -162,7 +161,6 @@ export function RenderControls() {
         value: crop_enabled,
         hint: 'Crop the viewport to the selected box',
         onChange: (value) => {
-          console.log('dispact value', value);
           dispatch_and_send(
             websocket,
             dispatch,

diff --git a/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx b/nerfstudio/viewer/app/src/modules/LoadPathModal/LoadPathModal.jsx
@@ -0,0 +1,149 @@
+/* eslint-disable react/jsx-props-no-spreading */
+import * as React from 'react';
+
+import { Box, Button, FormControl, Modal, Typography } from '@mui/material';
+import InputLabel from '@mui/material/InputLabel';
+import { useSelector } from 'react-redux';
+import { FileUpload } from '@mui/icons-material';
+import MenuItem from '@mui/material/MenuItem';
+import Select from '@mui/material/Select';
+
+interface LoadPathModalProps {
+  open: object;
+  setOpen: object;
+  pathUploadFunction: any;
+  loadCameraPathFunction: any;
+}
+
+export default function LoadPathModal(props: LoadPathModalProps) {
+  const open = props.open;
+  const setOpen = props.setOpen;
+  const uploadCameraPath = props.pathUploadFunction;
+  const loadCameraPath = props.loadCameraPathFunction;
+
+  const [existingPathSelect, setExistingPathSelect] = React.useState('');
+
+  // redux store state
+  const all_camera_paths = useSelector(
+    (state) => state.renderingState.all_camera_paths,
+  );
+
+  let camera_paths_arr = [];
+  if (typeof all_camera_paths === 'object' && all_camera_paths !== null) {
+    camera_paths_arr = Object.keys(all_camera_paths).map((key) => {
+      return {
+        name: key,
+        val: all_camera_paths[key],
+      };
+    });
+  }
+
+  const hiddenFileInput = React.useRef(null);
+  const handleFileUploadClick = () => {
+    hiddenFileInput.current.click();
+  };
+
+  // react state
+
+  const handleClose = () => setOpen(false);
+  const handlePathSelect = (event) => {
+    setExistingPathSelect(event.target.value);
+  };
+
+  const handleExistingLoadClick = () => {
+    loadCameraPath(existingPathSelect);
+    handleClose();
+  };
+
+  const handleFileInput = (event) => {
+    uploadCameraPath(event);
+    handleClose();
+  };
+
+  return (
+    <div className="LoadPathModal">
+      <Modal
+        open={open}
+        onClose={handleClose}
+        aria-labelledby="modal-modal-title"
+        aria-describedby="modal-modal-description"
+      >
+        <Box className="LoadPathModal-box">
+          <Typography
+            id="modal-modal-description"
+            component="div"
+            sx={{ mt: 2 }}
+          >
+            <div>
+              <h2>Load Camera Path</h2>
+              {camera_paths_arr.length > 0 && (
+                <>
+                  <p>
+                    Either upload a local file or select a saved camera path
+                  </p>
+                  <FormControl
+                    sx={{ minWidth: '100%' }}
+                    variant="filled"
+                    size="small"
+                  >
+                    <InputLabel id="ageInputLabel">Existing Path</InputLabel>
+                    <Select
+                      labelId="ageInputLabel"
+                      label="Camera Path"
+                      value={existingPathSelect}
+                      onChange={handlePathSelect}
+                    >
+                      {camera_paths_arr.map((obj) => {
+                        return <MenuItem value={obj.val}>{obj.name}</MenuItem>;
+                      })}
+                    </Select>
+                    <Button
+                      sx={{
+                        marginTop: '10px',
+                        marginLeft: 'auto',
+                        marginRight: 'auto',
+                        width: '60%',
+                      }}
+                      variant="outlined"
+                      size="medium"
+                      disabled={existingPathSelect === ''}
+                      onClick={handleExistingLoadClick}
+                    >
+                      Load
+                    </Button>
+                  </FormControl>
+                  <br />
+                  <p>
+                    <center>OR</center>
+                  </p>
+                </>
+              )}
+              {camera_paths_arr.length === 0 && (
+                <p>No existing saved paths found</p>
+              )}
+              <div className="LoadPathModal-upload_button">
+                <Button
+                  sx={{ width: '60%' }}
+                  variant="outlined"
+                  size="medium"
+                  startIcon={<FileUpload />}
+                  onClick={handleFileUploadClick}
+                >
+                  Upload Camera Path
+                  <input
+                    type="file"
+                    accept=".json"
+                    name="Camera Path"
+                    onChange={handleFileInput}
+                    hidden
+                    ref={hiddenFileInput}
+                  />
+                </Button>
+              </div>
+            </div>
+          </Typography>
+        </Box>
+      </Modal>
+    </div>
+  );
+}
diff --git a/nerfstudio/viewer/app/src/modules/LoadPathModal/index.jsx b/nerfstudio/viewer/app/src/modules/LoadPathModal/index.jsx
@@ -0,0 +1,3 @@
+import LoadPathModal from './LoadPathModal';
+
+export default LoadPathModal;
diff --git a/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx b/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx
@@ -19,9 +19,7 @@ export default function RenderModal(props: RenderModalProps) {
     (state) => state.renderingState.config_base_dir,
   );
 
-  const export_path = useSelector(
-    (state) => state.renderingState.export_path,
-  );
+  const export_path = useSelector((state) => state.renderingState.export_path);
 
   const data_base_dir = useSelector(
     (state) => state.renderingState.data_base_dir,
@@ -34,12 +32,14 @@ export default function RenderModal(props: RenderModalProps) {
   // Copy the text inside the text field
   const config_filename = `${config_base_dir}/config.yml`;
   const camera_path_filename = `${export_path}.json`;
-  const cmd = `ns-render --load-config ${config_filename} --traj filename --camera-path-filename ${data_base_dir}/camera_paths/${camera_path_filename} --output-path renders/${data_base_dir}/${export_path}.mp4`;
+  const data_base_dir_leaf = data_base_dir.split('/').pop();
+  const cmd = `ns-render --load-config ${config_filename} --traj filename --camera-path-filename ${data_base_dir}/camera_paths/${camera_path_filename} --output-path renders/${data_base_dir_leaf}/${export_path}.mp4`;
 
   const text_intro = `To render a full resolution video, run the following command in a terminal.`;
 
   const handleCopy = () => {
     navigator.clipboard.writeText(cmd);
+    handleClose();
   };
 
   return (
@@ -63,11 +63,11 @@ export default function RenderModal(props: RenderModalProps) {
                 <br />
                 The video will be saved to{' '}
                 <code className="RenderModal-inline-code">
-                  ./renders/{data_base_dir}/{export_path}.mp4
+                  ./renders/{data_base_dir_leaf}/{export_path}.mp4
                 </code>
                 .
               </p>
-              
+
               <div className="RenderModal-code">{cmd}</div>
               <div style={{ textAlign: 'center' }}>
                 <Button

diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx
@@ -47,6 +47,7 @@ import { CameraHelper } from './CameraHelper';
 import { get_curve_object_from_cameras, get_transform_matrix } from './curve';
 import { WebSocketContext } from '../../WebSocket/WebSocket';
 import RenderModal from '../../RenderModal';
+import LoadPathModal from '../../LoadPathModal';
 import CameraPropPanel from './CameraPropPanel';
 import LevaTheme from '../../../themes/leva_theme.json';
 
@@ -572,6 +573,7 @@ export default function CameraPanel(props) {
   const [seconds, setSeconds] = React.useState(4);
   const [fps, setFps] = React.useState(24);
   const [render_modal_open, setRenderModalOpen] = React.useState(false);
+  const [load_path_modal_open, setLoadPathModalOpen] = React.useState(false);
   const [animate, setAnimate] = React.useState(new Set());
   const [globalFov, setGlobalFov] = React.useState(DEFAULT_FOV);
   const [globalRenderTime, setGlobalRenderTime] =
@@ -1160,6 +1162,19 @@ export default function CameraPanel(props) {
     }
   };
 
+  const open_load_path_modal = () => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      const data = {
+        type: 'write',
+        path: 'populate_paths_payload',
+        data: true,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+    setLoadPathModalOpen(true);
+  };
+
   const isAnimated = (property) => animate.has(property);
 
   const toggleAnimate = (property) => {
@@ -1195,21 +1210,21 @@ export default function CameraPanel(props) {
     <div className="CameraPanel">
       <div>
         <div className="CameraPanel-path-row">
+          <LoadPathModal
+            open={load_path_modal_open}
+            setOpen={setLoadPathModalOpen}
+            pathUploadFunction={uploadCameraPath}
+            loadCameraPathFunction={load_camera_path}
+          />
           <Button
             size="small"
             className="CameraPanel-top-button"
             component="label"
             variant="outlined"
             startIcon={<FileUploadOutlinedIcon />}
+            onClick={open_load_path_modal}
           >
             Load Path
-            <input
-              type="file"
-              accept=".json"
-              name="Camera Path"
-              onChange={uploadCameraPath}
-              hidden
-            />
           </Button>
         </div>
         <div className="CameraPanel-path-row">

diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js
@@ -12,6 +12,7 @@ const initialState = {
   },
   // for sending actual commands to the client
   camera_path_payload: null,
+  populate_paths_payload: false,
   // the rendering state
   renderingState: {
     // cameras
@@ -25,7 +26,10 @@ const initialState = {
     camera_type: 'perspective',
 
     data_base_dir: 'data_base_dir', // the base directory of the images for saving camera path with the data
-    export_path: 'export_path',
+    export_path: 'export_path', // export name for render and camera_path
+
+    all_camera_paths: null, // object containing camera paths and names
+
 
     isTraining: true,
     output_options: ['rgb'], // populated by the possible Graph outputs

diff --git a/nerfstudio/viewer/app/src/themes/theme.ts b/nerfstudio/viewer/app/src/themes/theme.ts
@@ -169,5 +169,12 @@ export const appTheme = createTheme({
         },
       },
     },
+    MuiSelect: {
+      styleOverrides: {
+        icon: {
+          color: '#eeeeee',
+        },
+      },
+    },
   },
 });
diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py
@@ -388,6 +388,21 @@ def _check_camera_path_payload(self, trainer, step: int):
             write_to_json(Path(os.path.join(camera_paths_directory, camera_path_filename)), camera_path)
             self.vis["camera_path_payload"].delete()
 
+    def _check_populate_paths_payload(self, trainer, step: int):
+        populate_paths_payload = self.vis["populate_paths_payload"].read()
+        if populate_paths_payload:
+            # save a model checkpoint
+            trainer.save_checkpoint(step)
+            # get all camera paths
+            camera_path_dir = os.path.join(self.datapath, "camera_paths")
+            camera_path_files = os.listdir(camera_path_dir)
+            all_path_dict = {}
+            for i in camera_path_files:
+                if i[-4:] == "json":
+                    all_path_dict[i[:-5]] = load_from_json(Path(os.path.join(camera_path_dir, i)))
+            self.vis["renderingState/all_camera_paths"].write(all_path_dict)
+            self.vis["populate_paths_payload"].delete()
+
     def _check_webrtc_offer(self):
         """Check if there is a webrtc offer to respond to."""
         data = self.vis["webrtc/offer"].read()
@@ -463,14 +478,14 @@ def update_scene(self, trainer, step: int, graph: Model, num_rays_per_batch: int
             step: iteration step of training
             graph: the current checkpoint of the model
         """
-
         has_temporal_distortion = getattr(graph, "temporal_distortion", None) is not None
         self.vis["model/has_temporal_distortion"].write(str(has_temporal_distortion).lower())
 
         is_training = self.vis["renderingState/isTraining"].read()
         self.step = step
 
         self._check_camera_path_payload(trainer, step)
+        self._check_populate_paths_payload(trainer, step)
         self._check_webrtc_offer()
 
         camera_object = self._get_camera_object()
@@ -515,6 +530,7 @@ def update_scene(self, trainer, step: int, graph: Model, num_rays_per_batch: int
                     self._render_image_in_viewer(camera_object, graph, is_training)
                     camera_object = self._get_camera_object()
                 is_training = self.vis["renderingState/isTraining"].read()
+                self._check_populate_paths_payload(trainer, step)
                 self._check_camera_path_payload(trainer, step)
                 self._check_webrtc_offer()
                 run_loop = not is_training
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		import LoadPathModal from './LoadPathModal';

		export default LoadPathModal;