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

【问题标题】：How to convert 3D RGB label image (in semantic segmentation) to 2D gray image, and class indices start from 0?如何将 3D RGB 标签图像（在语义分割中）转换为 2D 灰度图像，并且类索引从 0 开始？
【发布时间】：2019-04-03 04:38:53
【问题描述】：

我有一个 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 中广播。

【问题讨论】：

标签： python numpy array-broadcasting

【解决方案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]])

【讨论】：

【解决方案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]

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

【讨论】：

【解决方案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

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

【讨论】：

【解决方案4】：

我试过这个...

首先，我注意到在下面的 RGB 值表中，绿色值都是相同的，因此没有必要检查它们。

其次，如果你将数组中的值除以 255，你会得到非常接近你需要的标签的零和一。所以，如果你做一点数学：

t = R/255 + 2B/255 -1

然后你会得到字典中的值：

  R   G   B    t
==================
255 255   0    0
 0  255 255    1
255 255 255    2

与其他几个答案进行比较的代码如下所示：

#!/usr/bin/env python3

import numpy as np

def me(img): 
    """Return R + 2B - 1 as label"""
    return np.uint8((img[:,:,0]/255) + 2*(img[:,:,2]/255) - 1) 

def deepak(img):
    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.uint8)
    label_seg[(img==r).all(axis=2)] = 0
    label_seg[(img==g).all(axis=2)] = 1
    label_seg[(img==b).all(axis=2)] = 2
    return label_seg

def marios(label):
    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]

    return label

# Generate a sample image
img = np.zeros((480,640,3), dtype=np.uint8)
img[:160,:,:]    = [255,255,0]
img[160:320,:,:] = [0,255,255]
img[320:,:,:]    = [255,255,255]

时间是这样出来的：

In [134]: %timeit deepak(img)
15.4 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [135]: %timeit marios(img)
15.4 ms ± 166 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [172]: %timeit me(img)                                                                           
869 µs ± 8.93 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)each)

18 倍加速是否值得可读性较差的代码是另一个论点，尽管 cmets 可以提供帮助:-)

注意，为了公平起见，他的时间可以减少到 0 10.3 毫秒，方法是删除下面将零数组中的一些元素归零的不必要行：

label_seg[(img==r).all(axis=2)] = 0

【讨论】：