Optimizing double loop in 4d numpy tensor slicing

Question

Let's say that we have two numpy ndarrays with shapes:

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

We also have an array with shape img.shape = (h,w) that is integer and tells me which "frame" v to pick for each position h,w. To do this, one can use the loop:

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

However, this is very slow. Is it possible to do it without loops? Maybe reshaping the 2D coordinates into one and then reshaping it back?

Paul Panzer · Accepted Answer · 2017-03-27 05:08:10Z

3

Here is one straightforward way

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

answered Mar 27, 2017 at 5:08

Paul Panzer

53.3k3 gold badges60 silver badges103 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Optimizing double loop in 4d numpy tensor slicing

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related