Map values to higher dimension with Numpy

Question

I'm trying to apply a color map to a two dimensional gray-scale image in Numpy (the image is loaded/generated by OpenCV).

I have a 256 entries long list with RGB values, which is my colormap:

cMap = [np.array([0, 0, 0], dtype=np.uint8),
        np.array([0, 1, 1], dtype=np.uint8),
        np.array([0, 0, 4], dtype=np.uint8),
        np.array([0, 0, 6], dtype=np.uint8),
        np.array([0, 1, 9], dtype=np.uint8),
        np.array([ 0,  0, 12], dtype=np.uint8),
        # Many more entries here
       ]

When I input a gray scale image (shape (y,x,1)), I would like to output a color image (shape (y,x,3)), where the gray scale value for each input pixel is used as an index in cMap to find the appropriate color for the output pixel. My feeble attempt so far (inspired by Fast replacement of values in a numpy array) looks like this:

colorImg = np.zeros(img.shape[:2] + (3,), dtype=np.uint8)

for k, v in enumerate(cMap):
    colorImg[img==k] = v

This gives me the error ValueError: array is not broadcastable to correct shape. I think I have narrowed the problem down: The slicing I do results in a 1D boolean array with an entry for each entry in img. colorImg, however has three times more entries than img, so the boolean array won't be long enough.

I have tried all kinds of reshaping and other slicings, but I'm pretty stuck now. I'm sure there is an elegant way to solve this :-)

Answer 1

The function you are looking for is numpy.take(). The only tricky bit is to specify axis = 0 so that the returned image has the correct number of channels.

colorImg = np.take(np.asarray(cMap), img, axis = 0)