如何使用Numpy向量化方式从n维图像数组提取对应行生成新图像

需求背景 #

现有维度为(480, 480, 640, 3)的Numpy数组image_array，其中包含480张480行、640列的3通道彩色图像。需要生成一张新图像，要求新图像的第N行，完全对应原数组中第N张图像的第N行。

原有实现（循环版） #

最初通过Python层循环实现需求，代码如下：

for i, img in enumerate(image_array):
    new_image[i,:] = img[i,:]

该方案逻辑直观，但Python循环效率较低，处理高分辨率图像时性能表现较差。

向量化优化方案 #

利用Numpy 高级整数数组索引 特性，可以完全替换循环实现，代码如下：

idx = np.arange(len(image_array))
new_image = image_array[idx, idx, ...]

实现原理 #

• np.arange(len(image_array))会生成一个长度为480、取值为0~479的一维整数数组idx
• 同时将idx作为数组前两个维度的索引时，Numpy会自动按位置匹配索引值，依次取出image_array[0,0,...]、image_array[1,1,...]……image_array[479,479,...]的结果
• 末尾的...是Numpy的省略索引语法，等价于后续维度全部取全量值，这里对应保留每张图像对应行的640列像素、3个颜色通道信息
• 最终得到的new_image形状为(480, 640, 3)，完全符合需求，且所有运算都在Numpy底层实现，执行效率远高于Python层循环

内容的提问来源于stack exchange，提问作者T H