Numpy 访问 nD 数组值 vetorized 方式

Question

在 python 3.9 中使用 numpy。

我有一组图片（480 张图片，每张图片的大小为 640x480）。

我想创建一个新图像，其中该图像的每一行 N 是我的数组中图像 #N 的行 N 的副本。

我可以用循环来做到这一点：

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

有没有办法使用向量化的方式使用 numpy 来完成它，因为根据文档它更快？

对于上下文： 我的图像数组最初是用这个创建的：

image_array = np.full((480,480,640,3), blank_image, np.uint8)

然后从视频流中更新数组中的每个图像。

我迷失了 np.diagonal、np.compress 的逻辑......

谢谢

编辑：

感谢第一个回答：

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

专心工作。

Answer 1

这样的事情怎么样（例如两个 2x2 像素的图像）：
要根据您的需要调整它，您必须更改“arange”函数的参数。

import numpy as np

imgs = np.arange(2*2*2*3).reshape(2, 2, 2, 3)

print(f'images = \n{imgs}')

new_imgs = np.zeros((2, 2, 2, 3))

new_imgs[np.arange(2), np.arange(2), ...] = imgs[np.arange(2), np.arange(2), ...]

print(f'new_imgs = \n{new_imgs}')

输出：

images = 
[[[[ 0  1  2]
   [ 3  4  5]]

  [[ 6  7  8]
   [ 9 10 11]]]


 [[[12 13 14]
   [15 16 17]]

  [[18 19 20]
   [21 22 23]]]]
new_imgs = 
[[[[ 0.  1.  2.]
   [ 3.  4.  5.]]

  [[ 0.  0.  0.]
   [ 0.  0.  0.]]]


 [[[ 0.  0.  0.]
   [ 0.  0.  0.]]

  [[18. 19. 20.]
   [21. 22. 23.]]]]

Answer 2

只需在前两个维度中使用相同的索引数组即可。

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