python - numpy.equal 与嵌套列表

Question

我想在图片中搜索一个矩形。图片来自太平船务。这意味着我将得到一个二维数组，其中每个项目都是一个包含三个颜色条目的列表。

我使用 np.equal 来获取具有搜索颜色的矩形在哪里。这是一个缩小的例子:

>>> l = np.array([[1,1], [2,1], [2,2], [1,0]])
>>> np.equal(l, [2,1])  # where [2,1] is the searched color
array([[False,  True],
   [ True,  True],
   [ True, False],
   [False, False]], dtype=bool)

但我已经预料到:

array([False, True, False, False], dtype=bool)

或

array([[False,  False],
   [ True,  True],
   [ False, False],
   [False, False]], dtype=bool)

如何使用 numpy 实现嵌套列表比较？

注意:然后我想用 np.where 从 np.equal 的结果中提取矩形的索引。

Answer 1

您可以沿第二个轴使用 all 方法:

>>> result = numpy.array([[1, 1], [2, 1], [2, 2], [1, 0]]) == [2, 1]
>>> result.all(axis=1)
array([False,  True, False, False], dtype=bool)

并获取索引:

>>> result.all(axis=1).nonzero()
(array([1]),)

我更喜欢 nonzero 而不是 where，因为 where 做 two very different things取决于传递给它的参数数量。当我需要它独特的功能时，我使用 where；当我需要 nonzero 的行为时，我明确地使用 nonzero。