IIUC, вы можете сделать это, чтобы найти, где массив имеет хотя бы одно ненулевое значение:
x[np.any(x != 0, axis = 2)]
# Or, using np.where as you were:
# np.where(np.any(x!=0, axis=2))
# (array([0, 0, 0, 0, 1, 1, 1, 1]), array([0, 5, 8, 9, 1, 2, 6, 9]))
Пример:
>>> x
array([[[ 0., 0., 255.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[255., 0., 0.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[255., 0., 0.],
[255., 0., 0.]],
[[ 0., 0., 0.],
[255., 0., 0.],
[ 0., 0., 255.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[255., 0., 0.],
[ 0., 0., 0.],
[ 0., 0., 0.],
[255., 0., 0.]]])
>>> x[np.any(x != 0, axis = 2)]
array([[ 0., 0., 255.],
[255., 0., 0.],
[255., 0., 0.],
[255., 0., 0.],
[255., 0., 0.],
[ 0., 0., 255.],
[255., 0., 0.],
[255., 0., 0.]])