将 RGB 图像转换为索引图像答案

【问题标题】：Convert RGB image to index image将 RGB 图像转换为索引图像
【发布时间】：2017-08-02 17:14:09
【问题描述】：

我想用 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
    }

我实现了一个简单的函数：

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

输入im 是一个由cv2.imread() 创建的numpy 数组。但是，这段代码真的很慢。由于 im 在 numpy 数组中，我首先尝试了 numpy 的 ufunc，如下所示：

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

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

那么还有其他方法可以进行优化吗？如果存在更有效的数据结构，则映射不必如此。我注意到 Python dict 的访问不会花费太多时间，但是从 numpy 数组 到 tuple（它是可散列的）的转换确实如此。

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

UPDATA1： 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)

【问题讨论】：

所以输入将只包含您列出的那 6 种不同的颜色？如果是这样，从 RBG 到灰度的转换将为您提供以下灰度值：[255,29,179,150,226,76] - 6 个不同的值。然后通过 cv2.LUT 运行它，将其重新映射到 0-5。
更改顺序或 for 循环，即使这样也会加快您的代码速度
@DanMašek 谢谢！您的解决方案工作正常。在我没有意识到转换为灰度时 RGB 的权重不同之前。灰度图像的体积为0-255。这意味着最大类数是 256。不过，在大多数情况下都可以。问题可能是某些颜色可能具有相同的灰度值。
@smttsp 感谢您的评论。但这无济于事
@DanMašek 没错。这就是我要找的。我试图实现一个 numpy ufunc。但它只需要一个元素一次数组。如果OpenCV提供接口让我们自定义convert函数就好了。

标签： python opencv numpy deep-learning

【解决方案1】：

这个怎么样？

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
}

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 i, c in enumerate(values):
        try:
            mask[img_id==c] = color2index[tuple(img[img_id==c][0])] 
        except:
            pass
    return mask

然后只需调用：

mask = rgb2mask(ing)

【讨论】：

【解决方案2】：

其实for循环需要很多时间。

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

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

【讨论】：

【解决方案3】：

我实现了一个简单的功能：...... 我首先尝试了 numpy 的 ufunc ，如下所示：...

我建议使用一个更简单的函数，它只转换一个像素：

def rgb2index(rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    return color2index[tuple(rgb)]

然后使用 numpy 例程是个好主意，但我们不需要ufunc：

np.apply_along_axis(rgb2index, 2, im)

这里numpy.apply_along_axis() 用于将我们的rgb2index() 函数应用于整个图像im 沿三个轴（0、1、2）中最后一个轴的RGB 切片。

我们甚至可以不使用该函数而只写：

np.apply_along_axis(lambda rgb: color2index[tuple(rgb)], 2, im)

【讨论】：

【解决方案4】：

类似于 Armali 和 Mendrika 的提议，我不得不稍微调整一下以使其正常工作（也许完全是我的错）。所以我只想分享一个有效的sn-p。

COLORS = np.array([
    [0, 0, 0],
    [0, 0, 255],
    [255, 0, 0]
])
W = np.power(255, [0, 1, 2])

HASHES = np.sum(W * COLORS, axis=-1)
HASH2COLOR = {h : c for h, c in zip(HASHES, COLORS)}
HASH2IDX = {h: i for i, h in enumerate(HASHES)}


def rgb2index(segmentation_rgb):
    """
    turn a 3 channel RGB color to 1 channel index color
    """
    s_shape = segmentation_rgb.shape
    s_hashes = np.sum(W * segmentation_rgb, axis=-1)
    func = lambda x: HASH2IDX[int(x)]
    segmentation_idx = np.apply_along_axis(func, 0, s_hashes.reshape((1, -1)))
    segmentation_idx = segmentation_idx.reshape(s_shape[:2])
    return segmentation_idx

segmentation = np.array([[0, 0, 0], [0, 0, 255], [255, 0, 0]] * 3).reshape((3, 3, 3))
rgb2index(segmentation)

Example plot

代码也可以在这里找到： https://github.com/theRealSuperMario/supermariopy/blob/dev/scripts/rgb2labels.py

【讨论】：

【解决方案5】：

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

【讨论】：

我的问题不是从一个颜色空间到另一个颜色空间的正常颜色转换，而是将颜色映射到索引。不过谢谢你的回答。
我也在寻找同样的东西。但是在 rgb 和 index 之间找不到很好的转换。

【解决方案6】：

这是一个将图像 (np.array) 转换为每像素标签（索引）的小实用函数，它也可以是 one-hot 编码：

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，它会计算（并返回）颜色码本。

【讨论】：

【解决方案7】：

如果您对使用 MATLAB 感到满意 - 可以将结果保存为 *.mat 并使用 scipy.io.loadmat 加载 - MATLAB 中有 rgb2ind 函数，它完全符合您的要求。如果没有，它可以作为 Python 中类似实现的灵感。

【讨论】：