在任意维度上用 numpy 数组索引 numpy 数组

Question

假设 img 是一个二维 numpy 数组。还假设 x 和 y 是与 img 形状相同的整数值二维 numpy 数组。考虑：

newImg = img[x, y]

newImg 现在是一个与 img 形状相同的二维数组，其中 newImg[i,j] == img[ x[i,j], y[i,j] ] 表示所有 i 和 j。

我想将此过程推广到任意数量的维度。也就是说，设 img 是一个 d 维 numpy 数组，并取 x[i]，对于 range(0, d) 中的 i，是一个与 img 形状相同的整数值 d 维 numpy 数组。我基本上想要的是：

newImg = img[x[0], x[1], ..., x[d-1]]

这显然是伪代码，预计不会工作。

如何使用 NumPy 做到这一点？

Answer 1

你试过简单吗

newImg = img[x]

看起来应该可以了！

我假设x 是满足以下条件的整数数组的列表或元组

len(x) == img.ndim
all(a.shape == img.shape for a in x)

这似乎与您所描述的相符。