如何将3D RGB标签图像（在语义分割中）转换为2D灰度图像，并且类索引从0开始？

Question

我有一个rgb语义分段标签，如果其中有3个类，并且每个RGB值是以下之一：

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

然后，我想根据dict将RGB文件中的所有值映射到新的2D标签图像中：

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

之后，新的灰色标签文件中的所有值都是0、1或2之一。是否有解决此问题的有效方法？例如在NumPy中 广播 。

Answer 1

你可以这样做：

# 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

因此，如果

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

然后，

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