numpy фильтр по размерам для заданных значений c

Question

Я пытаюсь отфильтровать изображение с формой (224, 224, 3) в RGB, которое было загружено в массив numpy, где последняя ось является каналами.

Простая логика c проста : найдите значения c пикселей по 3 каналам и назначьте метку для нового массива 1-Dim / 1-Channel. Это используется для сетей сегментации, где пиксельная метка назначается для представления области объекта на изображении.

Текущий l oop выглядит следующим образом:

mask     = Image.open(args.input).convert('RGB')
array    = np.array(img_mask)

h, w, c = array.shape
bg = np.zeros((w, h))
fg = np.zeros((w, h))

for i in range(0, h):
    for k in range(0, w):
        if np.array_equal(array[i,k], [0, 0, 0]):
            bg[i,k] = 1
        elif np.array_equal(array[i,k], [255, 200, 0]):
            fg[i,k] = 2

label = np.stack((bg, fg), axis=0)

Это создает (2, 224, 224), где axis = 0 - это каналы, а затем axis=1 имеет метки пикселей.

Я попытался заменить два цикла for, используя np.where, однако результат все еще трехмерный массив, тогда как мне нужен одномерный массив, созданный для каждой пиксельной координаты, которая соответствует условию:

bg    = np.where(array == [0, 0, 0],     1, 0)
fg    = np.where(array == [255, 200, 0], 2, 0)
label = np.stack((bg, fg), axis=0)

Однако здесь я получаю: (2, 224, 224, 3)

Текущий for l oop реализация мучительно медленная, и я подозреваю, что это можно сделать с помощью встроенных numpy функций, но мне пока не удалось достичь того же результата.

Answer 1

Я уверен, что вы можете использовать select для маскировки и ravel для получения 1-D массива:

array = np.array([[[0,0,0],[128,128,128]],
                  [[128,128,128],[255,255,0]]
                 ])

label = np.select([(array==(0,0,0)).all(-1),      # all make sure all channels are identical
                   (array==(255,255,0)).all(-1)],
                  [1,2], 0)

print(label.ravel())

Вывод:

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