将RGB图像转换为索引图像

Question

我想用Python将3通道RGB图像转换为索引图像。它用于处理训练深层网络进行语义分割的标签。我的意思是它有一个通道，每个像素都是索引，应该从零开始。当然，它们应该有相同的尺寸。转换基于Python dict中的以下映射：

color2index = {
        (255, 255, 255) : 0,
        (0,     0, 255) : 1,
        (0,   255, 255) : 2,
        (0,   255,   0) : 3,
        (255, 255,   0) : 4,
        (255,   0,   0) : 5
    }

我实现了一个天真的函数：

^{pr2}$

输入im是由cv2.imread()创建的一个numpy数组。但是，这个代码非常慢。 由于im在numpy数组中，我首先用如下方法尝试了numpy的ufunc：

RGB2index = np.frompyfunc(lambda x: color2index(tuple(x)))
indices = RGB2index(im)

但是结果证明ufunc每次只需要一个元素。我不能给函数一次三个参数（RGB值）。在

那么有没有其他的方法来优化呢？ 如果存在更有效的数据结构，映射就不必是那样的。我注意到Python dict的访问不会花费太多时间，但是从numpy array转换到tuple（它是散列的）需要花费很多时间。在

附言： 我的一个想法是在CUDA中实现一个内核。但这会更复杂。在

更新1:Dan Mašek's Answer效果不错。但首先我们要把RGB图像转换成灰度。当两种颜色的灰度值相同时，可能会出现问题。在

我把工作代码贴在这里。希望它能帮助别人。在

lut = np.ones(256, dtype=np.uint8) * 255
lut[[255,29,179,150,226,76]] = np.arange(6, dtype=np.uint8)
im_out = cv2.LUT(cv2.cvtColor(im, cv2.COLOR_BGR2GRAY), lut)

Answer 1

这里有一个小的实用函数来转换图像(np.数组)到每像素标签（索引），也可以是一个热编码：

def rgb2label(img, color_codes = None, one_hot_encode=False):
    if color_codes is None:
        color_codes = {val:i for i,val in enumerate(set( tuple(v) for m2d in img for v in m2d ))}
    n_labels = len(color_codes)
    result = np.ndarray(shape=img.shape[:2], dtype=int)
    result[:,:] = -1
    for rgb, idx in color_codes.items():
        result[(img==rgb).all(2)] = idx

    if one_hot_encode:
        one_hot_labels = np.zeros((img.shape[0],img.shape[1],n_labels))
        # one-hot encoding
        for c in range(n_labels):
            one_hot_labels[: , : , c ] = (result == c ).astype(int)
        result = one_hot_labels

    return result, color_codes


img = cv2.imread("input_rgb_for_labels.png")
img_labels, color_codes = rgb2label(img)
print(color_codes) # e.g. to see what the codebook is

img1 = cv2.imread("another_rgb_for_labels.png")
img1_labels, _ = rgb2label(img1, color_codes) # use the same codebook

如果提供了None，它将计算（并返回）颜色码本。在

Answer 2

实际上for循环需要很多时间。在

binary_mask = (im_array[:,:,0] == 255) & (im_array[:,:,1] == 255) & (im_array[:,:,2] == 0) 

也许上面的代码可以帮助你

Answer 3

你查过枕头库了吗？我记得，它有一些类和方法来处理颜色转换。请参见：https://pillow.readthedocs.io/en/4.0.x/reference/Image.html#PIL.Image.Image.convert