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)

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.