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

Question

简而言之，我想做的与这个问题类似：Convert RGB image to index image，但我想要获取 n 通道图像，而不是 1 通道索引图像，其中img[h, w]是 one-hot 编码向量。例如，如果输入图像为[[[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

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

[
    [[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 数组更改为 one-hot 编码值的 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

我们可以生成每个像素颜色的十进制当量。每个通道都具有0or255作为值，那么就有全部8可能性，但似乎我们只对其中四种颜色感兴趣。

那么，我们有两种方法来解决它：

• 一种方法是从这些十进制等值中创建唯一索引，从开始0到最终颜色，所有这些都按顺序进行，最后初始化一个输出数组并分配给它。

• 另一种方法是将broadcasted这些十进制等值与颜色进行比较。

接下来列出这两种方法 -

def indexing_based(a):
    b = (a == 255).dot([4,2,1])  # Decimal equivalents
    colors = np.array([7,4,1,6]) # Define colors decimal equivalents here
    idx = np.empty(colors.max()+1,dtype=int)
    idx[colors] = np.arange(len(colors))
    m,n,r = a.shape
    out = np.zeros((m,n,len(colors)), dtype=int)
    out[np.arange(m)[:,None], np.arange(n), idx[b]] = 1
    return out

def broadcasting_based(a):
    b = (a == 255).dot([4,2,1])  # Decimal equivalents
    colors = np.array([7,4,1,6]) # Define colors decimal equivalents here
    return (b[...,None] == colors).astype(int)

样本运行 -

>>> a = np.array([
...     [[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]],
...     [[255, 255, 255], [255,   0,   0], [255, 255,   0], [255,  0 ,   0]]])
>>> indexing_based(a)
array([[[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]],

       [[1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 0, 1],
        [0, 1, 0, 0]]])
>>> np.allclose(broadcasting_based(a), indexing_based(a))
True