From fef6b861be0fcc986165543d1bebf0ca880e2ab3 Mon Sep 17 00:00:00 2001
From: Rohan Mathur <rohanmathur@berkeley.edu>
Date: Thu, 19 Jan 2023 17:41:05 -0500
Subject: [PATCH] Camera path (#1251)

* added eta to base config

* added code to save camera paths to dated name in data folder

* updated camera path name customization

* removed unnecessary comments

* added new version to viewer package.json

* more clean up and also I left a print statement in

* fixed validation message
---
 nerfstudio/engine/trainer.py                  |  4 ++-
 nerfstudio/viewer/app/package.json            |  2 +-
 .../src/modules/RenderModal/RenderModal.jsx   | 30 ++++++++-----------
 .../SidePanel/CameraPanel/CameraPanel.jsx     |  9 +++---
 .../SidePanel/CameraPanel/CameraPropPanel.jsx | 26 ++++++++++++++++
 nerfstudio/viewer/app/src/reducer.js          |  3 ++
 nerfstudio/viewer/server/viewer_utils.py      | 26 ++++++++++++----
 7 files changed, 70 insertions(+), 30 deletions(-)

diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
index 30945a4c7..49879c31d 100644
--- a/nerfstudio/engine/trainer.py
+++ b/nerfstudio/engine/trainer.py
@@ -122,7 +122,9 @@ class Trainer:
         viewer_log_path = self.base_dir / config.viewer.relative_log_filename
         self.viewer_state, banner_messages = None, None
         if self.config.is_viewer_enabled() and local_rank == 0:
-            self.viewer_state, banner_messages = viewer_utils.setup_viewer(config.viewer, log_filename=viewer_log_path)
+            self.viewer_state, banner_messages = viewer_utils.setup_viewer(
+                config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data
+            )
         self._check_viewer_warnings()
         # set up writers/profilers if enabled
         writer_log_path = self.base_dir / config.logging.relative_log_dir
diff --git a/nerfstudio/viewer/app/package.json b/nerfstudio/viewer/app/package.json
index 78d4997f3..ccaa2fa1c 100644
--- a/nerfstudio/viewer/app/package.json
+++ b/nerfstudio/viewer/app/package.json
@@ -1,7 +1,7 @@
 {
   "name": "viewer",
   "homepage": ".",
-  "version": "22-12-16-0",
+  "version": "23-01-19-0",
   "private": true,
   "dependencies": {
     "@emotion/react": "^11.10.4",
diff --git a/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx b/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx
index 6305d4043..c9cea8f1b 100644
--- a/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx
+++ b/nerfstudio/viewer/app/src/modules/RenderModal/RenderModal.jsx
@@ -1,7 +1,7 @@
 /* eslint-disable react/jsx-props-no-spreading */
 import * as React from 'react';
 
-import { Box, Button, Modal, TextField, Typography } from '@mui/material';
+import { Box, Button, Modal, Typography } from '@mui/material';
 import { useSelector } from 'react-redux';
 import ContentCopyRoundedIcon from '@mui/icons-material/ContentCopyRounded';
 
@@ -19,15 +19,22 @@ export default function RenderModal(props: RenderModalProps) {
     (state) => state.renderingState.config_base_dir,
   );
 
+  const export_path = useSelector(
+    (state) => state.renderingState.export_path,
+  );
+
+  const data_base_dir = useSelector(
+    (state) => state.renderingState.data_base_dir,
+  );
+
   // react state
-  const [filename, setFilename] = React.useState('render_output');
 
   const handleClose = () => setOpen(false);
 
   // Copy the text inside the text field
   const config_filename = `${config_base_dir}/config.yml`;
-  const camera_path_filename = `${config_base_dir}/camera_path.json`;
-  const cmd = `ns-render --load-config ${config_filename} --traj filename --camera-path-filename ${camera_path_filename} --output-path renders/${filename}.mp4`;
+  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 text_intro = `To render a full resolution video, run the following command in a terminal.`;
 
@@ -56,22 +63,11 @@ export default function RenderModal(props: RenderModalProps) {
                 <br />
                 The video will be saved to{' '}
                 <code className="RenderModal-inline-code">
-                  ./renders/{filename}.mp4
+                  ./renders/{data_base_dir}/{export_path}.mp4
                 </code>
                 .
               </p>
-              <TextField
-                label="Output Name"
-                inputProps={{
-                  inputMode: 'numeric',
-                  pattern: '[+-]?([0-9]*[.])?[0-9]+',
-                }}
-                onChange={(e) => {
-                  setFilename(e.target.value);
-                }}
-                value={filename}
-                variant="standard"
-              />
+              
               <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
index 6dc11d30a..27fbf5e36 100644
--- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx
+++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx
@@ -556,9 +556,10 @@ export default function CameraPanel(props) {
   ]);
 
   // redux store state
-  const config_base_dir = useSelector(
-    (state) => state.renderingState.config_base_dir,
+  const export_path = useSelector(
+    (state) => state.renderingState.export_path,
   );
+
   const websocket = useContext(WebSocketContext).socket;
   const DEFAULT_FOV = 50;
   const DEFAULT_RENDER_TIME = 0.0;
@@ -1083,10 +1084,8 @@ export default function CameraPanel(props) {
     setRenderModalOpen(true);
 
     const camera_path_object = get_camera_path();
-    const camera_path_filename = `${config_base_dir}/camera_path.json`;
-
     const camera_path_payload = {
-      camera_path_filename,
+      camera_path_filename: export_path,
       camera_path: camera_path_object,
     };
 
diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPropPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPropPanel.jsx
index 354361573..84a5cfaed 100644
--- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPropPanel.jsx
+++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPropPanel.jsx
@@ -7,6 +7,7 @@ export default function CameraPropPanel(props) {
   const fps = props.fps;
   const set_fps = props.set_fps;
 
+  // redux store state
   const store = useStoreContext();
 
   const dispatch = useDispatch();
@@ -20,6 +21,16 @@ export default function CameraPropPanel(props) {
   );
   const camera_type = useSelector((state) => state.renderingState.camera_type);
 
+  const export_path = useSelector((state) => state.renderingState.export_path);
+
+  const setExportPath = (value) => {
+    dispatch({
+      type: 'write',
+      path: 'renderingState/export_path',
+      data: value,
+    });
+  };
+
   const setResolution = (value) => {
     dispatch({
       type: 'write',
@@ -43,6 +54,20 @@ export default function CameraPropPanel(props) {
 
   const [, setControls] = useControls(
     () => ({
+      path: {
+        label: 'Export Name',
+        value: export_path,
+        onChange: (v) => {
+          const valid_filename_reg = /^([a-z]|[A-Z]|[0-9]|-|_)+$/g;
+          if(!valid_filename_reg.test(v)){
+            alert("Please only use letters, numbers, and hyphens");
+          }
+          else {
+            setExportPath(v);
+          }
+        },
+        
+      },
       camera_resolution: {
         label: 'Resolution',
         value: { width: render_width, height: render_height },
@@ -84,6 +109,7 @@ export default function CameraPropPanel(props) {
     { store },
   );
 
+  setControls({path: export_path});
   setControls({ video_fps: fps });
   setControls({ video_duration: seconds });
   setControls({
diff --git a/nerfstudio/viewer/app/src/reducer.js b/nerfstudio/viewer/app/src/reducer.js
index e5a3c9907..b9ebe42a8 100644
--- a/nerfstudio/viewer/app/src/reducer.js
+++ b/nerfstudio/viewer/app/src/reducer.js
@@ -24,6 +24,9 @@ const initialState = {
     field_of_view: 50,
     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',
+
     isTraining: true,
     output_options: ['rgb'], // populated by the possible Graph outputs
     output_choice: 'rgb', // the selected output
diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py
index 9fa47aec4..d7907c476 100644
--- a/nerfstudio/viewer/server/viewer_utils.py
+++ b/nerfstudio/viewer/server/viewer_utils.py
@@ -19,6 +19,7 @@ from __future__ import annotations
 import asyncio
 import enum
 import os
+import re
 import sys
 import threading
 import time
@@ -68,13 +69,13 @@ def get_viewer_version() -> str:
 
 
 @check_main_thread
-def setup_viewer(config: cfg.ViewerConfig, log_filename: Path):
+def setup_viewer(config: cfg.ViewerConfig, log_filename: Path, datapath: str):
     """Sets up the viewer if enabled
 
     Args:
         config: the configuration to instantiate viewer
     """
-    viewer_state = ViewerState(config, log_filename=log_filename)
+    viewer_state = ViewerState(config, log_filename=log_filename, datapath=datapath)
     banner_messages = [f"Viewer at: {viewer_state.viewer_url}"]
     return viewer_state, banner_messages
 
@@ -225,11 +226,12 @@ class ViewerState:
         config: viewer setup configuration
     """
 
-    def __init__(self, config: cfg.ViewerConfig, log_filename: Path):
+    def __init__(self, config: cfg.ViewerConfig, log_filename: Path, datapath: str):
         self.config = config
         self.vis = None
         self.viewer_url = None
         self.log_filename = log_filename
+        self.datapath = datapath
         if self.config.launch_bridge_server:
             # start the viewer bridge server
             assert self.config.websocket_port is not None
@@ -306,6 +308,14 @@ class ViewerState:
         # set the config base dir
         self.vis["renderingState/config_base_dir"].write(str(self.log_filename.parents[0]))
 
+        # set the data base dir
+        self.vis["renderingState/data_base_dir"].write(str(self.datapath))
+
+        # get the timestamp of the train run to set default export path name
+        timestamp_reg = re.compile("[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{6}")
+        timestamp_match = timestamp_reg.findall(str(self.log_filename.parents[0]))
+        self.vis["renderingState/export_path"].write(timestamp_match[-1])
+
         # clear the current scene
         self.vis["sceneState/sceneBox"].delete()
         self.vis["sceneState/cameras"].delete()
@@ -341,10 +351,14 @@ class ViewerState:
         if camera_path_payload:
             # save a model checkpoint
             trainer.save_checkpoint(step)
-            # write to json file
-            camera_path_filename = camera_path_payload["camera_path_filename"]
+            # write to json file in datapath directory
+            camera_path_filename = camera_path_payload["camera_path_filename"] + ".json"
             camera_path = camera_path_payload["camera_path"]
-            write_to_json(Path(camera_path_filename), camera_path)
+            camera_paths_directory = os.path.join(self.datapath, "camera_paths")
+            if not os.path.exists(camera_paths_directory):
+                os.mkdir(camera_paths_directory)
+
+            write_to_json(Path(os.path.join(camera_paths_directory, camera_path_filename)), camera_path)
             self.vis["camera_path_payload"].delete()
 
     def _check_webrtc_offer(self):