Render with colormaps (#1893)

* render with colormaps

* black

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)

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:

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(),

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,
     ):

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

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,
         )