python - 有什么方法可以在多个值而不是一个值上运行 np.where 吗？

Question

我想知道我是否有一个 numpy 数组中的图像，比如 250x250x3(3 个 channel )，是否可以使用 np.where 快速找出大小为 3 的 250x250 数组是否等于 [143 , 255, 0] 或其他由 rgb 表示的颜色并得到一个 250x250 的 bool 数组？

当我在 4x4x3 的代码中尝试它时，我得到了一个 3x3 的数组，但我不完全确定这个形状是从哪里来的。

import numpy as np

test = np.arange(4,52).reshape(4,4,3)
print(np.where(test == [4,5,6]))

-------------------------------------------

Result:

array([[0, 0, 0],
       [0, 0, 0],
       [0, 1, 2]])


What I'm trying to get:

array([[1, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])

Answer 1

解决方案

您根本不需要np.where(或任何特别复杂的东西)。您可以利用 bool 数组的强大功能:

print(np.all(test == [4,5,6], axis=-1).astype(int))
# output:
#     [[1 0 0 0]
#      [0 0 0 0]
#      [0 0 0 0]
#      [0 0 0 0]]

等效的替代方法是使用logical_and:

print(np.logical_and.reduce(test == [4,5,6], axis=-1).astype(int))
# output:
#     [[1 0 0 0]
#      [0 0 0 0]
#      [0 0 0 0]
#      [0 0 0 0]]

重载测试

import numpy as np
np.random.seed(0)

# the subarray we'll search for
pattern = [143, 255, 0]

# generate a random test array
arr = np.random.randint(0, 255, size=(255,255,3))

# insert the pattern array at ~10000 random indices
ix = np.unique(np.random.randint(np.prod(arr.shape[:-1]), size=10000))
arr.reshape(-1, arr.shape[-1])[ix] = pattern

# find all instances of the pattern array (ignore partial matches)
loc = np.all(arr==pattern, axis=-1).astype(int)

# test that the found locs are equivalent to the test ixs
locix = np.ravel_multi_index(loc.nonzero(), arr.shape[:-1])
np.testing.assert_array_equal(np.sort(ix), np.sort(locix))
# test has been run, the above assert passes