How to convert 3D RGB label image (in semantic segmentation) to 2D gray image, and class indices start from 0?

Question

I have a rgb semantic segmentation label, if there exists 3 classes in it, and each RGB value is one of:

[255, 255, 0], [0, 255, 255], [255, 255, 255]

respectively, then I want to map all values in RGB file into a new 2D label image according to the dict:

{(255, 255, 0): 0, (0, 255, 255): 1, (255, 255, 255): 2}

after that, all values in the new gray label file is one of 0, 1 or 2. Is there an efficient way to solve this problem? For example broadcasting in NumPy.

Answer 1

You can do this:

# the three channels
r = np.array([255, 255, 0])
g = np.array([0, 255, 255])
b = np.array([255, 255, 255])

label_seg = np.zeros((img.shape[:2]), dtype=np.int)
label_seg[(img==r).all(axis=2)] = 0
label_seg[(img==g).all(axis=2)] = 1
label_seg[(img==b).all(axis=2)] = 2

So that, if

img = np.array([[r,g,b],[r,r,r],[b,g,r],[b,g,r]])

then,

label_seg = array([[0, 1, 2],
                   [0, 0, 0],
                   [2, 1, 0],
                   [2, 1, 0]])

Answer 2

How about this one:

mask_mapping = {
    (255, 255, 0):   0,
    (0, 255, 255):   1,
    (255, 255, 255): 2,
}
for k in mask_mapping:
    label[(label == k).all(axis=2)] = mask_mapping[k]

I think it's based on the same idea as the accepted method, but it looks more clear.

Answer 3

I've also answered this question here: Convert RGB image to index image

Basically:

cmap = {(255, 255, 0): 0, (0, 255, 255): 1, (255, 255, 255): 2}

def rgb2mask(img):

    assert len(img.shape) == 3
    height, width, ch = img.shape
    assert ch == 3

    W = np.power(256, [[0],[1],[2]])

    img_id = img.dot(W).squeeze(-1) 
    values = np.unique(img_id)

    mask = np.zeros(img_id.shape)

    for c in enumerate(values):
        try:
            mask[img_id==c] = cmap[tuple(img[img_id==c][0])] 
        except:
            pass
    return mask

You can extend extend the dictionary as you want.