優化雙迴路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, :])