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

Question

简单地说，我正在尝试做的是类似于这个问题： 将RGB图像转换为索引图像 ，但不是单通道索引图像，我想得到n通道图像，其中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

，我有一个由这四种颜色组成的图像，其中图像是一个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索引值数组更改为单个编码值的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

我们可以生成每个像素颜色的十进制等值。 每个通道的值为0或255 ，总共有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

Answer 2

我的解决方案看起来像这样，应该适用于任意颜色：

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)

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