Implementation:Kornia Kornia Colormap

Knowledge Sources	Kornia
Domains	Vision, Color_Processing
Last Updated	2026-02-09 15:00 GMT

Overview

This module provides colormap creation and application utilities for mapping grayscale tensors to RGB using built-in or custom colormaps.

Description

colormap.py is a module in the Kornia library's color subpackage that defines the ColorMapType enumeration of 19 built-in colormaps, the ColorMap class for creating and storing colormap tensors with configurable resolution, and the apply_colormap function and ApplyColorMap nn.Module for applying colormaps to grayscale image tensors. The colormaps are generated by linearly interpolating base color data to a target number of colors, producing a torch.Tensor of shape (3, num_colors).

Usage

Import this module when you need to pseudocolor grayscale images or scalar-valued tensors (e.g. depth maps, heatmaps, attention maps) into RGB images for visualization. Both functional and nn.Module interfaces are provided.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/color/colormap.py
• Lines: 1-305

Signature

class ColorMapType(Enum):
    autumn = 1
    bone = 2
    jet = 3
    winter = 4
    rainbow = 5
    ocean = 6
    summer = 7
    spring = 8
    cool = 9
    hsv = 10
    brg = 11
    pink = 12
    hot = 13
    plasma = 14
    viridis = 15
    cividis = 16
    twilight = 17
    turbo = 18
    seismic = 19

class ColorMap:
    def __init__(
        self,
        base: Union[list[RGBColor], str, ColorMapType],
        num_colors: int = 64,
        device: Optional[torch.device] = None,
        dtype: Optional[torch.dtype] = None,
    ) -> None: ...

def apply_colormap(input_tensor: torch.Tensor, colormap: ColorMap) -> torch.Tensor: ...

class ApplyColorMap(nn.Module):
    def __init__(self, colormap: ColorMap) -> None: ...
    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: ...

Import

from kornia.color import ColorMap, ColorMapType, apply_colormap, ApplyColorMap

I/O Contract

Inputs

Name	Type	Required	Description
base (ColorMap)	Union[list[RGBColor], str, ColorMapType]	Yes	Built-in colormap name (str), ColorMapType enum, or a custom list of RGB triplets.
num_colors (ColorMap)	int	No	Number of colors in the interpolated colormap. Default: 64.
device (ColorMap)	Optional[torch.device]	No	Device for the colormap tensor. Default: None.
dtype (ColorMap)	Optional[torch.dtype]	No	Data type for the colormap tensor. Default: None.
input_tensor (apply_colormap)	torch.Tensor	Yes	Grayscale image tensor of shape (*, C, H, W). Values must be integers in [0, 255] or floats in [0, 1].
colormap (apply_colormap)	ColorMap	Yes	A ColorMap object to apply.

Outputs

Name	Type	Description
ColorMap.colors	torch.Tensor	Tensor of shape (3, num_colors) representing the interpolated colormap.
apply_colormap return	torch.Tensor	RGB tensor of shape (B, C*3, H, W) with the colormap applied.

Usage Examples

Basic Usage

import torch
from kornia.color import ColorMap, ColorMapType, apply_colormap

# Create a viridis colormap with 256 colors
cmap = ColorMap(base='viridis', num_colors=256)
print(cmap.colors.shape)  # torch.Size([3, 256])

# Apply colormap to a grayscale tensor (values 0-255)
gray = torch.randint(0, 256, (1, 1, 64, 64))
rgb = apply_colormap(gray, cmap)
print(rgb.shape)  # torch.Size([1, 3, 64, 64])

# Using ColorMapType enum
cmap_jet = ColorMap(base=ColorMapType.jet, num_colors=128)

# Using a custom colormap from two endpoints
custom_cmap = ColorMap(base=[[0., 0., 1.], [1., 0., 0.]], num_colors=64)

# Using the nn.Module interface
from kornia.color import ApplyColorMap
module = ApplyColorMap(colormap=cmap)
rgb_out = module(gray)

Related Pages

• Principle:Kornia_Kornia_Colormap_Visualization