How to vectorize extraction in multidimensional numpy array

Question

I am struggling with a vectorization problem which seems to be simple at first glance:

Let's say I have 100 images of size (7, 7) with 2 channels represented by a numpy array of size (100, 2, 7, 7). I would like to extract on all images small patches (let say of size (2, 3, 3)), but these patches are not located at the same place on every image. Locations of the patch are described by a matrix of size (2, 100) (one x and one y for each image). I am able to do that with a for loop over all images, but it takes time.

Here is an example code :

data = np.arange(9800).reshape((100, 2, 7, 7))
size = 3
pos = np.random.randint(0, 7-size, (2, 100))
for i in range(a.shape[0]):
    patch = data[i, :, pos[0,i]:(pos[0,i]+size), pos[1,i]:pos[1,i]+size]

In other words, I'd like to reproduce this code without the for loop. Does anyone have a clue to do that?

Only to precise, in real life my dataset is bigger than 100 images... that is why I would like to get rid of the for loop — Victor Boutin, Feb 05 '18 at 11:08
Yes I am storing the output. The small code here was only to illustrate my problem — Victor Boutin, Feb 05 '18 at 12:42

score 2 · Accepted Answer · edited Feb 07 '18 at 23:16

2

The standard trick with numpy.lib.stride_tricks.as_strided should work here:

das = np.lib.stride_tricks.as_strided(data, (5, 5, 100, 2, 3, 3), data.strides[-2:] + data.strides)
patches = das[(*pos, np.arange(100))]

You can verify it works using:

for i in range(data.shape[0]):
    assert np.all(patches[i]==data[i,:,pos[0,i]:pos[0,i]+size,pos[1,i]:pos[1,i]+size])

edited Feb 07 '18 at 23:16

meduz

3,903
1
28
40

answered Feb 05 '18 at 11:29

Paul Panzer

51,835
3
54
99

Thanks, it works. Even If thought I would have been simpler !! Thanks ! – Victor Boutin Feb 05 '18 at 12:47

How to vectorize extraction in multidimensional numpy array

1 Answers1