将 RGB 图像转换为索引图像

Question

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

\n\n

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

\n\n

我已经实现了一个简单的功能：

\n\n

def im2index(im):\n    """\n    turn a 3 channel RGB image to 1 channel index image\n    """\n    assert len(im.shape) == 3\n    height, width, ch = im.shape\n    assert ch == 3\n    m_lable = np.zeros((height, width, 1), dtype=np.uint8)\n    for w in range(width):\n        for h in range(height):\n            b, g, r = im[h, w, :]\n            m_lable[h, w, :] = color2index[(r, g, b)]\n    return m_lable\n

\n\n

输入im是由 . 创建的numpycv2.imread()数组。然而，这段代码真的很慢。\n由于它im位于 numpy 数组中，所以我首先尝试使用如下所示的numpyufunc：

\n\n

RGB2index = np.frompyfunc(lambda x: color2index(tuple(x)))\nindices = RGB2index(im)\n

\n\n

但事实证明ufunc每次只需要一个元素。我无法一次为函数提​​供三个参数（RGB 值）。

\n\n

那么还有其他方法可以进行优化吗？\n如果存在更有效的数据结构，则映射不必是这样的。我注意到 Python 字典的访问不会花费太多时间，但从numpy 数组到元组（可散列）的转换却需要花费很多时间。

\n\n

PS：\n我的一个想法是在 CUDA 中实现内核。但情况会更复杂。

\n\n

UPDATA1： \n Dan Ma\xc5\xa1ek\ 的答案效果很好。但首先我们必须将 RGB 图像转换为灰度图像。当两种颜色具有相同的灰度值时可能会出现问题。

\n\n

我将工作代码粘贴在这里。希望它可以帮助其他人。

\n\n

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

\n

Answer 1

您检查过 Pillow 库https://python-pillow.org/吗？我记得它有一些类和方法来处理颜色转换。请参阅：https ://pillow.readthedocs.io/en/4.0.x/reference/Image.html#PIL.Image.Image.convert