From 2d3452e5a260107b3b43d66abc22ad248f3217a9 Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 12 Aug 2022 17:35:53 -0700 Subject: [PATCH] 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 | 8 -- nerfactory/models/instant_ngp.py | 11 -- .../src/modules/ConfigPanel/ConfigPanel.jsx | 40 ++++++ .../modules/ConfigPanel/ConfigPanelSlice.js | 1 + nerfactory/viewer/app/src/reducer.js | 6 + nerfactory/viewer/server/viewer_utils.py | 118 ++++++++++++++++-- 6 files changed, 153 insertions(+), 31 deletions(-) diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index 3edc64e21..ad9653762 100644 --- a/nerfactory/models/base.py +++ b/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]: ... diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index f852cf685..36a4c6f66 100644 --- a/nerfactory/models/instant_ngp.py +++ b/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"] diff --git a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 9347e6245..a97b80957 100644 --- a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/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, diff --git a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanelSlice.js b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanelSlice.js index 30db097df..49f86a065 100644 --- a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanelSlice.js +++ b/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; diff --git a/nerfactory/viewer/app/src/reducer.js b/nerfactory/viewer/app/src/reducer.js index 7d084abe6..9af63cc65 100644 --- a/nerfactory/viewer/app/src/reducer.js +++ b/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 { diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 21a2006ce..62b03415d 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/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)