Быстрый способ извлечения ROI в массиве 3D NumPy

Question

У меня есть одномерный массив, который содержит некоторые видеоданные:

data=np.random.randint(0,high=255,size=(500*500*100),dtype=np.uint8)
imgs=np.reshape(data,(100, 500,500)) # 100 frames, 500x500 pixels

Я хочу выделить определенные области интереса (roi) вдоль всех кадров

idx=np.random.randint(0,high=500*500,size=(49, 300)) #300 rois, 49 points each
rois=imgs.reshape(100, -1)[:,idx]

Я сплющил каждыйкадр, а затем взял rois вдоль 1-го измерения.Реальный массив imgs больше показанного здесь, и предыдущая операция индексирования может быть немного медленной.Если я изменю imgs другим способом (см. Ниже), rois.size будет таким же, и индексирование будет намного быстрее, но это приведет к неверным данным

%timeit imgs.reshape(100, -1)[:,idx] # 13 ms
%timeit imgs.reshape(-1, 100)[idx, :] # 1.2 ms, much faster but wrong data

В моем реальном коде разница будетпочти в 50 раз.Есть ли способ индексировать imgs быстро?

Answer 1

Кажется, можно сэкономить хотя бы немного времени, отсортировав пиксели ROI и работая с транспонированными координатами:

>>> def f_pp(im2D, idx):
...     s = np.argsort(idx.ravel())
...     out = np.empty((*idx.shape, im2D.shape[0]), im2D.dtype)
...     out.reshape(-1, im2D.shape[0])[s] = im2D.T[idx.ravel()[s]]
...     return out
... 

# results are the same:
>>> np.all(f_pp(imgs.reshape(100, -1), idx) == np.moveaxis(imgs.reshape(100, -1)[:, idx], 0, 2))
True

>>> timeit("imgs.reshape(100, -1)[:, idx]", globals=globals(), number=100)
1.3392871069954708
# transposing alone is not enough:
>>> timeit("imgs.reshape(100, -1).T[idx]", globals=globals(), number=100)
1.3336799899989273
# but together with sorting I see a 2x speedup
>>> timeit("f_pp(imgs.reshape(100, -1), idx)", globals=globals(), number=100)
0.5874412529956317
# still much worse than if we had a more favorable memory layout in
# the first place
>>> timeit("imgs.reshape(-1, 100)[idx]", globals=globals(), number=100)
0.06296327701420523