Dynamic resolution in viewer (#397)

Dynamic resolution

nerfactory/engine/trainer.py

@@ -137,7 +137,8 @@ class Trainer:
                         callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)
 
                     with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t:
-                        self.visualizer_state.update_scene(step, self.pipeline.model)
+                        num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
+                        self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
 
                 if step != 0 and step % self.config.logging.steps_per_log == 0:
                     writer.put_dict(name="Loss/train-loss_metrics_dict", scalar_dict=loss_metric_dict, step=step)

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

@@ -21,37 +21,18 @@ export function RenderControls() {
   const colormapChoice = useSelector(
     (state) => state.renderingState.colormap_choice,
   );
-  const min_resolution = useSelector(
-    (state) => state.renderingState.minResolution,
-  );
   const max_resolution = useSelector(
     (state) => state.renderingState.maxResolution,
   );
+  const target_train_util = useSelector(
+    (state) => state.renderingState.targetTrainUtil,
+  );
   const field_of_view = useSelector(
     (state) => state.renderingState.field_of_view,
   );
 
   const dispatch = useDispatch();
 
-  const set_min_resolution = (value) => {
-    if (websocket.readyState === WebSocket.OPEN) {
-      dispatch({
-        type: 'write',
-        path: 'renderingState/minResolution',
-        data: value,
-      });
-      const cmd = 'write';
-      const path = 'renderingState/minResolution';
-      const data = {
-        type: cmd,
-        path,
-        data: value,
-      };
-      const message = msgpack.encode(data);
-      websocket.send(message);
-    }
-  };
-
   const set_max_resolution = (value) => {
     if (websocket.readyState === WebSocket.OPEN) {
       dispatch({
@@ -71,6 +52,25 @@ export function RenderControls() {
     }
   };
 
+  const set_target_train_util = (value) => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      dispatch({
+        type: 'write',
+        path: 'renderingState/targetTrainUtil',
+        data: value,
+      });
+      const cmd = 'write';
+      const path = 'renderingState/targetTrainUtil';
+      const data = {
+        type: cmd,
+        path,
+        data: value,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
   const set_output_choice = (value) => {
     if (websocket.readyState === WebSocket.OPEN) {
       dispatch({
@@ -130,6 +130,15 @@ export function RenderControls() {
 
   const [, setControls] = useControls(
     () => ({
+      // training speed
+      'Train Speed': buttonGroup({
+        Fast: () =>
+          setControls({ target_train_util: 0.9, max_resolution: 512 }),
+        Balanced: () =>
+          setControls({ target_train_util: 0.7, max_resolution: 1024 }),
+        Slow: () =>
+          setControls({ target_train_util: 0.1, max_resolution: 2048 }),
+      }),
       // output_options
       output_options: {
         label: 'Output Render',
@@ -149,27 +158,29 @@ export function RenderControls() {
         },
         disabled: colormapOptions.length === 1,
       },
-      // resolution
-      min_resolution: {
-        label: 'Min Res.',
-        value: min_resolution,
-        min: 10,
-        max: 256,
-        step: 1,
+      // FOV
+      'Camera FoV': {
+        value: field_of_view,
         onChange: (v) => {
-          set_min_resolution(v);
+          set_fov(v);
         },
       },
-      ' ': buttonGroup({
-        '32px': () => setControls({ min_resolution: 32 }),
-        '64px': () => setControls({ min_resolution: 64 }),
-        '128px': () => setControls({ min_resolution: 128 }),
-        '256px': () => setControls({ min_resolution: 256 }),
-      }),
+      // Dynamic Resolution
+      target_train_util: {
+        label: 'Train Util.',
+        value: target_train_util,
+        min: 0,
+        max: 1,
+        step: 0.05,
+        onChange: (v) => {
+          set_target_train_util(v);
+        },
+      },
+      // resolution
       max_resolution: {
         label: 'Max Res.',
         value: max_resolution,
-        min: 10,
+        min: 256,
         max: 2048,
         step: 1,
         onChange: (v) => {
@@ -177,23 +188,10 @@ export function RenderControls() {
         },
       },
       '  ': buttonGroup({
-        '128px': () => setControls({ max_resolution: 128 }),
         '256px': () => setControls({ max_resolution: 256 }),
         '512px': () => setControls({ max_resolution: 512 }),
         '1024px': () => setControls({ max_resolution: 1024 }),
-      }),
-      'Camera FoV': {
-        value: field_of_view,
-        onChange: (v) => {
-          set_fov(v);
-        },
-      },
-      // training speed
-      'Train Speed': buttonGroup({
-        Fast: () => setControls({ min_resolution: 10, max_resolution: 10 }),
-        Balanced: () =>
-          setControls({ min_resolution: 128, max_resolution: 512 }),
-        Slow: () => setControls({ min_resolution: 258, max_resolution: 1024 }),
+        '2048px': () => setControls({ max_resolution: 2048 }),
       }),
     }),
     [
@@ -201,15 +199,14 @@ export function RenderControls() {
       outputChoice,
       colormapOptions,
       colormapChoice,
-      min_resolution,
       max_resolution,
+      target_train_util,
       field_of_view,
       websocket, // need to re-render when websocket changes to use the new websocket
     ],
   );
 
   useEffect(() => {
-    setControls({ min_resolution });
     setControls({ max_resolution });
     setControls({ output_options: outputChoice });
     setControls({ colormap_options: colormapChoice });
@@ -220,8 +217,8 @@ export function RenderControls() {
     outputChoice,
     colormapOptions,
     colormapChoice,
-    min_resolution,
     max_resolution,
+    target_train_util,
     field_of_view,
   ]);
 

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

@@ -118,7 +118,7 @@ interface BasicTabsProps {
 export function BasicTabs(props: BasicTabsProps) {
   const sceneTree = props.sceneTree;
 
-  const [value, setValue] = React.useState(2);
+  const [value, setValue] = React.useState(0);
 
   const handleChange = (event: React.SyntheticEvent, newValue: number) => {
     setValue(newValue);

nerfactory/viewer/app/src/reducer.js

@@ -18,8 +18,8 @@ const initialState = {
     output_choice: 'rgb', // the selected output
     colormap_options: ['default'], // populated by the output choice
     colormap_choice: 'default', // the selected colormap
-    minResolution: 128,
-    maxResolution: 512,
+    maxResolution: 1024,
+    targetTrainUtil: 0.9,
     field_of_view: 80,
     eval_fps: 'Paused',
     train_eta: 'Paused',

nerfactory/viewer/server/viewer_utils.py

@@ -19,12 +19,11 @@ import enum
 import logging
 import sys
 import threading
-import time
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import numpy as np
 import torch
-from rich import print  # pylint: disable=redefined-builtin
+from rich.console import Console
 
 from nerfactory.cameras.cameras import Cameras
 from nerfactory.cameras.rays import RayBundle
@@ -38,6 +37,8 @@ from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subp
 from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
 from nerfactory.viewer.server.visualizer import Viewer
 
+console = Console(width=120)
+
 
 class OutputTypes(str, enum.Enum):
     """Noncomprehsnive list of output render types"""
@@ -140,11 +141,14 @@ class CheckThread(threading.Thread):
             data = self.state.vis["renderingState/camera"].read()
             if data is not None:
                 camera_object = data["object"]
-                if self.state.prev_camera_matrix is None or not np.allclose(
-                    camera_object["matrix"], self.state.prev_camera_matrix
+                if self.state.prev_camera_matrix is None or (
+                    not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix)
+                    and not self.state.prev_moving
                 ):
                     self.state.check_interrupt_vis = True
+                    self.state.prev_moving = True
                     return
+                self.state.prev_moving = False
 
             # check output type
             output_type = self.state.vis["renderingState/output_choice"].read()
@@ -169,13 +173,6 @@ class CheckThread(threading.Thread):
                     self.state.check_interrupt_vis = True
                     return
 
-            # check min render
-            min_resolution = self.state.vis["renderingState/minResolution"].read()
-            if min_resolution is not None:
-                if self.state.min_resolution != min_resolution:
-                    self.state.check_interrupt_vis = True
-                    return
-
 
 @decorate_all([check_visualizer_enabled])
 class VisualizerState:
@@ -202,18 +199,18 @@ class VisualizerState:
                 # TODO(ethan): move this into the writer such that it's at the bottom
                 # of the logging stack and easy to see and click
                 # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go
-                print("\n")
+                console.line()
                 self.viewer_url = (
                     f"https://viewer.nerfactory.com/branch/master/?websocket_url=localhost:{websocket_port}"
                 )
                 viewer_url_local = f"http://localhost:4000/?websocket_url=localhost:{websocket_port}"
                 pub_open_viewer_instructions_string = f"[Public] Open the viewer at {self.viewer_url}"
                 dev_open_viewer_instructions_string = f"[Local] Open the viewer at {viewer_url_local}"
-                print("-" * len(pub_open_viewer_instructions_string))
-                print(pub_open_viewer_instructions_string)
-                print(dev_open_viewer_instructions_string)
-                print("-" * len(pub_open_viewer_instructions_string))
-                print("\n")
+                console.rule(characters="=")
+                console.print(pub_open_viewer_instructions_string)
+                console.print(dev_open_viewer_instructions_string)
+                console.rule(characters="=")
+                console.line()
             self.vis = Viewer(zmq_url=self.config.zmq_url)
         else:
             logging.info("Continuing without viewer.")
@@ -222,15 +219,15 @@ class VisualizerState:
         self.prev_camera_matrix = None
         self.prev_output_type = OutputTypes.INIT
         self.prev_colormap_type = ColormapTypes.INIT
+        self.prev_moving = False
         self.output_type_changed = True
-        self.min_resolution = 50
         self.max_resolution = 1000
-        self.res_upscale_factor = 1
         self.check_interrupt_vis = False
         self.check_done_render = True
-        self.last_render_time = time.time()
-        self.min_wait_time = 0.5  # 1.0 is on high side and will cause lag
         self.step = 0
+        self.static_fps = 1
+        self.moving_fps = 24
+        self.camera_moving = False
 
         self.outputs_set = False
 
@@ -268,7 +265,7 @@ class VisualizerState:
         # K = camera.get_intrinsics_matrix()
         # set_persp_intrinsics_matrix(self.vis, K.double().numpy())
 
-    def update_scene(self, step: int, graph: Model) -> None:
+    def update_scene(self, step: int, graph: Model, num_rays_per_batch: int) -> None:
         """updates the scene based on the graph weights
 
         Args:
@@ -279,17 +276,46 @@ class VisualizerState:
         is_training = self.vis["renderingState/isTraining"].read()
         self.step = step
 
+        camera_object = self._get_camera_object()
+        if camera_object is None:
+            return
+
         if is_training is None or is_training:
-            # in training mode, render every few steps
-            if self._is_render_step(step):
-                self._render_image_in_viewer(graph)
+            # in training mode
+
+            if self.camera_moving:
+                # if the camera is moving, then we pause training and update camera continuously
+
+                while self.camera_moving:
+                    self._render_image_in_viewer(camera_object, graph, is_training)
+                    camera_object = self._get_camera_object()
+            else:
+                # if the camera is not moving, then we approximate how many training steps need to be taken
+                # to render at a FPS defined by self.static_fps.
+
+                if EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]:
+                    train_rays_per_sec = GLOBAL_BUFFER["events"][EventName.TRAIN_RAYS_PER_SEC.value]["avg"]
+                    target_train_util = self.vis["renderingState/targetTrainUtil"].read()
+                    if target_train_util is None:
+                        target_train_util = 0.9
+
+                    batches_per_sec = train_rays_per_sec / num_rays_per_batch
+
+                    num_steps = int(1 / self.static_fps * batches_per_sec)
+                else:
+                    num_steps = 1
+
+                if step % num_steps == 0:
+                    self._render_image_in_viewer(camera_object, graph, is_training)
+
         else:
             # in pause training mode, enter render loop with set graph
             local_step = step
             run_loop = not is_training
             while run_loop:
-                if self._is_render_step(local_step) and step > 0:
-                    self._render_image_in_viewer(graph)
+                # if self._is_render_step(local_step) and step > 0:
+                if step > 0:
+                    self._render_image_in_viewer(camera_object, graph, is_training)
                 is_training = self.vis["renderingState/isTraining"].read()
                 run_loop = not is_training
                 local_step += 1
@@ -299,33 +325,25 @@ class VisualizerState:
         Used in conjunction with SetTrace.
         """
         if event == "line":
-            if self.check_interrupt_vis and self.res_upscale_factor > 1:
-                self.res_upscale_factor = 1
+            if self.check_interrupt_vis and not self.camera_moving:
                 raise IOChangeException
         return self.check_interrupt
 
-    def _is_render_step(self, step: int, default_steps: int = 2, max_steps: int = 10) -> bool:
-        """dynamically calculate when to render grapic based on resolution of image
+    def _get_camera_object(self):
+        """Gets the camera object from the viewer and updates the movement state if it has changed."""
 
-        Args:
-            step: current train iteration step
-            default_steps: base multiple number of steps
-            max_steps: maximum number of steps in between renders
-        """
-        if step != 0:
-            if self.res_upscale_factor == 1:
-                return True
-            steps_per_render_image = min(default_steps * self.res_upscale_factor, max_steps)
-            steps_condition = step % steps_per_render_image == 0
-            if steps_condition:
-                if self.res_upscale_factor > 3:
-                    if time.time() - self.last_render_time >= self.min_wait_time:
-                        self.last_render_time = time.time()
-                        return True  # if higher res, and minimum wait time achieved
-                    return False  # if higher res, and minimum wait time NOT achieved
-                self.last_render_time = time.time()
-                return True  # if not higher res, but steps met
-        return False  # if init
+        data = self.vis["renderingState/camera"].read()
+        if data is None:
+            return None
+
+        camera_object = data["object"]
+
+        if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix):
+            self.camera_moving = False
+        else:
+            self.prev_camera_matrix = camera_object["matrix"]
+            self.camera_moving = True
+        return camera_object
 
     def _apply_colormap(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
         """Determines which colormap to use based on set colormap type
@@ -435,8 +453,50 @@ class VisualizerState:
             self.vis["renderingState/train_eta"].write("Paused")
             self.vis["renderingState/vis_train_ratio"].write("100% spent on viewer")
 
+    def _calculate_image_height(self, camera_object, is_training: bool) -> Optional[int]:
+        """Calculate the maximum image height that can be rendered in the time budget
+
+        Args:
+            camera_object: the camera object to use for rendering
+            is_training: whether or not we are training
+        Returns:
+            image_height: the maximum image height that can be rendered in the time budget
+        """
+        max_resolution = self.vis["renderingState/maxResolution"].read()
+        if max_resolution:
+            self.max_resolution = max_resolution
+
+        if self.camera_moving or not is_training:
+            target_train_util = 0
+        else:
+            target_train_util = self.vis["renderingState/targetTrainUtil"].read()
+            if target_train_util is None:
+                target_train_util = 0.9
+
+        if EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]:
+            train_rays_per_sec = GLOBAL_BUFFER["events"][EventName.TRAIN_RAYS_PER_SEC.value]["avg"]
+        else:
+            return None
+        if EventName.VIS_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]:
+            vis_rays_per_sec = GLOBAL_BUFFER["events"][EventName.VIS_RAYS_PER_SEC.value]["avg"]
+        else:
+            vis_rays_per_sec = train_rays_per_sec
+
+        current_fps = self.moving_fps if self.camera_moving else self.static_fps
+
+        # calculate number of rays that can be rendered given the target fps
+        num_vis_rays = vis_rays_per_sec / current_fps * (1 - target_train_util)
+
+        aspect_ratio = camera_object["aspect"]
+
+        image_height = (num_vis_rays / aspect_ratio) ** 0.5
+        image_height = int(round(image_height, -1))
+        image_height = min(self.max_resolution, image_height)
+
+        return image_height
+
     @profiler.time_function
-    def _render_image_in_viewer(self, graph: Model) -> None:
+    def _render_image_in_viewer(self, camera_object, graph: Model, is_training: bool) -> None:
         """
         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.
@@ -444,18 +504,6 @@ class VisualizerState:
         Args:
             graph: current checkpoint of model
         """
-        # check and perform camera updates
-        data = self.vis["renderingState/camera"].read()
-        if data is None:
-            return
-        camera_object = data["object"]
-        # hacky way to prevent overflow check to see if < 100; TODO(make less hacky)
-        if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix):
-            self.res_upscale_factor = min(self.res_upscale_factor * 2, 100)
-        else:
-            self.prev_camera_matrix = camera_object["matrix"]
-            self.res_upscale_factor = 1
-
         # check and perform output type updates
         output_type = self.vis["renderingState/output_choice"].read()
         output_type = OutputTypes.INIT if output_type is None else output_type
@@ -467,20 +515,11 @@ class VisualizerState:
         colormap_type = ColormapTypes.INIT if colormap_type is None else colormap_type
         self.prev_colormap_type = colormap_type
 
-        # check and perform min/max update
-        max_resolution = self.vis["renderingState/maxResolution"].read()
-        if max_resolution:
-            self.max_resolution = max_resolution
+        # Calculate camera pose and intrinsics
+        image_height = self._calculate_image_height(camera_object, is_training)
+        if image_height is None:
+            return
 
-        min_resolution = self.vis["renderingState/minResolution"].read()
-        if min_resolution:
-            self.min_resolution = min_resolution
-
-        damped_upsacale_factor = 1 if self.res_upscale_factor < 8 else self.res_upscale_factor
-        image_height = min(
-            self.min_resolution * damped_upsacale_factor,
-            self.max_resolution,
-        )
         intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h(
             camera_object, image_height=image_height
         )