通过numpy功能用python循环优化或替换数组迭代

Question

我有以下代码可以按预期工作，但我很好奇循环是否可以用本机 numpy 函数/方法替换以获得更好的性能。我所拥有的是一个保存 RGB 值的数组，我用作查找表和两个保存灰度值（0-255）的二维数组。这两个数组的每个值对应查找表的一个轴的值​​。

如前所述，真正好的是摆脱 python 中的（慢）循环并使用更快的 numpy 方法。

#!/usr/bin/env python3
from PIL import Image
import numpy as np

dim = (2000, 2000)
rows, cols = dim

# holding a 256x256 RGB color lookup table
color_map = np.random.random_integers(0, 255, (256,256,3))
# image 1 greyscale values
color_map_idx_row = np.random.randint(0, 255, dim)
# image 2 greyscale values
color_map_idx_col = np.random.randint(0, 255, dim)

# output image data
result_data = np.zeros((rows, cols, 3), dtype=np.uint8)

# is there any built in function in numpy that could
# replace this loop?
# -------------------------------------------------------

for i in range(rows):
    for j in range(cols):
        row_idx = color_map_idx_row.item(i, j)
        col_idx = color_map_idx_col.item(i, j)
        rgb_color = color_map[row_idx,col_idx]
        result_data[i,j] = rgb_color


img = Image.fromarray(result_data, 'RGB')
img.save('result.png')

Answer 1

您可以用花式索引替换双 for 循环：

In [33]: result_alt = color_map[color_map_idx_row, color_map_idx_col]

这证实了结果是一样的：

In [36]: np.allclose(result_data, result_alt)
Out[36]: True

Answer 2

您可以将 3D 阵列重塑为 2D 阵列，axis=1 持有三个通道。然后，使用row-slicing，行索引从行和列索引数组计算为linear indices。请注意，重构后的数组只是一个视图，不会对任何工作区内存造成负担。因此，我们会有 -

m = color_map.shape[0]
out = color_map.reshape(-1,3)[color_map_idx_row*m + color_map_idx_col]