Question

我有以下代码可以像预期的那样工作，但我很好奇是否可以用本机numpy函数/方法替换循环以获得更好的性能。我所拥有的是一个包含RGB值的数组，我将其用作查找表，两个2d数组保存灰度值（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]

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

2 个答案: