Optimierung der Doppelschleife in 4d numpy Tensor

Question

Aufschneiden

Lassen Sie uns sagen, dass wir zwei numpy ndarrays mit Formen haben:

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

wir eine Reihe haben auch mit Form img.shape = (h,w), die integer ist und sagt mir, was "frame" v für jede Position h, w auswählen. Dazu kann man die Schleife verwenden:

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

Dies ist jedoch sehr langsam. Ist es möglich, ohne Schleifen zu machen? Vielleicht die 2D-Koordinaten in eine umformen und sie dann neu gestalten?

Answer 1

Hier ist eine einfache Möglichkeit,

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