将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

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