Python 四维numpy张量切片中的双环优化

假设我们有两个带形状的numpy Ndarray：

video.shape=（v，h，w，3）

和

image.shape=（高、宽、3）

我们还有一个shape

img.shape=（h，w）

为整数的数组，它告诉我为每个位置h，w选择哪个“帧”v。为此，可以使用循环：

for j in range(w):
    for i in range(h):
        image[i, j, :] = video[img[i, j], i, j, :]

---

然而，这是非常缓慢的。有没有可能不使用循环？也许将二维坐标重新整形为一个，然后再重新整形？

这里有一个简单的方法

import numpy as np

v, h, w = 40, 50, 60

video = np.random.random((v,h,w,3))
img = np.random.randint(0, v, (h, w))

i, j = img.shape
i, j = np.ogrid[:i, :j]

image = video[img, i, j, :]

# check

for j in range(w):
    for i in range(h):
        assert np.all(image[i, j, :] == video[img[i, j], i, j, :])

---

这里有一个简单的方法

import numpy as np

v, h, w = 40, 50, 60

video = np.random.random((v,h,w,3))
img = np.random.randint(0, v, (h, w))

i, j = img.shape
i, j = np.ogrid[:i, :j]

image = video[img, i, j, :]

# check

for j in range(w):
    for i in range(h):
        assert np.all(image[i, j, :] == video[img[i, j], i, j, :])