From 28532c583f963dc1694b1d09a3414bdea4c9383b Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 9 May 2023 12:26:44 -0700 Subject: [PATCH] Render with colormaps (#1893) * render with colormaps * black --- nerfstudio/utils/colormaps.py | 68 +++++++++++++++++-- nerfstudio/viewer/server/control_panel.py | 17 ++--- .../viewer/server/render_state_machine.py | 12 +++- nerfstudio/viewer/server/viewer_elements.py | 15 ++-- nerfstudio/viewer/server/viewer_utils.py | 62 +---------------- scripts/render.py | 42 ++++++++++-- 6 files changed, 130 insertions(+), 86 deletions(-) diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index aae1b2a92..19b5922ef 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -14,7 +14,7 @@ """ Helper functions for visualizing outputs """ -from typing import Optional +from typing import Literal, Optional import torch from matplotlib import cm @@ -22,19 +22,73 @@ from torchtyping import TensorType from nerfstudio.utils import colors +Colormaps = Literal["default", "turbo", "viridis", "magma", "inferno", "cividis"] -def apply_colormap(image: TensorType["bs":..., 1], cmap="viridis") -> TensorType["bs":..., "rgb":3]: + +def apply_colormap( + image: TensorType["bs":..., 1], + colormap: Colormaps = "turbo", + normalize: bool = False, + colormap_min: float = 0, + colormap_max: float = 1, + invert: bool = False, + eps: float = 1e-9, +) -> TensorType["bs":..., "rgb":3]: + """ + Applies a colormap to a tensor image. Currently only supports 1 and 3 channel inputs + + Args: + image: Input tensor image. + colormap: Colormap to apply. + normalize: Whether to normalize the input tensor image. + colormap_min: Minimum value for the output colormap. + colormap_max: Maximum value for the output colormap. + invert: Whether to invert the output colormap. + eps: Epsilon value for numerical stability. + + Returns: + Tensor with the colormap applied. + """ + + # default for rgb images + if image.shape[-1] == 3: + return image + + # rendering depth outputs + if image.shape[-1] == 1 and torch.is_floating_point(image): + output = image + if normalize: + output = output - torch.min(output) + output = output / (torch.max(output) + eps) + output = output * (colormap_max - colormap_min) + colormap_min + output = torch.clip(output, 0, 1) + if invert: + output = 1 - output + return apply_float_colormap(output, colormap=colormap) + + # rendering boolean outputs + if image.dtype == torch.bool: + return apply_boolean_colormap(image) + + raise NotImplementedError + + +def apply_float_colormap( + image: TensorType["bs":..., 1], colormap: Colormaps = "viridis" +) -> TensorType["bs":..., "rgb":3]: """Convert single channel to a color image. Args: image: Single channel image. - cmap: Colormap for image. + colormap: Colormap for image. Returns: TensorType: Colored image with colors in [0, 1] """ + if colormap == "default": + colormap = "turbo" - colormap = cm.get_cmap(cmap) + colormap = cm.get_cmap(colormap) colormap = torch.tensor(colormap.colors).to(image.device) # type: ignore image = torch.nan_to_num(image, 0) image_long = (image * 255).long() @@ -50,7 +104,7 @@ def apply_depth_colormap( accumulation: Optional[TensorType["bs":..., 1]] = None, near_plane: Optional[float] = None, far_plane: Optional[float] = None, - cmap="turbo", + colormap: Colormaps = "turbo", ) -> TensorType["bs":..., "rgb":3]: """Converts a depth image to color for easier analysis. @@ -59,7 +113,7 @@ def apply_depth_colormap( accumulation: Ray accumulation used for masking vis. near_plane: Closest depth to consider. If None, use min image value. far_plane: Furthest depth to consider. If None, use max image value. - cmap: Colormap to apply. + colormap: Colormap to apply. Returns: Colored depth image with colors in [0, 1] @@ -72,7 +126,7 @@ def apply_depth_colormap( depth = torch.clip(depth, 0, 1) # depth = torch.nan_to_num(depth, nan=0.0) # TODO(ethan): remove this - colored_image = apply_colormap(depth, cmap=cmap) + colored_image = apply_colormap(depth, colormap=colormap) if accumulation is not None: colored_image = colored_image * accumulation + (1 - accumulation) diff --git a/nerfstudio/viewer/server/control_panel.py b/nerfstudio/viewer/server/control_panel.py index 3856fc70e..74f305035 100644 --- a/nerfstudio/viewer/server/control_panel.py +++ b/nerfstudio/viewer/server/control_panel.py @@ -14,11 +14,11 @@ """ Control panel for the viewer """ from collections import defaultdict -from typing import Callable, DefaultDict, List, Tuple +from typing import Callable, DefaultDict, List, Tuple, get_args import torch -from nerfstudio.viewer.server import viewer_utils +from nerfstudio.utils.colormaps import Colormaps from nerfstudio.viewer.server.viewer_elements import ( ViewerCheckbox, ViewerDropdown, @@ -55,7 +55,7 @@ class ControlPanel: ["not set"], cb_hook=lambda han: [self.update_control_panel(), update_output_cb(han), rerender_cb(han)], ) - self._colormap = ViewerDropdown("Colormap", "default", ["default"], cb_hook=rerender_cb) + self._colormap = ViewerDropdown[Colormaps]("Colormap", "default", ["default"], cb_hook=rerender_cb) self._invert = ViewerCheckbox("| Invert", False, cb_hook=rerender_cb) self._normalize = ViewerCheckbox("| Normalize", False, cb_hook=rerender_cb) self._min = ViewerNumber("| Min", 0.0, cb_hook=rerender_cb) @@ -150,11 +150,11 @@ class ControlPanel: dimensions: the number of dimensions of the render dtype: the data type of the render """ - colormap_options = [] + colormap_options: List[Colormaps] = [] if dimensions == 3: - colormap_options = [viewer_utils.ColormapTypes.DEFAULT.value] + colormap_options = ["default"] if dimensions == 1 and dtype == torch.float: - colormap_options = [c.value for c in list(viewer_utils.ColormapTypes)[1:]] + colormap_options = [c for c in list(get_args(Colormaps)) if c != "default"] self._colormap.set_options(colormap_options) @@ -169,9 +169,10 @@ class ControlPanel: return self._output_render.value @property - def colormap(self) -> str: + def colormap(self) -> Colormaps: """Returns the current colormap""" - return self._colormap.value + colormap = self._colormap.value + return colormap @property def colormap_invert(self) -> bool: diff --git a/nerfstudio/viewer/server/render_state_machine.py b/nerfstudio/viewer/server/render_state_machine.py index 56ebad425..f3af3e3fe 100644 --- a/nerfstudio/viewer/server/render_state_machine.py +++ b/nerfstudio/viewer/server/render_state_machine.py @@ -24,7 +24,7 @@ import torch from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.model_components.renderers import background_color_override_context -from nerfstudio.utils import writer +from nerfstudio.utils import colormaps, writer from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter from nerfstudio.viewer.server import viewer_utils from nerfstudio.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h @@ -238,7 +238,15 @@ class RenderStateMachine(threading.Thread): self.viewer.update_colormap_options( dimensions=outputs[output_render].shape[-1], dtype=outputs[output_render].dtype ) - selected_output = (viewer_utils.apply_colormap(self.viewer.control_panel, outputs) * 255).type(torch.uint8) + selected_output = colormaps.apply_colormap( + image=outputs[self.viewer.control_panel.output_render], + colormap=self.viewer.control_panel.colormap, + normalize=self.viewer.control_panel.colormap_normalize, + colormap_min=self.viewer.control_panel.colormap_min, + colormap_max=self.viewer.control_panel.colormap_max, + invert=self.viewer.control_panel.colormap_invert, + ) + selected_output = (selected_output * 255).type(torch.uint8) self.viewer.viser_server.set_background_image( selected_output.cpu().numpy(), diff --git a/nerfstudio/viewer/server/viewer_elements.py b/nerfstudio/viewer/server/viewer_elements.py index 19d9e7588..fbf60bfb6 100644 --- a/nerfstudio/viewer/server/viewer_elements.py +++ b/nerfstudio/viewer/server/viewer_elements.py @@ -12,12 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=too-few-public-methods,access-member-before-definition,attribute-defined-outside-init + """ Viewer GUI elements for the nerfstudio viewer """ from __future__ import annotations from abc import abstractmethod -from typing import Any, Callable, Generic, List, Optional, Tuple, TypeVar +from typing import Any, Callable, Generic, List, Optional, Tuple + +from typing_extensions import TypeVar from nerfstudio.viewer.viser import GuiHandle, GuiSelectHandle, ViserServer @@ -266,7 +270,10 @@ class ViewerCheckbox(ViewerParameter[bool]): self.gui_handle.set_disabled(self.disabled) -class ViewerDropdown(ViewerParameter[str]): +TString = TypeVar("TString", default=str) + + +class ViewerDropdown(ViewerParameter[TString], Generic[TString]): """A dropdown in the viewer Args: @@ -282,8 +289,8 @@ class ViewerDropdown(ViewerParameter[str]): def __init__( self, name: str, - default_value: str, - options: List[str], + default_value: TString, + options: List[TString], disabled: bool = False, cb_hook: Callable[[ViewerDropdown], Any] = lambda element: None, ): diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index be3b480c7..e47f652c9 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -17,21 +17,18 @@ """Code to interface with the `vis/` (the JS viewer).""" from __future__ import annotations -import enum import os import socket import sys from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple import torch from rich.console import Console from nerfstudio.data.scene_box import SceneBox from nerfstudio.models.base_model import Model -from nerfstudio.utils import colormaps from nerfstudio.utils.io import load_from_json -from nerfstudio.viewer.server.control_panel import ControlPanel if TYPE_CHECKING: from nerfstudio.engine.trainer import Trainer @@ -59,17 +56,6 @@ def get_viewer_url(websocket_port: int) -> str: return f"https://viewer.nerf.studio/versions/{version}/?websocket_url={websocket_url}" -class ColormapTypes(str, enum.Enum): - """List of colormap render types""" - - DEFAULT = "default" - TURBO = "turbo" - VIRIDIS = "viridis" - MAGMA = "magma" - INFERNO = "inferno" - CIVIDIS = "cividis" - - class IOChangeException(Exception): """Basic camera exception to interrupt viewer""" @@ -149,49 +135,3 @@ def update_render_aabb( model.render_aabb = SceneBox(aabb=torch.stack([crop_min, crop_max], dim=0)) else: model.render_aabb = None - - -def apply_colormap( - control_panel: ControlPanel, outputs: Dict[str, Any], colors: Optional[torch.Tensor] = None, eps=1e-6 -): - """Determines which colormap to use based on set colormap type - - Args: - control_panel: control panel object - outputs: the output tensors for which to apply colormaps on - colors: is only set if colormap is for semantics. Defaults to None. - eps: epsilon to handle floating point comparisons - """ - colormap_type = control_panel.colormap - output_type = control_panel.output_render - - # default for rgb images - if colormap_type == ColormapTypes.DEFAULT and outputs[output_type].shape[-1] == 3: - return outputs[output_type] - - # rendering depth outputs - if outputs[output_type].shape[-1] == 1 and outputs[output_type].dtype == torch.float: - output = outputs[output_type] - if control_panel.colormap_normalize: - output = output - torch.min(output) - output = output / (torch.max(output) + eps) - output = output * (control_panel.colormap_max - control_panel.colormap_min) + control_panel.colormap_min - output = torch.clip(output, 0, 1) - if control_panel.colormap_invert: - output = 1 - output - if colormap_type == ColormapTypes.DEFAULT: - return colormaps.apply_colormap(output, cmap=ColormapTypes.TURBO.value) - return colormaps.apply_colormap(output, cmap=colormap_type) - - # rendering semantic outputs - if outputs[output_type].dtype == torch.int: - logits = outputs[output_type] - labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1) # type: ignore - assert colors is not None - return colors[labels] - - # rendering boolean outputs - if outputs[output_type].dtype == torch.bool: - return colormaps.apply_boolean_colormap(outputs[output_type]) - - raise NotImplementedError diff --git a/scripts/render.py b/scripts/render.py index 4549a9bdd..bdff762d9 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -40,7 +40,7 @@ from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.data.scene_box import SceneBox from nerfstudio.model_components import renderers from nerfstudio.pipelines.base_pipeline import Pipeline -from nerfstudio.utils import install_checks +from nerfstudio.utils import colormaps, install_checks from nerfstudio.utils.eval_utils import eval_setup from nerfstudio.utils.rich_utils import ItersPerSecColumn @@ -57,6 +57,11 @@ def _render_trajectory_video( seconds: float = 5.0, output_format: Literal["images", "video"] = "video", camera_type: CameraType = CameraType.PERSPECTIVE, + colormap: colormaps.Colormaps = "turbo", + normalize: bool = True, + colormap_min: float = 0.0, + colormap_max: float = 1.0, + invert: bool = False, ) -> None: """Helper function to create a video of the spiral trajectory. @@ -70,6 +75,11 @@ def _render_trajectory_video( seconds: Length of output video. output_format: How to save output data. camera_type: Camera projection format type. + colormap: Colormap to use if single channel output. + normalize: Whether to normalize the output images, only works with colormaps. + colormap_min: Minimum colormap value. + colormap_max: Maximum colormap value. + invert: Whether to invert the colormap. """ CONSOLE.print("[bold green]Creating trajectory " + output_format) cameras.rescale_output_resolution(rendered_resolution_scaling_factor) @@ -126,9 +136,19 @@ def _render_trajectory_video( f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center" ) sys.exit(1) - output_image = outputs[rendered_output_name].cpu().numpy() - if output_image.shape[-1] == 1: - output_image = np.concatenate((output_image,) * 3, axis=-1) + output_image = outputs[rendered_output_name] + output_image = ( + colormaps.apply_colormap( + image=output_image, + colormap=colormap, + normalize=normalize, + colormap_min=colormap_min, + colormap_max=colormap_max, + invert=invert, + ) + .cpu() + .numpy() + ) render_image.append(output_image) render_image = np.concatenate(render_image, axis=1) if output_format == "images": @@ -293,6 +313,15 @@ class RenderTrajectory: """Number of interpolation steps between eval dataset cameras.""" eval_num_rays_per_chunk: Optional[int] = None """Specifies number of rays per chunk during eval.""" + colormap: colormaps.Colormaps = "turbo" + normalize: bool = False + """Whether to normalize the output if using a colormap.""" + colormap_min: float = 0.0 + """Minimum value of the colormap.""" + colormap_max: float = 1.0 + """Maximum value of the colormap.""" + invert: bool = False + """Whether to invert the colormap.""" def main(self) -> None: """Main function.""" @@ -345,6 +374,11 @@ class RenderTrajectory: seconds=seconds, output_format=self.output_format, camera_type=camera_type, + colormap=self.colormap, + normalize=self.normalize, + colormap_min=self.colormap_min, + colormap_max=self.colormap_max, + invert=self.invert, )