如何将3D RGB标签图像(语义分割)转换为二维灰度图像，类索引从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]])

Answer 2

这个怎么样？

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]

我认为它是基于与公认的方法相同的想法，但它看起来更清晰。

Answer 3

我在这里也回答了这个问题：Convert RGB image to index image

基本上：

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

您可以根据需要扩展字典。