using 2d array as indices of a 4d array

Question

I have a Numpy 2D array (4000,8000) from a tensor.max() operation, that stores the indices of the first dimension of a 4D array (30,4000,8000,3). I need to obtain a (4000,8000,3) array that uses the indices over this set of images and extract the pixels of each position in the 2D max array.

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)

Is there any vectorised way to do that? I am fighting with the RAM usage using loops. Thanks

Answer 1

You can use np.take_along_axis.

First let's create some data (you should have provided a reproducible example):

>>> 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))

Then, we'll use np.take_along_axis. But for that the indices array must be of the same shape than the arr array. So we are using np.newaxis to insert axis where shapes don't match.

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

It already gives usable output, but with a singleton dimension on first axis:

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

So we can squeeze that:

>>> res = np.squeeze(res)

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

And eventually check if the data is as we wanted:

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

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