Custom Colormap using Matplotlib.image

Question

I am using matplotlib.image.imsave('file.png',file,cmap=cmap) to save a .png of a 2d Numpy array where the array can only have values of 0, 1 or 10. I would like 0 to be white, 1 to be green and 10 to be red. I saw something similar at this question: Matplotlib: Custom colormap with three colors . The problem is that imsave does not take norm as an argument but using pyplot is too slow for my application. Any assistance would be appreciated!

Answer 1

The input array consisting of values [0,1,10] is not really an image array. Image arrays would go from 0 to 255 or from 0. to 1..

a. using a LinearSegmentedColormap

An idea can be to normalize your array im to 1.: im = im/im.max(). It is then possible to create a colormap with the values 0 -> white, 0.1 -> green, 1 -> red using matplotlib.colors.LinearSegmentedColormap.from_list.

import matplotlib.image
import numpy as np

im = np.random.choice([0,1,10], size=(90, 90), p=[0.5,0.3,0.2])

im2 = im/10.
clist = [(0,"white"), (1./10.,"green"), (1, "red")]
cmap = matplotlib.colors.LinearSegmentedColormap.from_list("name", clist)
matplotlib.image.imsave(__file__+'.png', im, cmap=cmap)

The corresponding pyplot plot

import matplotlib.pyplot as plt
plt.imshow(im, cmap=cmap)
plt.colorbar(ticks=[0,1,10])
plt.show()

would look like this

b. using a ListedColormap

A ListedColormap can be used to generate a colormap of the three colors white, green and red. In this colormap the colors are equally spaced, so one needs to map the image array to equally spaced values as well. This can be done using np.unique(im,return_inverse=True)[1].reshape(im.shape), which returns an array containing only the values [0,1,2]. We again need to normalize to 1.

im = np.random.choice([0,1,10], size=(90, 90), p=[0.5,0.3,0.2])

im2 = np.unique(im,return_inverse=True)[1].reshape(im.shape)
im3 = im2/float(im2.max())

clist = ["white", "green","red"]
cmap = matplotlib.colors.ListedColormap(clist)
matplotlib.image.imsave(__file__+'2.png',im3, cmap=cmap) 

While the output image looks exactly the same as above, the corresponding matplotlib plot would have a different colorbar.

import matplotlib.pyplot as plt
plt.imshow(im2, cmap=cmap)
cb = plt.colorbar(ticks=[0,1,2])
cb.ax.set_yticklabels([0,1,10])
plt.show()

Answer 2

Just build a (N, M, 3) array and treat it as image-pixels in RGB-mode. Then it's enough to map your 3 unique values to those 3 colors.

Code:

import numpy as np
from scipy.misc import imsave

raw = np.random.choice((0,1,10), size=(500, 500))
img = np.empty((raw.shape[0], raw.shape[1], 3))

img[raw == 0] = (255, 255, 255)  # RGB -> white
img[raw == 1] = (0,255,0)        #        green
img[raw == 10] = (255,0,0)       #        red

imsave('image.png', img)

I'm using scipy's imsave here, but the matplotlib one probably works the same.