python - 如何使用 numpy 将 RGB 图像转换为基于颜色的单热编码 3d 数组？

Question

简单地说，我想做的是类似于这个问题:Convert RGB image to index image ，但我想获取 n channel 图像，而不是 1 channel 索引图像，其中 img[h, w] 是一个单热编码向量。例如，如果输入图像是 [[[0, 0, 0], [255, 255, 255]]，索引 0 分配给黑色，索引 1 分配给白色，那么所需的输出是 [[[1, 0], [0, 1]]]。

就像前面那个人问的问题一样，我很天真地实现了这个，但是代码运行得相当慢，我相信使用 numpy 的适当解决方案会快得多。

此外，正如上一篇文章中所建议的，我可以将每张图像预处理为灰度图像，然后对图像进行单热编码，但我想要一个更通用的解决方案。

例子

假设我想将白色分配给 0，将红色分配给 1，将蓝色分配给 2，将黄色分配给 3:

(255, 255, 255): 0
(255, 0, 0): 1
(0, 0, 255): 2
(255, 255, 0): 3

，我有一个由这四种颜色组成的图像，其中图像是一个包含每个像素的 R、G、B 值的 3D 数组:

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

，这就是我想要将每个像素更改为索引的单热编码值的地方。 (由于将 2d 索引值数组更改为 3d 单热编码值数组很容易，因此获取 2d 索引值数组也很好。)

[
    [[1, 0, 0, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 1, 0, 0]],
    [[0, 0, 1, 0], [1, 0, 0, 0], [0, 1, 0, 0], [0, 1, 0, 0]],
    [[0, 0, 1, 0], [0, 0, 1, 0], [1, 0, 0, 0], [1, 0, 0, 0]],
    [[1, 0, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 0, 1]]
]

在此示例中，我使用了 RGB 分量为 255 或 0 的颜色，但我不希望解决方案依赖于该事实。

Answer 1

我的解决方案如下所示，适用于任意颜色:

color_dict = {0: (0,   255, 255),
              1: (255, 255,   0),
              ....}


def rgb_to_onehot(rgb_arr, color_dict):
    num_classes = len(color_dict)
    shape = rgb_arr.shape[:2]+(num_classes,)
    arr = np.zeros( shape, dtype=np.int8 )
    for i, cls in enumerate(color_dict):
        arr[:,:,i] = np.all(rgb_arr.reshape( (-1,3) ) == color_dict[i], axis=1).reshape(shape[:2])
    return arr


def onehot_to_rgb(onehot, color_dict):
    single_layer = np.argmax(onehot, axis=-1)
    output = np.zeros( onehot.shape[:2]+(3,) )
    for k in color_dict.keys():
        output[single_layer==k] = color_dict[k]
    return np.uint8(output)

我还没有测试它的速度，但至少，它是有效的:)