Есть ли способ запустить np.where для нескольких значений, а не только для одного?

Question

Мне интересно, если у меня есть изображение в массиве, скажем, 250x250x3 (3 канала), возможно ли использовать np.where, чтобы быстро выяснить, равен ли какой-либо из массивов 250x250 размера 3 [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 (или что-то особенно сложное) вообще. Вы можете просто использовать мощь логических массивов:

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

Answer 2

Для простоты предположим, что мы ищем все места, где все 3 канала равны 1. Это сделает работу:

np.random.seed(0)
a=np.random.randint(0,2,(3,5,5))
print(a)
np.where((a[0]==1)*(a[1]==1)*(a[2]==1))

Это выдаст

[[[0 1 1 0 1]
  [1 1 1 1 1]
  [1 0 0 1 0]
  [0 0 0 0 1]
  [0 1 1 0 0]]

 [[1 1 1 1 0]
  [1 0 1 0 1]
  [1 0 1 1 0]
  [0 1 0 1 1]
  [1 1 1 0 1]]

 [[0 1 1 1 1]
  [0 1 0 0 1]
  [1 0 1 0 1]
  [0 0 0 0 0]
  [1 1 0 0 0]]]

(array([0, 0, 1, 2, 4], dtype=int64), array([1, 2, 4, 0, 1], dtype=int64))

И действительноесть 5 координат, в которых все 3 канала равны 1. Если вы хотите получить более легкое для чтения представление, замените последнюю строку на

tuple(zip(*np.where((a[0]==1)*(a[1]==1)*(a[2]==1))))

. Это выведет

((0, 1), (0, 2), (1, 4), (2, 0), (4, 1))

это все 5 мест, где все 3 канала равны 1.

Обратите внимание, что (a[0]==1)*(a[1]==1)*(a[2]==1) - это просто

array([[False,  True,  True, False, False],
       [False, False, False, False,  True],
       [ True, False, False, False, False],
       [False, False, False, False, False],
       [False,  True, False, False, False]])

логическое представление, которое вы искали.

Если вы хотите получить любой другой триплет, скажем, [143, 255, 0], просто используйте (a[0]==143)*(a[1]==255)*(a[2]==0).