использование 2d-массива в качестве индексов 4d-массива

Question

У меня есть Numpy 2D-массив (4000,8000) из операции tensor.max(), в котором хранятся индексы первого измерения 4D-массива (30,4000,8000,3). Мне нужно получить массив (4000,8000,3), который использует индексы по этому набору изображений и извлекает пиксели каждой позиции в массиве 2D max.

A = np.random.randint( 0, 29, (4000,8000), dtype=int)
B = np.random.randint(0,255,(30,4000,8000,3),dtype=np.uint8)

final = np.zeros((B.shape[1],B.shape[2],3))

r = 0
c = 0
for row in A:
  c = 0
  for col in row:
    x = A[r,c]
    final[r,c] = B[x,r,c]
    c=c+1
  r=r+1

print(final.shape)

Есть ли векторизованный способ сделай это? Я борюсь с использованием ОЗУ с помощью циклов. Спасибо

Answer 1

Вы можете использовать np.take_along_axis.

Сначала давайте создадим некоторые данные (вы должны были предоставить воспроизводимый пример ):

>>> N, H, W, C = 10, 20, 30, 3
>>> arr = np.random.randn(N, H, W, C)
>>> indices = np.random.randint(0, N, size=(H, W))

Затем мы будем использовать np.take_along_axis. Но для этого массив indices должен иметь ту же форму, что и массив arr. Итак, мы используем np.newaxis для вставки оси там, где формы не совпадают.

>>> res = np.take_along_axis(arr, indices[np.newaxis, ..., np.newaxis], axis=0)

Он уже дает полезный вывод, но с одноэлементным размером на первой оси:

>>> res.shape
(1, 20, 30, 3)

Итак, мы можем сжать это:

>>> res = np.squeeze(res)

>>> res.shape
(20, 30, 3)

И в конечном итоге проверить, соответствуют ли данные нам:

>>> np.all(res[0, 0] == arr[indices[0, 0], 0, 0])
True

>>> np.all(res[5, 3] == arr[indices[5, 3], 5, 3])
True