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

How to efficiently vectorize (i.e. avoid explicit loops) a mapping of some colours into others in a Numpy array

我输入了一个三阶 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: operands could not be broadcast together with shapes (480,480) (4,) (480,480,4)

我想这是预期的行为,但我认为我提出的映射是明确的,因此可以通过 Numpy 实现。

当我将你的解决方案与这个解决方案进行比较时:

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

在笔记本中使用 %%timeit 可以得到 87 毫秒,而对于 1000x1000x3 data.

,您的笔记本可以得到 218 毫秒

编辑:删除了关于 'problem' 的评论以及我通过仅在一个地方更改为 new_data 创建的解决方案。

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

idata = data.view(np.uint32)

其次,您可以将灰度图从 (K,) 数组转换为适当的 (K, 4) 数组:

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 returns 沿每个轴的索引。前两个索引是图像中匹配的坐标,而第三个索引正是你想要的灰度值。

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

对于更大的 K,您可以采用完全不同的方法:

  1. 通过 np.uniquereturn_inverse=True 通过 idata。这将为您提供一组已排序的唯一值,以及一个将它们替换为图像顺序的索引。
  2. 运行 np.searchsortedipalette 放入唯一值中。
  3. 丢弃任何不代表完全匹配的索引。
  4. 用灰度值替换完全匹配。
  5. 使用灰色替换从唯一值重建图像。

像这样:

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 时使用默认设置显示的像素,这样您就可以立即看到该方法有效。

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

创建字典 d 从您的调色板构建 {palette[i]: grey_palette[i]}(由琐碎的条目完成 {i:i} 否则),对其进行矢量化并应用于您的数据

numpy.vectorize(d.get)(data)

它应该很快,但我没有用你的数据类型测试它。