Converting a grayscale image to an RGB heatmap image with matplotlib [duplicate]

Question

How do I convert an M x N grayscale image, or in other words a matrix or 2-D array, into an RGB heatmap, or in other words an M x N x 3 array?

Example:

 [[0.9, 0.3], [0.2, 0.1]] 

should become

[[red, green-blue], [green-blue, blue]] 

where red is [1, 0, 0], blue is [0, 0, 1],etc.

Answer 1

import matplotlib.pyplot as plt

img = [[0.9, 0.3], [0.2, 0.1]]

cmap = plt.get_cmap('jet')

rgba_img = cmap(img)
rgb_img = np.delete(rgba_img, 3, 2)

cmap is an instance of matplotlib's LinearSegmentedColormap class, which is derived from the Colormap class. It works because of the __call__ function defined in Colormap. Here is the docstring from matplotlib's git repo for reference, since it's not described in the API.

def __call__(self, X, alpha=None, bytes=False):
    """
    *X* is either a scalar or an array (of any dimension).
    If scalar, a tuple of rgba values is returned, otherwise
    an array with the new shape = oldshape+(4,). If the X-values
    are integers, then they are used as indices into the array.
    If they are floating point, then they must be in the
    interval (0.0, 1.0).
    Alpha must be a scalar between 0 and 1, or None.
    If bytes is False, the rgba values will be floats on a
    0-1 scale; if True, they will be uint8, 0-255.
    """

A simpler option is to display img, using plt.imshow or plt.matshow, and then copy or save the result as an RGB or RGBA image. This was too slow for my application (~ 30 times slower on my machine).