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

Question

我有一个rgb语义分割标签，如果里面有3个类，每个RGB值都是其中之一：

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

分别，然后我想根据dict将RGB文件中的所有值映射到一个新的二维标签图像中：

{(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

我试过这个...

首先，我注意到在下面的 table 个 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 倍的加速是否值得可读性较差的代码是另一个论点，尽管注释可以提供帮助:-)

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

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

Answer 4

我也在这里回答了这个问题：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

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

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

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

python

numpy

array-broadcasting