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

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

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.

Solution

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]])