如何有效地向量化(即避免显式循环)将某些颜色映射到Numpy数组中的其他颜色

Question

我有一个三级Numpy矩阵作为输入(即图像:水平、垂直和4个彩色通道)。我想阅读这个矩阵元素在他们的前两个指数，并只映射到其他颜色，定义在各自的数组。性能是非常重要的，因为这种映射将被多次应用，可能会成为程序的瓶颈。

确切地说，我到目前为止掌握的代码是：

# data is the rank-3 Numpy array with the image (obtained using the library PIL)
# palette has shape (11, 4) and defines the 11 colours to map
# palette_grey has shape (11,) and defines the 11 tones of grey to apply

for i in range(palette.shape[0]):  # loop in colours
    match = (data[:,:,0] == palette[i,0]) & (data[:,:,1] == palette[i,1]) & (data[:,:,2] == palette[i,2]) # build matrix only True when a given pixel has the right color
    for j in range(3): # loop to apply the mapping to the three channels (because it's just grey, so all channels are equal)
        data[:,:,j] = np.where(match, grey_palette[i], data[:,:,j]) # the mapping itself

虽然主要任务是矢量化的(通过np.where)，但为了提高性能，我仍然希望避免两个显式循环。

有什么办法做到这一点吗？

编辑：

我已经尝试删除第二个循环(在通道中)，通过定义这两个调色板具有相同的形状(11,4)。然后，我试过这样做：

for i in range(palette.shape[0]):
    match = (data[:,:,0] == palette[i,0]) & (data[:,:,1] == palette[i,1]) & (data[:,:,2] == palette[i,2])
    data[:,:,:] = np.where(match, grey_palette[i], data[:,:,:])

但是它引发了一个错误：

ValueError:操作数不能与形状(480,480) (4 ) (480,480,4)一起广播

我想这是预期的行为，但我认为我提议的映射是明确的，因此可以由Numpy完成。

Answer 1

假设您有一个(M, N, 4)数组的uint8值。您可以将其视为32位整数的数组，如下所示：

idata = data.view(np.uint32)

其次，您可以将灰色映射从(K, 4)数组转换为适当的(K,)数组：

grey = np.repeat(grey[:, None], 4, axis=-1)

当您完成它时，可以将它转换为整数和调色板：

ipalette = palette.view(np.uint32).squeeze()
igrey = grey.view(np.uint32).squeeze()

现在您有了一个(M, N, 1)图像、(K,)调色板和灰色地图。对于足够小的K，您可以这样做：

ind = np.nonzero(idata == ipalette)
idata[ind[0], ind[1], 0] = igrey[ind[-1]]

这将直接将灰色值写入原始的data数组。掩码idata == ipalette向(M, N, K)广播。np.nonzero沿着每个轴返回索引。前两个索引是图像中匹配的坐标，而第三个索引恰好是您想要的灰色值。

这也适用于uint16图像，因为您可以使用uint64来聚合数据。

对于更大的K，您可以采取完全不同的方法：

1. 通过np.unique和return_inverse=True来传递idata。这将为您提供一个排序的唯一值数组，以及将其替换为图像顺序的索引。

就像这样：

uniques, reverse_index = np.unique(idata.ravel(), return_inverse=True)
palette_index = np.searchsorted(uniques, ipalette)
# clipping won't affect the match but avoids errors
palette_mask = (ipalette == uniques[palette_index.clip(0, uniques.size - 1)])
uniques[palette_index[palette_mask]] = igrey[palette_mask]
new_data = uniques[reverse_index].view(np.uint8).reshape(data.shape)

下面是第一种方法的完整示例：

np.random.seed(0)
data = np.random.randint(0, 255, size=(100, 100, 4), dtype=np.uint8)
palette = np.array([[117, 166,  22, 183],
                    [ 38,  28,  93, 140],
                    [ 63, 214,   9,  84],
                    [185,  51,   1,  24],
                    [131, 210, 145,   3],
                    [111, 180, 165, 245],
                    [ 62, 220, 102, 144],
                    [177,  97, 158, 135],
                    [202,  67, 169,  10],
                    [ 23, 177,  26,  15],
                    [ 19, 100,  25,  66],
                    [ 86, 227, 222, 182],
                    [255, 255, 255, 255]], dtype=np.uint8)
grey = np.arange(10, 10 + len(palette), dtype=np.uint8)

idata = data.view(np.uint32)
ipalette = palette.view(np.uint32).squeeze()
igrey = np.repeat(grey[:, None], data.shape[-1], axis=-1).view(np.uint32).squeeze()

ind = np.nonzero(idata == ipalette)
idata[ind[0], ind[1], 0] = igrey[ind[-1]]

我选择了调色板作为在执行>>> data时使用默认设置显示的像素，这样您就可以立即看到该方法是否有效。

使用第二种方法，您可以看到裁剪使调色板中的值0xFFFFFF成为可能，它获得和插入索引大于数组的大小，而无需崩溃或引入任何特殊情况。

Answer 2

当我比较你的解决方案和这个方案时：

for i in range(palette.shape[0]):
    new_data[data == palette[i]] = grey_palette[i]

在笔记本中使用%%timeit，为您的1000 x1000x3 data提供87 vs和218 vs。

编辑:删除关于您的解决方案的“问题”的评论，我只在一个地方更改为new_data。

Answer 3

从调色板中创建一个字典d {palette[i]: grey_palette[i]} (由普通条目{i:i}完成)，将其矢量化并应用于您的数据

numpy.vectorize(d.get)(data)

它应该是快速的，但我还没有用您的数据类型进行测试。