优化双回路4D numpy的张切片

Question

比方说，我们有形状的两个numpy的ndarrays：

video.shape = (v, h, w, 3) 和 image.shape = (h, w, 3)

我们也有形状img.shape = (h,w)是整数，并告诉我一个数组， “框架”v为每个位置选择h，w。要做到这一点，我们可以使用循环：

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

然而，这是非常慢。有没有可能没有循环？也许将二维坐标重塑为一个，然后重新塑造它？

Answer 1

这里有一个简单的方法

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, :])