Adding colormap customization on viewer (#272)

* set up logic for colormap defined by user on viewer

* ooops typo

* nit

* done with defaults

* lint

nerfactory/models/base.py

@@ -125,14 +125,6 @@ class Model(nn.Module):
             Outputs of model. (ie. rendered colors)
         """
 
-    def process_outputs_as_images(self, outputs):  # pylint:disable=no-self-use
-        """Process output images into visualizable colored images"""
-        # TODO override this function elsewhere or do something else for processing images
-        for k, v in outputs.items():
-            if v.shape[-1] == 1:
-                v = torch.tile(v, (1, 1, 3))
-            outputs[k] = v
-
     @overload
     def forward(self, ray_bundle: RayBundle, batch: None = None) -> Dict[str, torch.Tensor]:
         ...

nerfactory/models/instant_ngp.py

@@ -137,17 +137,6 @@ class NGPModel(Model):
         loss_dict = misc.scale_dict(loss_dict, loss_coefficients)
         return loss_dict
 
-    def process_outputs_as_images(self, outputs):  # pylint:disable=no-self-use
-        """Do preprocessing to make images valid"""
-        # TODO: make log_test_image_outputs use this directly
-        # TODO: implement across all the different model implementations
-        outputs["accumulation"] = visualization.apply_colormap(outputs["accumulation"])
-        outputs["depth"] = visualization.apply_depth_colormap(
-            outputs["depth"],
-            accumulation=outputs["accumulation"],
-        )
-        outputs["alive_ray_mask"] = visualization.apply_boolean_colormap(outputs["alive_ray_mask"])
-
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
         image = batch["image"]
         rgb = outputs["rgb"]

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

@@ -22,6 +22,12 @@ export function RenderControls() {
   const outputChoice = useSelector(
     (state) => state.renderingState.output_choice,
   );
+  const colormapOptions = useSelector(
+    (state) => state.renderingState.colormap_options,
+  );
+  const colormapChoice = useSelector(
+    (state) => state.renderingState.colormap_choice,
+  );
   const min_resolution = useSelector(
     (state) => state.renderingState.minResolution,
   );
@@ -96,6 +102,25 @@ export function RenderControls() {
     }
   };
 
+  const set_colormap_choice = (value) => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      dispatch({
+        type: 'write',
+        path: 'renderingState/colormap_choice',
+        data: value,
+      });
+      const cmd = 'write';
+      const path = 'renderingState/colormap_choice';
+      const data = {
+        type: cmd,
+        path,
+        data: value,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
   const set_fov = (value) => {
     if (websocket.readyState === WebSocket.OPEN) {
       dispatch({
@@ -174,6 +199,16 @@ export function RenderControls() {
           set_output_choice(v);
         },
       },
+      // colormap_options
+      colormap_options: {
+        label: 'Colormap',
+        options: colormapOptions,
+        value: colormapChoice,
+        onChange: (v) => {
+          set_colormap_choice(v);
+        },
+        disabled: colormapOptions.length === 1
+      },
       // resolution
       min_resolution: {
         label: 'Min Res.',
@@ -225,6 +260,8 @@ export function RenderControls() {
       isTraining,
       outputOptions,
       outputChoice,
+      colormapOptions,
+      colormapChoice,
       min_resolution,
       max_resolution,
       field_of_view,
@@ -273,12 +310,15 @@ export function RenderControls() {
     setControls({ min_resolution });
     setControls({ max_resolution });
     setControls({ output_options: outputChoice });
+    setControls({ colormap_options: colormapChoice });
     setControls({ 'Camera FoV': field_of_view });
   }, [
     setControls,
     isTraining,
     outputOptions,
     outputChoice,
+    colormapOptions,
+    colormapChoice,
     min_resolution,
     max_resolution,
     field_of_view,

nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanelSlice.js

@@ -4,3 +4,4 @@
 export const selectTrainingState = (state) =>
   state.shared.rendering.training_state;
 export const selectOutputOptions = (state) => state.shared.output_options;
+export const selectColormapOptions = (state) => state.shared.colormap_options;

nerfactory/viewer/app/src/reducer.js

@@ -15,6 +15,8 @@ const initialState = {
     isTraining: true,
     output_options: ['rgb'], // populated by the possible Graph outputs
     output_choice: 'rgb', // the selected output
+    colormap_options: ['default'], // populated by the output choice
+    colormap_choice: 'default', // the selected colormap
     minResolution: 50,
     maxResolution: 512,
     field_of_view: 80,
@@ -22,6 +24,7 @@ const initialState = {
     train_eta: 'Paused',
     vis_train_ratio: 'Paused',
   },
+  
   // the scene state
   sceneState: {
     sceneBounds: null,
@@ -33,6 +36,9 @@ const initialState = {
 };
 
 function setData(newState, state, path, data) {
+  if (path === 'colormap_options') {
+    newState.colormap_choice = 'default'; // eslint-disable-line no-param-reassign
+  }
   if (path.length === 1) {
     newState[path[0]] = data; // eslint-disable-line no-param-reassign
   } else {

nerfactory/viewer/server/viewer_utils.py

@@ -15,6 +15,7 @@
 """Code to interface with the `vis/` (the JS visualizer).
 """
 
+import enum
 import logging
 import sys
 import threading
@@ -29,7 +30,7 @@ from nerfactory.cameras.rays import RayBundle
 from nerfactory.dataloaders.image_dataset import ImageDataset
 from nerfactory.dataloaders.structs import DatasetInputs
 from nerfactory.models.base import Model
-from nerfactory.utils import profiler
+from nerfactory.utils import profiler, visualization
 from nerfactory.utils.config import ViewerConfig
 from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all
 from nerfactory.utils.writer import GLOBAL_BUFFER, EventName
@@ -37,6 +38,25 @@ from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_w
 from nerfactory.viewer.server.visualizer import Viewer
 
 
+class OutputTypes(str, enum.Enum):
+    """Noncomprehsnive list of output render types"""
+
+    INIT = "init"
+    RGB = "rgb"
+    RGB_FINE = "rgb_fine"
+
+
+class ColormapTypes(str, enum.Enum):
+    """Noncomprehsnive list of colormap render types"""
+
+    INIT = "init"
+    DEFAULT = "default"
+    TURBO = "turbo"
+    DEPTH = "depth"
+    SEMANTIC = "semantic"
+    BOOLEAN = "boolean"
+
+
 class IOChangeException(Exception):
     """Basic camera exception to interrupt visualizer"""
 
@@ -85,7 +105,6 @@ class RenderThread(threading.Thread):
             self.exc = e
 
         if outputs:
-            self.graph.process_outputs_as_images(outputs)
             self.vis_outputs = outputs
 
         self.state.check_done_render = True
@@ -129,11 +148,19 @@ class CheckThread(threading.Thread):
             # check output type
             output_type = self.state.vis["renderingState/output_choice"].read()
             if output_type is None:
-                output_type = "default"
+                output_type = OutputTypes.INIT
             if self.state.prev_output_type != output_type:
                 self.state.check_interrupt_vis = True
                 return
 
+            # check colormap type
+            colormap_type = self.state.vis["renderingState/colormap_choice"].read()
+            if colormap_type is None:
+                colormap_type = ColormapTypes.INIT
+            if self.state.prev_colormap_type != colormap_type:
+                self.state.check_interrupt_vis = True
+                return
+
             # check max render
             max_resolution = self.state.vis["renderingState/maxResolution"].read()
             if max_resolution is not None:
@@ -169,7 +196,9 @@ class VisualizerState:
 
         # visualizer specific variables
         self.prev_camera_matrix = None
-        self.prev_output_type = "default"
+        self.prev_output_type = OutputTypes.INIT
+        self.prev_colormap_type = ColormapTypes.INIT
+        self.output_type_changed = True
         self.min_resolution = 50
         self.max_resolution = 1000
         self.res_upscale_factor = 1
@@ -263,20 +292,79 @@ class VisualizerState:
                 return True  # if not higher res, but steps met
         return False  # if init
 
-    def _send_output_to_viewer(self, outputs: Dict[str, Any]):
+    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
+
+        Args:
+            outputs: the output tensors for which to apply colormaps on
+            stuff_colors: is only set if colormap is for semantics. Defaults to None.
+        """
+        # default for rgb images
+        if self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].shape[-1] == 3:
+            return outputs[self.prev_output_type]
+
+        # rendering depth outputs
+        if self.prev_colormap_type == ColormapTypes.DEPTH or (
+            self.prev_colormap_type == ColormapTypes.DEFAULT
+            and outputs[self.prev_output_type].dtype == torch.float
+            and (torch.max(outputs[self.prev_output_type]) - 1.0) > eps  # handle floating point arithmetic
+        ):
+            return visualization.apply_depth_colormap(
+                outputs[self.prev_output_type], accumulation=outputs["accumulation"]
+            )
+
+        # rendering accumulation outputs
+        if self.prev_colormap_type == ColormapTypes.TURBO or (
+            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.float
+        ):
+            return visualization.apply_colormap(outputs[self.prev_output_type])
+
+        # rendering semantic outputs
+        if self.prev_colormap_type == ColormapTypes.SEMANTIC or (
+            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.int
+        ):
+            logits = outputs[self.prev_output_type]
+            labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1)  # type: ignore
+            assert stuff_colors is not None
+            return stuff_colors[labels]
+
+        # rendering boolean outputs
+        if self.prev_colormap_type == ColormapTypes.BOOLEAN or (
+            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.bool
+        ):
+            return visualization.apply_boolean_colormap(outputs[self.prev_output_type])
+
+        raise NotImplementedError
+
+    def _send_output_to_viewer(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
         """Chooses the correct output and sends it to the viewer
 
         Args:
             outputs: the dictionary of outputs to choose from, from the graph
+            stuff_colors: is only set if colormap is for semantics. Defaults to None.
         """
         if not self.outputs_set:
-            # set_output_options(self.vis, list(outputs.keys()))
             self.vis["renderingState/output_options"].write(list(outputs.keys()))
             self.outputs_set = True
         # gross hack to get the image key, depending on which keys the graph uses
-        if self.prev_output_type == "default":
-            self.prev_output_type = "rgb" if "rgb" in outputs else "rgb_fine"
-        image_output = outputs[self.prev_output_type].cpu().numpy() * 255
+        if self.prev_output_type == OutputTypes.INIT:
+            self.prev_output_type = OutputTypes.RGB if OutputTypes.RGB in outputs else OutputTypes.RGB_FINE
+        # re-register colormaps and send to viewer
+        if self.output_type_changed or self.prev_colormap_type == ColormapTypes.INIT:
+            self.prev_colormap_type = ColormapTypes.DEFAULT
+            colormap_options = [ColormapTypes.DEFAULT]
+            if (
+                outputs[self.prev_output_type].shape[-1] != 3
+                and outputs[self.prev_output_type].dtype == torch.float
+                and (torch.max(outputs[self.prev_output_type]) - 1.0) <= eps  # handle floating point arithmetic
+            ):
+                # accumulation can also include depth
+                colormap_options.extend(["depth"])
+            self.output_type_changed = False
+            self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type)
+            self.vis["renderingState/colormap_options"].write(colormap_options)
+        selected_output = self._apply_colormap(outputs, stuff_colors).cpu().numpy()
+        image_output = selected_output * 255
         image = (image_output).astype("uint8")
         self.vis.set_image(image)
 
@@ -327,10 +415,15 @@ class VisualizerState:
 
         # check and perform output type updates
         output_type = self.vis["renderingState/output_choice"].read()
-        if output_type is None:
-            output_type = "default"
+        output_type = OutputTypes.INIT if output_type is None else output_type
+        self.output_type_changed = self.prev_output_type != output_type
         self.prev_output_type = output_type
 
+        # check and perform colormap type updates
+        colormap_type = self.vis["renderingState/colormap_choice"].read()
+        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:
@@ -381,7 +474,8 @@ class VisualizerState:
 
         outputs = render_thread.vis_outputs
         if outputs is not None:
-            self._send_output_to_viewer(outputs)
+            stuff_colors = graph.stuff_colors if hasattr(graph, "stuff_colors") else None
+            self._send_output_to_viewer(outputs, stuff_colors=stuff_colors)
             self._update_viewer_stats(render_duration, image_height)