python - Convert multi-dimensional Numpy array to 2-dimensional array based on color values

Question

I have an image which is read as a uint8 array with the shape (512,512,3). Now I would like to convert this array to a uint8 array of shape (512,512,1), where each pixel value in the third axis are converted from a color value [255,0,0] to a single class label value [3], based on the following color/class encoding:

1 : [0, 0, 0], 
 2 : [0, 0, 255], 
 3 : [255, 0, 0], 
 4 : [150, 30, 150], 
 5 : [255, 65, 255], 
 6 : [150, 80, 0], 
 7 : [170, 120, 65], 
 8 : [125, 125, 125], 
 9 : [255, 255, 0], 
 10 : [0, 255, 255], 
 11 : [255, 150, 0], 
 12 : [255, 225, 120], 
 13 : [255, 125, 125], 
 14 : [200, 100, 100], 
 15 : [0, 255, 0], 
 16 : [0, 150, 80], 
 17 : [215, 175, 125], 
 18 : [220, 180, 210], 
 19 : [125, 125, 255]

What is the most efficient way to do this? I thought of looping through all classes and using numpy.where, but this is obviously time-consuming.

Answer 1

You could use giant lookup table. Let cls be [[0,0,0], [0,0,255], ...] of dtype=np.uint8.

LUT = np.zeros(size=(256,256,256), dtype='u1')
LUT[cls[:,0],cls[:,1],cls[:,2]] = np.arange(cls.shape[1])+1
img_as_cls = LUT[img[...,0],img[...,1], img[...,2]]

This solution is O(1) per pixel. It is also quite cache efficient because a small part of entries in LUT are actually used. It takes circa 10ms to process 1000x1000 image on my machine.

The solution can be slightly improved by converting 3-color channels to 24-bit integers. Here is the code

def scalarize(x):
    # compute x[...,2]*65536+x[...,1]*256+x[...,0] in efficient way
    y = x[...,2].astype('u4')
    y <<= 8
    y +=x[...,1]
    y <<= 8
    y += x[...,0]
    return y
LUT = np.zeros(2**24, dtype='u1')
LUT[scalarize(cls)] = 1 + np.arange(cls.shape[0])
simg = scalarize(img)
img_to_cls = LUT[simg]

After optimization it takes about 5ms to process 1000x1000 image.