矢量化批量图像像素查找numpy数组

Question

假设我有一个(B, H, W, C)形式的输入numpy数组，并且有一个由(B, W)形式的x值和(B, H)形式的y值组成的网格。你知道吗

我的目标是使用x和y值来获取图像数组的C值。当B = 1也就是说当我的图像是(H, W, C)的形式，当x和y是(H\W, )的形式时，我能够做到这一点，但是我不知道如何将它扩展到包含在单个numpy数组中的一批图像。你知道吗

有什么我可以利用的整形方法吗？你知道吗

示例

假设我有两个大小为(400, 400, 3)的cat图像。然后input_img.shape = (2, 400, 400, 3)。你知道吗

我有一个x.shape = (2, 400)和y.shape = (2, 400)的列表，我想索引到input_img，这样我就可以为每个(x_i, y_i)获得一个数组z_i.shape = (C, )，总计z.shape = (400, C)。你知道吗

我要向量化的内容

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

Answer 1

下面是一个使用^{}-

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