Mapping RGB values in an image to a corresponding ID using a dictionary

I am working on a segmentation problem where given an image, each RGB value corresponds to a class label. The problem I have is to efficiently map RGB values from an image (numpy array) to a corresponding class label image.

Let’s provide the following simplified example:

color2IdMap 
{(100,0,100):0, (0,200,0):2}

labelOld 
array([[[100,0,100],
         [0,200,0]],
       [[100,0,100],
        [0,200,0]]], dtype=uint8)

(in a real example the colorIdMap will have about 20 entries and labelOld will be an array of shape: (1024,512,3))

Now I want the result to be the following mapped array. with shape: (1024,512)

labelNew
array([[ 0, 2],
       [ 0, 2]])

I attempted to do this with loops and list comprehensions but both methods are quite slow (about ~10seconds per image, which is a big number for 250K images). And I am wondering if there is a faster way of doing it.

---

Attempted method 1:

labelNew = np.empty((1052,1914), dtype=np.uint8)
for i in range(1052):
    for j in range(1914):
        labelNew[i, j] = color2IdMap[tuple(labelOld[i, j])]

Attempted method 2:

labelNew = [[color2IdMap[tuple(x)] for x in y] for y in labelOld]

---

So, my question is if there is any faster and more efficient way of doing this?

Answer

Here’s one approach based on dimensionality-reduction –

# Get keys and values
k = np.array(list(color2IdMap.keys()))
v = np.array(list(color2IdMap.values()))

# Setup scale array for dimensionality reduction
s = 256**np.arange(3)

# Reduce k to 1D
k1D = k.dot(s)

# Get sorted k1D and correspondingly re-arrange the values array
sidx = k1D.argsort()
k1Ds = k1D[sidx]
vs = v[sidx]

# Reduce image to 2D
labelOld2D = np.tensordot(labelOld, s, axes=((-1),(-1)))

# Get the positions of 1D sorted keys and get the correspinding values by
# indexing into re-arranged values array
out = vs[np.searchsorted(k1Ds, labelOld2D)]

Alternatively, we could use sidx as sorter input arg for np.searchsorted to get the final output –

out = v[sidx[np.searchsorted(k1D, labelOld2D, sorter=sidx)]]