added basic stats to viewer (#259)

nerfactory/utils/writer.py

@@ -116,6 +116,7 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True,
         remain_iter = GLOBAL_BUFFER["max_iter"] - step
         remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"]
         put_scalar(EventName.ETA, remain_time, step)
+        GLOBAL_BUFFER["viewer/train_eta"] = _format_time(remain_time)
 
 
 @check_main_thread
@@ -152,6 +153,7 @@ def setup_event_writers(config: LoggingConfig, max_iter: int) -> None:
     GLOBAL_BUFFER["max_buffer_size"] = config.max_buffer_size
     GLOBAL_BUFFER["steps_per_log"] = config.steps_per_log
     GLOBAL_BUFFER["events"] = {}
+    GLOBAL_BUFFER["viewer/train_eta"] = "Paused"
 
 
 class Writer:

nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx

@@ -31,6 +31,12 @@ export function RenderControls() {
   const field_of_view = useSelector(
     (state) => state.renderingState.field_of_view,
   );
+  let eval_fps = useSelector(
+    (state) => state.renderingState.eval_fps,
+  );
+  let train_eta = useSelector(
+    (state) => state.renderingState.train_eta,
+  );
 
   const dispatch = useDispatch();
 
@@ -111,18 +117,8 @@ export function RenderControls() {
   };
 
   const [, setControls] = useControls(
-    'Rendering State',
+    'Rendering Controls',
     () => ({
-      // WebSocket isConnected
-      // button does nothing except be an indicator
-      'WebSocket Connected': button(() => {}, {
-        disabled: !isWebsocketConnected,
-      }),
-      // webRTC isConnected
-      // button does nothing except be an indicator
-      'WebRTC Connected': button(() => {}, {
-        disabled: !isWebrtcConnected,
-      }),
       // isTraining
       'Pause Training': button(
         () => {
@@ -220,8 +216,6 @@ export function RenderControls() {
       },
     }),
     [
-      isWebsocketConnected,
-      isWebrtcConnected,
       isTraining,
       outputOptions,
       outputChoice,
@@ -231,6 +225,38 @@ export function RenderControls() {
     ],
   );
 
+  const [, setState] = useControls(
+    'Rendering Monitor',
+    () => ({
+      // WebSocket isConnected
+      // button does nothing except be an indicator
+      'WebSocket Connected': button(() => {}, {
+        disabled: !isWebsocketConnected,
+      }),
+      // webRTC isConnected
+      // button does nothing except be an indicator
+      'WebRTC Connected': button(() => {}, {
+        disabled: !isWebrtcConnected,
+      }),
+      eval_fps: {
+        label: 'Eval FPS',
+        value: eval_fps,
+        disabled: true,
+      },
+      train_eta: {
+        label: 'Train ETA',
+        value: train_eta,
+        disabled: true,
+      },
+    }),
+    [
+      isWebsocketConnected,
+      isWebrtcConnected,
+      eval_fps,
+      train_eta,
+    ],
+  );
+
   useEffect(() => {
     setControls({ min_resolution });
     setControls({ max_resolution });
@@ -238,8 +264,6 @@ export function RenderControls() {
     setControls({ 'Camera FoV': field_of_view });
   }, [
     setControls,
-    isWebsocketConnected,
-    isWebrtcConnected,
     isTraining,
     outputOptions,
     outputChoice,
@@ -248,5 +272,28 @@ export function RenderControls() {
     field_of_view,
   ]);
 
+  useEffect(() => {}, [
+    setState,
+    isWebsocketConnected,
+    isWebrtcConnected,
+  ]);
+
+  useEffect(() => {
+    websocket.addEventListener('message', (originalCmd) => {
+      // set the remote description when the offer is received
+      const cmd = msgpack.decode(new Uint8Array(originalCmd.data));
+      if (cmd.path === '/renderingState/eval_fps') {
+        eval_fps = cmd.data;
+        console.log("setting state", eval_fps);
+        setState({ eval_fps });
+      }
+      if (cmd.path === '/renderingState/train_eta') {
+        train_eta = cmd.data;
+        console.log("setting state", train_eta)
+        setState({ train_eta });
+      }
+    });
+  }, []);
+
   return null;
 }

nerfactory/viewer/app/src/reducer.js

@@ -18,6 +18,8 @@ const initialState = {
     minResolution: 50,
     maxResolution: 512,
     field_of_view: 80,
+    eval_fps: 'Paused',
+    train_eta: 'Paused',
   },
   // the scene state
   sceneState: {

nerfactory/viewer/server/viewer_utils.py

@@ -32,6 +32,7 @@ from nerfactory.graphs.base import Graph
 from nerfactory.utils import profiler
 from nerfactory.utils.config import ViewerConfig
 from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all
+from nerfactory.utils.writer import GLOBAL_BUFFER
 from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
 from nerfactory.viewer.server.visualizer import Viewer
 
@@ -279,6 +280,16 @@ class VisualizerState:
         image = (image_output).astype("uint8")
         self.vis.set_image(image)
 
+    def _update_viewer_stats(self, render_time, image_height) -> None:
+        is_training = self.vis["renderingState/isTraining"].read()
+        if is_training is None or is_training:
+            eval_fps = f"{1 / render_time:.2f} fps at {image_height} res"
+            self.vis["renderingState/eval_fps"].write(eval_fps)
+            self.vis["renderingState/train_eta"].write(GLOBAL_BUFFER["viewer/train_eta"])
+        else:
+            self.vis["renderingState/eval_fps"].write("Paused")
+            self.vis["renderingState/train_eta"].write("Paused")
+
     @profiler.time_function
     def _render_image_in_viewer(self, graph: Graph) -> None:
         """
@@ -336,21 +347,25 @@ class VisualizerState:
         camera_ray_bundle.num_rays_per_chunk = self.config.num_rays_per_chunk
 
         graph.eval()
+
         check_thread = CheckThread(state=self)
         render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle)
+        start_time = time.time()
         check_thread.start()
         render_thread.start()
-
         try:
             render_thread.join()
             check_thread.join()
         except Exception:  # pylint: disable=broad-except
             pass
+        render_duration = time.time() - start_time
 
         graph.train()
+
         outputs = render_thread.vis_outputs
         if outputs is not None:
             self._send_output_to_viewer(outputs)
+            self._update_viewer_stats(render_duration, image_height)
 
 
 def get_default_vis() -> Viewer: