A utility function for TensorFlow that maps a grayscale image to a matplotlib colormap for use with TensorBoard image summaries.

colorize.py

import matplotlib
import matplotlib.cm

import tensorflow as tf

def colorize(value, vmin=None, vmax=None, cmap=None):
    """
    A utility function for TensorFlow that maps a grayscale image to a matplotlib
    colormap for use with TensorBoard image summaries.

    By default it will normalize the input value to the range 0..1 before mapping
    to a grayscale colormap.

    Arguments:
      - value: 2D Tensor of shape [height, width] or 3D Tensor of shape
        [height, width, 1].
      - vmin: the minimum value of the range used for normalization.
        (Default: value minimum)
      - vmax: the maximum value of the range used for normalization.
        (Default: value maximum)
      - cmap: a valid cmap named for use with matplotlib's `get_cmap`.
        (Default: 'gray')

    Example usage:

    ```
    output = tf.random_uniform(shape=[256, 256, 1])
    output_color = colorize(output, vmin=0.0, vmax=1.0, cmap='viridis')
    tf.summary.image('output', output_color)
    ```
    
    Returns a 3D tensor of shape [height, width, 3].
    """

    # normalize
    vmin = tf.reduce_min(value) if vmin is None else vmin
    vmax = tf.reduce_max(value) if vmax is None else vmax
    value = (value - vmin) / (vmax - vmin) # vmin..vmax

    # squeeze last dim if it exists
    value = tf.squeeze(value)

    # quantize
    indices = tf.to_int32(tf.round(value * 255))

    # gather
    cm = matplotlib.cm.get_cmap(cmap if cmap is not None else 'gray')
    colors = tf.constant(cm.colors, dtype=tf.float32)
    value = tf.gather(colors, indices)

    return value

@DEKHTIARJonathan, the object returned by matplotlib.cm.get_cmap indeed has no attribute 'colors'; however, it may be called as a function accepting numpy arrays of grayscale values in [0, 255] (and returning pixel intensities in [0., 1.] for 4 channels (RGBA), not 3). So substituting line 48 of the gist with the following may work for you:

colors = cm(np.arange(256))[:, :3]
colors = tf.constant(colors, dtype=tf.float32)

Hope that helps.

This is very cool, thanks!
Something that would help out is if you make it compatible with the normal (batch_size, width, height, channels=1) input for keras / tensorflow. The output then would be (batch_size, width, height, 3), creating one image per batch input.
Does this require much change?

Edit: I added this line which effectively did what I was describing:
out = tensorflow.map_fn(lambda img: colorize(img, cmap='jet'), out)

Very helpful, thanks alot. Came here from Pytorch. This works with numpy/torch arrays and exports a (h,w,4)- int8 instead.

import matplotlib
import matplotlib.cm

def colorize(value, vmin=None, vmax=None, cmap=None):
    """
    A utility function for Torch/Numpy that maps a grayscale image to a matplotlib
    colormap for use with TensorBoard image summaries.
    By default it will normalize the input value to the range 0..1 before mapping
    to a grayscale colormap.
    Arguments:
      - value: 2D Tensor of shape [height, width] or 3D Tensor of shape
        [height, width, 1].
      - vmin: the minimum value of the range used for normalization.
        (Default: value minimum)
      - vmax: the maximum value of the range used for normalization.
        (Default: value maximum)
      - cmap: a valid cmap named for use with matplotlib's `get_cmap`.
        (Default: Matplotlib default colormap)
    
    Returns a 4D uint8 tensor of shape [height, width, 4].
    """

    # normalize
    vmin = value.min() if vmin is None else vmin
    vmax = value.max() if vmax is None else vmax
    if vmin!=vmax:
        value = (value - vmin) / (vmax - vmin) # vmin..vmax
    else:
        # Avoid 0-division
        value = value*0.
    # squeeze last dim if it exists
    value = value.squeeze()

    cmapper = matplotlib.cm.get_cmap(cmap)
    value = cmapper(value,bytes=True) # (nxmx4)
    return value

from tensorboardX import SummaryWriter
import torch
import numpy as np
# x = np.ones([100,100])
x = torch.ones([100,100])
colorized = colorize(x)

writer = SummaryWriter('runs/imagetest/')
writer.add_image('TEST', colorized, 1)

when you remove line 41 value = tf.squeeze(value) it works on arbitrary input shapes and returns [original_shape, ] + [3]. Could be useful to write whole batches.