向量化批处理图像像素查找numpy数组

Question

Suppose I have an input numpy array of the form (B, H, W, C) and have a grid comprised of x values of the form (B, W) and y values of the form (B, H) .

My goal is to use the x and y values to grab the C values of the image array. I'm able to do it when B = 1 that is when my image is of the form (H, W, C) and when x and y are of the form (H\\W, ) but have no clue how to expand this to a batch of images contained in a single numpy array.

Is there some sort of reshaping I can take advantage of?

Example

Suppose I have 2 cat images of size (400, 400, 3) . Then input_img.shape = (2, 400, 400, 3) .

I have a list x.shape = (2, 400) and y.shape = (2, 400) and would like to index into input_img such that I obtain for each (x_i, y_i) an array z_i.shape = (C, ) totalling z.shape = (400, C) .

What I'm trying to Vectorize

for i in range(batch_size):
    z_i = input_img[i, x[i], y[i]]

Answer 1

这是一种使用advanced-indexing的方法-

input_img[np.arange(batch_size)[:,None], x, y]