python - 将 numpy 数组与例如交互时如何设置迭代深度下一个？

Question

假设我有以下数组，表示 rgb 图像的结构:

[[[ 0,  1,  2], [ 3,  4,  5]]
 [[ 6,  7,  8], [ 9, 10, 11]]
 [[12, 13, 14], [15, 16, 17]]]

我如何遍历像素，例如[0, 1, 2] 然后是 [3, 4, 5]，得到相应的索引？ 使用 numpys nditer 函数我无法定义它应该停止的深度/轴，所以它会迭代每个值，例如0 然后 1 然后 2 等等。

有没有我可以定义迭代深度的 numpy 方法？

Answer 1

如果我正确理解你的问题，你可以只使用一个简单的嵌套循环

A = np.array([[[ 0,  1,  2], [ 3,  4,  5]],
              [[ 6,  7,  8], [ 9, 10, 11]],
              [[12, 13, 14], [15, 16, 17]]])

for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        print(i, j, A[i,j,...])

0 0 [0 1 2]
0 1 [3 4 5]
1 0 [6 7 8]
1 1 [ 9 10 11]
2 0 [12 13 14]
2 1 [15 16 17]

或者，以更 numpythonic 的方式使用 np.ndindex

for i in np.ndindex(A.shape[:2]):
    print(i, A[i])

(0, 0) [0 1 2]
(0, 1) [3 4 5]
(1, 0) [6 7 8]
(1, 1) [ 9 10 11]
(2, 0) [12 13 14]
(2, 1) [15 16 17]