通过 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]