add render time in viewer (#1007)

* Add temporary render for viewer

* Add slider for render time

nerfstudio/viewer/app/package.json

@@ -1,7 +1,7 @@
 {
   "name": "viewer",
   "homepage": ".",
-  "version": "22-11-10-0",
+  "version": "22-11-28-0",
   "private": true,
   "dependencies": {
     "@emotion/react": "^11.10.4",

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

@@ -1,5 +1,5 @@
 import { buttonGroup, useControls } from 'leva';
-import { useContext, useEffect } from 'react';
+import { useContext, useEffect, useState } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
 import { WebSocketContext } from '../WebSocket/WebSocket';
@@ -27,6 +27,9 @@ export function RenderControls() {
   const target_train_util = useSelector(
     (state) => state.renderingState.targetTrainUtil,
   );
+  const render_time = useSelector(
+    (state) => state.renderingState.renderTime,
+  );
 
   const dispatch = useDispatch();
 
@@ -106,6 +109,36 @@ export function RenderControls() {
     }
   };
 
+  const [display_render_time, set_display_render_time] = useState(false);
+
+  const receive_temporal_dist = e => {
+      const msg = msgpack.decode(new Uint8Array(e.data));
+      if (msg.path === "/model/has_temporal_distortion") {
+        set_display_render_time(msg.data === "true");
+        websocket.removeEventListener("message", receive_temporal_dist);
+      }
+  }
+  websocket.addEventListener('message', receive_temporal_dist);
+
+  const set_render_time = (value) => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      const path = 'renderingState/render_time';
+      dispatch({
+        type: 'write',
+        path,
+        data: value,
+      });
+      const cmd = 'write';
+      const data = {
+        type: cmd,
+        path,
+        data: value,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
   const [, setControls] = useControls(
     () => ({
       // training speed
@@ -164,6 +197,17 @@ export function RenderControls() {
         '1024px': () => setControls({ max_resolution: 1024 }),
         '2048px': () => setControls({ max_resolution: 2048 }),
       }),
+      // Dynamic NeRF rendering time
+      ...(display_render_time ? {render_time: {
+        label: 'Render Timestep',
+        value: render_time,
+        min: 0,
+        max: 1,
+        step: 0.01,
+        onChange: (v) => {
+          set_render_time(v);
+        },
+      }} : {}),
     }),
     [
       outputOptions,
@@ -172,6 +216,8 @@ export function RenderControls() {
       colormapChoice,
       max_resolution,
       target_train_util,
+      render_time,
+      display_render_time,
       websocket, // need to re-render when websocket changes to use the new websocket
     ],
   );
@@ -188,6 +234,8 @@ export function RenderControls() {
     colormapChoice,
     max_resolution,
     target_train_util,
+    render_time,
+    display_render_time,
   ]);
 
   return null;

nerfstudio/viewer/app/src/reducer.js

@@ -34,6 +34,7 @@ const initialState = {
     train_eta: 'Paused',
     vis_train_ratio: 'Paused',
     log_errors: '',
+    renderTime: 0.0,
 
     // export options
     clipping_enabled: true,

nerfstudio/viewer/server/viewer_utils.py

@@ -325,6 +325,8 @@ class ViewerState:
         # set the initial state whether to train or not
         self.vis["renderingState/isTraining"].write(start_train)
 
+        self.vis["renderingState/render_time"].write(str(0))
+
         # set the properties of the camera
         # self.vis["renderingState/camera"].write(json_)
 
@@ -373,6 +375,9 @@ class ViewerState:
             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
 
@@ -700,6 +705,7 @@ class ViewerState:
 
     @profiler.time_function
     def _render_image_in_viewer(self, camera_object, graph: Model, is_training: bool) -> None:
+        # pylint: disable=too-many-statements
         """
         Draw an image using the current camera pose from the viewer.
         The image is sent of a TCP connection and then uses WebRTC to send it to the viewer.
@@ -750,12 +756,16 @@ class ViewerState:
             dim=0,
         )
 
+        times = float(self.vis["renderingState/render_time"].read())
+        times = torch.tensor([times])
+
         camera = Cameras(
             fx=intrinsics_matrix[0, 0],
             fy=intrinsics_matrix[1, 1],
             cx=intrinsics_matrix[0, 2],
             cy=intrinsics_matrix[1, 2],
             camera_to_worlds=camera_to_world[None, ...],
+            times=times,
         )
         camera = camera.to(graph.device)