Repository: kornia/kornia

Issue:
Implement `applyColorMap` from opencv
## 🚀 Feature

Implement `applyColorMap` from opencv

https://docs.opencv.org/ref/master/d3/d50/group__imgproc__colormap.html#gadf478a5e5ff49d8aa24e726ea6f65d15
______________________________________________________________________

#### Consider also to contribute to Kornia universe projects :)

<sub>

- [**Tutorials**](https://github.com/kornia/tutorials): our repository containing the tutorials.

</sub>



Discussion:
Should the colormap, have all the values (for example for 8 bits, one color for each value between 0-255), or should the function do something more complete like try to distribute the existing colors in the colormap among the existing values in the source image? 


cv2 code: https://github.com/opencv/opencv/blob/ebb6915e588fcee1e6664cce670f0253bac0e67b/modules/imgproc/src/colormap.cpp#L805
i would try to match whatever opencv this at first order -- later we can do more fancy things, even thinking how to differentiate, but later 

Resolving patch:
diff --git a/docs/generate_example_images.py b/docs/generate_example_images.py
--- a/docs/generate_example_images.py
+++ b/docs/generate_example_images.py
@@ -142,6 +142,7 @@ def main():
         "rgb_to_ycbcr": ((), 1),
         "rgb_to_yuv": ((), 1),
         "rgb_to_linear_rgb": ((), 1),
+        "apply_colormap": ((K.color.colormap.AUTUMN(255),), 1),
     }
     # ITERATE OVER THE TRANSFORMS
     for fn_name, (args, num_samples) in color_transforms_list.items():
@@ -149,6 +150,9 @@ def main():
         fn = getattr(mod, fn_name)
         if fn_name == "grayscale_to_rgb":
             out = fn(K.color.rgb_to_grayscale(img2), *args)
+        elif fn_name == "apply_colormap":
+            gray_image = (K.color.rgb_to_grayscale(img2) * 255.0).round()
+            out = K.color.rgb_to_bgr(fn(gray_image, *args))
         else:
             out = fn(img2, *args)
         # perform normalization to visualize
@@ -162,7 +166,7 @@ def main():
         if out.shape[1] != 3:
             out = out.repeat(1, 3, 1, 1)
         # save the output image
-        if fn_name == "grayscale_to_rgb":
+        if fn_name in ("grayscale_to_rgb", "apply_colormap"):
             out = torch.cat(
                 [K.color.rgb_to_grayscale(img2[0]).repeat(3, 1, 1), *(out[i] for i in range(out.size(0)))], dim=-1
             )
@@ -173,6 +177,21 @@ def main():
         sig = f"{fn_name}({', '.join([str(a) for a in args])})"
         print(f"Generated image example for {fn_name}. {sig}")
 
+    colormaps_list = {"AUTUMN": (256,)}
+    bar_img_gray = torch.range(0, 255).repeat(1, 40, 1)  # 1x1x40x256
+    bar_img = K.color.grayscale_to_rgb(bar_img_gray)
+    # ITERATE OVER THE COLORMAPS
+    for colormap_name, args in colormaps_list.items():
+        cm = getattr(mod, colormap_name)(*args)
+        out = K.color.rgb_to_bgr(K.color.apply_colormap(bar_img_gray, cm))
+
+        out = torch.cat([bar_img, out], dim=-1)
+
+        out_np = K.utils.tensor_to_image((out * 255.0).byte())
+        cv2.imwrite(str(OUTPUT_PATH / f"{colormap_name}.png"), out_np)
+        sig = f"{colormap_name}({', '.join([str(a) for a in args])})"
+        print(f"Generated image example for {colormap_name}. {sig}")
+
     # korna.enhance module
     mod = importlib.import_module("kornia.enhance")
     transforms: dict = {
diff --git a/kornia/color/__init__.py b/kornia/color/__init__.py
--- a/kornia/color/__init__.py
+++ b/kornia/color/__init__.py
@@ -1,3 +1,4 @@
+from .colormap import AUTUMN, ApplyColorMap, ColorMap, RGBColor, apply_colormap
 from .gray import BgrToGrayscale, GrayscaleToRgb, RgbToGrayscale, bgr_to_grayscale, grayscale_to_rgb, rgb_to_grayscale
 from .hls import HlsToRgb, RgbToHls, hls_to_rgb, rgb_to_hls
 from .hsv import HsvToRgb, RgbToHsv, hsv_to_rgb, rgb_to_hsv
@@ -107,6 +108,11 @@
     "RgbToLinearRgb",
     "Sepia",
     "sepia",
+    "AUTUMN",
+    "ApplyColorMap",
+    "ColorMap",
+    "RGBColor",
+    "apply_colormap",
 ]
 
 sepia = sepia_from_rgb
diff --git a/kornia/color/colormap.py b/kornia/color/colormap.py
new file mode 100644
--- /dev/null
+++ b/kornia/color/colormap.py
@@ -0,0 +1,204 @@
+from abc import ABC
+from typing import List, Optional
+
+import torch
+from torch.nn.functional import interpolate
+
+from kornia.core import Module, Tensor, tensor
+from kornia.testing import KORNIA_CHECK_IS_GRAY
+
+RGBColor = List[float]
+
+
+def _list_color_to_tensor(
+    colors: List[RGBColor], dtype: Optional[torch.dtype], device: Optional[torch.device]
+) -> Tensor:
+    return tensor([color for color in colors], dtype=dtype, device=device).T
+
+
+class ColorMap(ABC):
+    r"""Abstract class to represents a color map."""
+    colors: Tensor
+
+    def __init__(
+        self,
+        base_colormap: List[RGBColor],
+        num_colors: int,
+        device: Optional[torch.device],
+        dtype: Optional[torch.dtype],
+    ):
+        self._dtype = dtype
+        self._device = device
+
+        self.colors = interpolate(
+            _list_color_to_tensor(base_colormap, dtype=self._dtype, device=self._device)[None, ...],
+            size=num_colors,
+            mode='linear',
+        )[0, ...]
+
+    def __len__(self) -> int:
+        return self.colors.shape[-1]
+
+
+# TODO: Add more colormaps
+_BASE_AUTUMN: List[RGBColor] = [
+    [1.0, 0.0, 0.0],
+    [1.0, 0.01587301587301587, 0.0],
+    [1.0, 0.03174603174603174, 0.0],
+    [1.0, 0.04761904761904762, 0.0],
+    [1.0, 0.06349206349206349, 0.0],
+    [1.0, 0.07936507936507936, 0.0],
+    [1.0, 0.09523809523809523, 0.0],
+    [1.0, 0.1111111111111111, 0.0],
+    [1.0, 0.126984126984127, 0.0],
+    [1.0, 0.1428571428571428, 0.0],
+    [1.0, 0.1587301587301587, 0.0],
+    [1.0, 0.1746031746031746, 0.0],
+    [1.0, 0.1904761904761905, 0.0],
+    [1.0, 0.2063492063492063, 0.0],
+    [1.0, 0.2222222222222222, 0.0],
+    [1.0, 0.2380952380952381, 0.0],
+    [1.0, 0.253968253968254, 0.0],
+    [1.0, 0.2698412698412698, 0.0],
+    [1.0, 0.2857142857142857, 0.0],
+    [1.0, 0.3015873015873016, 0.0],
+    [1.0, 0.3174603174603174, 0.0],
+    [1.0, 0.3333333333333333, 0.0],
+    [1.0, 0.3492063492063492, 0.0],
+    [1.0, 0.3650793650793651, 0.0],
+    [1.0, 0.3809523809523809, 0.0],
+    [1.0, 0.3968253968253968, 0.0],
+    [1.0, 0.4126984126984127, 0.0],
+    [1.0, 0.4285714285714285, 0.0],
+    [1.0, 0.4444444444444444, 0.0],
+    [1.0, 0.4603174603174603, 0.0],
+    [1.0, 0.4761904761904762, 0.0],
+    [1.0, 0.492063492063492, 0.0],
+    [1.0, 0.5079365079365079, 0.0],
+    [1.0, 0.5238095238095238, 0.0],
+    [1.0, 0.5396825396825397, 0.0],
+    [1.0, 0.5555555555555556, 0.0],
+    [1.0, 0.5714285714285714, 0.0],
+    [1.0, 0.5873015873015873, 0.0],
+    [1.0, 0.6031746031746031, 0.0],
+    [1.0, 0.6190476190476191, 0.0],
+    [1.0, 0.6349206349206349, 0.0],
+    [1.0, 0.6507936507936508, 0.0],
+    [1.0, 0.6666666666666666, 0.0],
+    [1.0, 0.6825396825396826, 0.0],
+    [1.0, 0.6984126984126984, 0.0],
+    [1.0, 0.7142857142857143, 0.0],
+    [1.0, 0.7301587301587301, 0.0],
+    [1.0, 0.746031746031746, 0.0],
+    [1.0, 0.7619047619047619, 0.0],
+    [1.0, 0.7777777777777778, 0.0],
+    [1.0, 0.7936507936507936, 0.0],
+    [1.0, 0.8095238095238095, 0.0],
+    [1.0, 0.8253968253968254, 0.0],
+    [1.0, 0.8412698412698413, 0.0],
+    [1.0, 0.8571428571428571, 0.0],
+    [1.0, 0.873015873015873, 0.0],
+    [1.0, 0.8888888888888888, 0.0],
+    [1.0, 0.9047619047619048, 0.0],
+    [1.0, 0.9206349206349206, 0.0],
+    [1.0, 0.9365079365079365, 0.0],
+    [1.0, 0.9523809523809523, 0.0],
+    [1.0, 0.9682539682539683, 0.0],
+    [1.0, 0.9841269841269841, 0.0],
+    [1.0, 1.0, 0.0],
+]
+
+
+# Generate complete color map
+class AUTUMN(ColorMap):
+    r"""The GNU Octave colormap `autumn`
+
+    .. image:: _static/img/AUTUMN.png
+    """
+
+    def __init__(
+        self, num_colors: int = 64, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None
+    ):
+        super().__init__(base_colormap=_BASE_AUTUMN, num_colors=num_colors, device=device, dtype=dtype)
+
+
+def apply_colormap(input_tensor: Tensor, colormap: ColorMap) -> Tensor:
+    r"""Apply to a gray tensor a colormap.
+
+        The image data is assumed to be integer values.
+
+    .. image:: _static/img/apply_colormap.png
+
+    Args:
+        input_tensor: the input tensor of a gray image.
+        colormap: the colormap desired to be applied to the input tensor.
+
+    Returns:
+        A RGB tensor with the applied color map into the input_tensor
+
+    Example:
+        >>> input_tensor = torch.tensor([[[0, 1, 2], [25, 50, 63]]])
+        >>> colormap = AUTUMN()
+        >>> apply_colormap(input_tensor, colormap)
+        tensor([[[1.0000, 1.0000, 1.0000],
+                 [1.0000, 1.0000, 1.0000]],
+        <BLANKLINE>
+                [[0.0000, 0.0159, 0.0317],
+                 [0.3968, 0.7937, 1.0000]],
+        <BLANKLINE>
+                [[0.0000, 0.0000, 0.0000],
+                 [0.0000, 0.0000, 0.0000]]])
+    """
+    # FIXME: implement to work with RGB images
+    # should work with KORNIA_CHECK_SHAPE(x, ["B","C", "H", "W"])
+
+    KORNIA_CHECK_IS_GRAY(input_tensor)
+
+    if len(input_tensor.shape) == 4 and input_tensor.shape[1] == 1:  # if (B x 1 X H x W)
+        input_tensor = input_tensor[:, 0, ...]  # (B x H x W)
+    elif len(input_tensor.shape) == 3 and input_tensor.shape[0] == 1:  # if (1 X H x W)
+        input_tensor = input_tensor[0, ...]  # (H x W)
+
+    keys = torch.arange(0, len(colormap), dtype=input_tensor.dtype, device=input_tensor.device)  # (num_colors)
+
+    index = torch.bucketize(input_tensor, keys)  # shape equals <input_tensor>: (B x H x W) or (H x W)
+
+    output = colormap.colors[:, index]  # (3 x B x H x W) or (3 x H x W)
+
+    if len(output.shape) == 4:
+        output = output.permute(1, 0, -2, -1)  # (B x 3 x H x W)
+
+    return output  # (B x 3 x H x W) or (3 x H x W)
+
+
+class ApplyColorMap(Module):
+    r"""Module that apply to a gray tensor (integer tensor) a colormap.
+
+    Args:
+        input_tensor: the input tensor of a gray image.
+        colormap: the colormap desired to be applied to the input tensor.
+
+    Returns:
+        A RGB tensor with the applied color map into the input_tensor
+
+    Example:
+        >>> input_tensor = torch.tensor([[[0, 1, 2], [25, 50, 63]]])
+        >>> colormap = AUTUMN()
+        >>> apply_cm = ApplyColorMap(colormap)
+        >>> apply_cm(input_tensor)
+        tensor([[[1.0000, 1.0000, 1.0000],
+                 [1.0000, 1.0000, 1.0000]],
+        <BLANKLINE>
+                [[0.0000, 0.0159, 0.0317],
+                 [0.3968, 0.7937, 1.0000]],
+        <BLANKLINE>
+                [[0.0000, 0.0000, 0.0000],
+                 [0.0000, 0.0000, 0.0000]]])
+    """
+
+    def __init__(self, colormap: ColorMap) -> None:
+        super().__init__()
+        self.colormap = colormap
+
+    def forward(self, input_tensor: Tensor) -> Tensor:
+        return apply_colormap(input_tensor, self.colormap)


Test patch:
diff --git a/test/color/test_colormap.py b/test/color/test_colormap.py
new file mode 100644
--- /dev/null
+++ b/test/color/test_colormap.py
@@ -0,0 +1,92 @@
+import pytest
+import torch
+from torch.autograd import gradcheck
+
+from kornia.color import AUTUMN, ApplyColorMap, apply_colormap
+from kornia.core import tensor
+from kornia.testing import BaseTester, assert_close, tensor_to_gradcheck_var
+
+
+def test_autumn(device, dtype):
+    cm = AUTUMN(num_colors=254, device=device, dtype=dtype)
+    colors = cm.colors
+
+    actual = colors[..., 0]
+
+    expected = tensor([1, 0, 0], device=device, dtype=dtype)
+
+    assert_close(actual, expected)
+
+    actual = colors[..., 127]
+    expected = tensor([1, 0.5019997358322144, 0], device=device, dtype=dtype)
+    assert_close(actual, expected)
+
+    actual = colors[..., -1]
+    expected = tensor([1, 1, 0], device=device, dtype=dtype)
+    assert_close(actual, expected)
+
+
+class TestApplyColorMap(BaseTester):
+    def test_smoke(self, device, dtype):
+        input_tensor = tensor([[[0, 1, 3], [25, 50, 63]]], device=device, dtype=dtype)
+
+        expected_tensor = tensor(
+            [
+                [[1, 1, 1], [1, 1, 1]],
+                [[0, 0.01587301587301587, 0.04761904761904762], [0.3968253968253968, 0.7936507936507936, 1]],
+                [[0, 0, 0], [0, 0, 0]],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+        cm = AUTUMN(device=device, dtype=dtype)
+        actual = apply_colormap(input_tensor, cm)
+
+        self.assert_close(actual, expected_tensor)
+
+    def test_exception(self, device, dtype):
+        cm = AUTUMN(device=device, dtype=dtype)
+        with pytest.raises(TypeError):
+            apply_colormap(torch.rand(size=(5, 1, 1), dtype=dtype, device=device), cm)
+
+    @pytest.mark.parametrize("shape", [(2, 1, 4, 4), (1, 4, 4), (4, 4)])
+    def test_cardinality(self, shape, device, dtype):
+        cm = AUTUMN(device=device, dtype=dtype)
+        input_tensor = torch.randint(0, 63, shape, device=device, dtype=dtype)
+        actual = apply_colormap(input_tensor, cm)
+
+        if len(shape) == 4:
+            expected_shape = (shape[0], 3, shape[-2], shape[-1])
+        else:
+            expected_shape = (3, shape[-2], shape[-1])
+
+        assert actual.shape == expected_shape
+
+    @pytest.mark.skip(reason='jacobian mismatch')
+    def test_gradcheck(self, device, dtype):
+        # TODO: implement differentiability
+        cm = AUTUMN(device=device, dtype=dtype)
+        input_tensor = torch.randint(0, 63, (1, 2, 1), device=device, dtype=dtype)
+
+        input_tensor = tensor_to_gradcheck_var(input_tensor)
+        cm.colors = tensor_to_gradcheck_var(cm.colors)
+
+        assert gradcheck(apply_colormap, (input_tensor, cm), raise_exception=True, fast_mode=True)
+
+    def test_jit(self, device, dtype):
+        op = apply_colormap
+        op_script = torch.jit.script(op)
+
+        cm = AUTUMN(device=device, dtype=dtype)
+        img = torch.ones(1, 3, 3, device=device, dtype=dtype)
+
+        self.assert_close(op(img, cm), op_script(img, cm))
+
+    def test_module(self, device, dtype):
+        op = apply_colormap
+        cm = AUTUMN(device=device, dtype=dtype)
+        op_module = ApplyColorMap(cm)
+
+        img = torch.ones(1, 3, 3, device=device, dtype=dtype)
+
+        self.assert_close(op(img, cm), op_module(img))