Update apply_colormap

[vgilabert94]

🚀 Feature -> Update apply_colormap

From #2794 (comment)

Discussion about ColorMaps -> apply_colormap function. (

kornia/kornia/color/colormap.py

Line 161 in 8df6c44

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

)

Motivation

• Right now the input tensor should be in uint8 (0-255) which not follow the float32 standard of kornia (0-1).
• Only grayscale images are accepted, we also can update to any channels. If input is with channel 3, output will be 9 (3 per every input channel).
• Try to accept batched tensors (batch >= 1).

Pitch

• Use image in float32 in range (0,1).
• Accept images with no restriction channels.

[edgarriba]

how are you planning to implement the LUT with f32 ?

[vgilabert94]

I've create two "premilinar" versions that allow to use batch > 1 and channels > 1 and range [0,1] in float:

def apply_colormap_v1(input_tensor, cmap):
    B, C, H, W = input_tensor.shape
    colormap = cmap.colors.permute(1, 0)
    num_colors, channels_cmap = colormap.shape
    input_tensor = input_tensor.reshape(B, C, -1)

    keys = torch.linspace(0.0, 1.0, num_colors-1, device=input_tensor.device)
    index = torch.bucketize(input_tensor, keys)
    output = colormap[index]
    # (B, C, H*W, channels_cmap) -> (B, C*channels_cmap, H, W)
    output = output.permute(0, 1, 3, 2).reshape(B, C * channels_cmap, H, W)

    return output


def apply_colormap_v2(input_tensor, cmap):
    B, C, H, W = input_tensor.shape
    cmap = cmap.colors.permute(1, 0)
    num_colors, channels_cmap = cmap.shape
    input_tensor = input_tensor.reshape(B, C, -1)

    keys = torch.linspace(0.0, 1.0, num_colors-1, device=input_tensor.device)
    index = torch.bucketize(input_tensor, keys).unsqueeze(-1).expand(-1, -1, -1, 3)
    output = torch.gather(cmap.expand(B, C, -1, -1), 2, index)
    # (B, C, H*W, channels_cmap) -> (B, C*channels_cmap, H, W)
    output = output.permute(0, 1, 3, 2).reshape(B, C*channels_cmap, H, W)

    return output

And also I create a a benchmark comparing both version with implemented in kornia (only uint8, grayscale images and batch==1). Results:

[------------------------- cpu --------------------------]
                      |   kornia  |     v1     |     v2   
1 threads: -----------------------------------------------
      1-128-float32   |   1020.5  |    1025.7  |     728.6
      1-512-float32   |  15653.2  |   23636.0  |   12356.3
      1-1024-float32  |  68119.8  |   65416.5  |   55136.0
      2-128-float32   |           |    3853.9  |    2427.8
      2-512-float32   |           |   34865.0  |   23163.8
      2-1024-float32  |           |  225147.9  |  174052.5
      5-128-float32   |           |    4790.0  |    3578.2
      5-512-float32   |           |   90941.5  |   60986.6
      5-1024-float32  |           |  557343.1  |  431017.7
      1-128-float64   |   1085.0  |    1984.4  |    1342.8
      1-512-float64   |  16091.4  |   19510.6  |   13638.2
      1-1024-float64  |  66923.7  |  110457.9  |   95593.2
      2-128-float64   |           |    2274.8  |    1593.7
      2-512-float64   |           |   36912.1  |   33436.7
      2-1024-float64  |           |  278575.3  |  236258.5
      5-128-float64   |           |    5420.1  |    4192.2
      5-512-float64   |           |  166984.2  |  146667.9
      5-1024-float64  |           |  738199.4  |  620192.5

Times are in microseconds (us).

[----------------------- cuda ----------------------]
                      |  kornia  |    v1    |    v2  
1 threads: ------------------------------------------
      1-128-float32   |   63.1   |    78.1  |    85.4
      1-512-float32   |   79.6   |   283.7  |   258.4
      1-1024-float32  |  294.9   |   832.3  |   762.7
      2-128-float32   |          |    80.1  |    89.0
      2-512-float32   |          |   386.8  |   384.2
      2-1024-float32  |          |  1510.2  |  1490.0
      5-128-float32   |          |    78.5  |    91.5
      5-512-float32   |          |   936.8  |   932.9
      5-1024-float32  |          |  3770.8  |  3732.2
      1-128-float64   |   59.6   |    76.4  |    86.2
      1-512-float64   |   61.8   |   334.9  |   334.5
      1-1024-float64  |  273.2   |  1297.5  |  1295.6
      2-128-float64   |          |    81.4  |    84.5
      2-512-float64   |          |   658.7  |   659.7
      2-1024-float64  |          |  2575.0  |  2574.4
      5-128-float64   |          |    97.4  |    88.4
      5-512-float64   |          |  1615.8  |  1613.9
      5-1024-float64  |          |  6431.6  |  6433.7

Times are in microseconds (us).

In my opinion v2 could be the best option.
Full benchmark code: https://colab.research.google.com/drive/1ciSwdxMWooYWSzlJk4B_P180O84E2BVf?usp=sharing

[vgilabert94]

In this benchmark, I have noticed that if the number of colors in the ColorMap is low (N=8), the current version may give some possibly wrong results. With the new versions this seems to be fixed.

N=256

N=8

[johnnv1]

Pretty nice cases, I liked it... about the supporting C > 1, the unique use case I think about is having one hot-coded tensor (categorical), do you see another use case?

on top of that, I would like to keep support for uint tensors, so that we can use the function to generate colored images from categorical masks

[vgilabert94]

I add channel support because you was talking about it (#2794 (comment)), but I am not sure what the other use cases might be.

Yes, we can keep support for uint tensors, simply by checking the input type and creating the keys with torch.linspace or torch.arange.