Заменить numpy массивов массивов значениями

Question

Я пытаюсь преобразовать массив numpy массивов. Пример ввода такой:

np.array([[[0, 0, 0], [255, 255, 255]], [[255, 255, 255], [0, 0, 0]]], np.uint8)

Я хочу заменить все массивы значениями [0, 0, 0] на 0 и [255, 255, 255] на 255. Я хочу, чтобы конечный результат был в следующем формате:

np.array([[0, 255], [255, 0]], np.uint8)

Массив всегда будет иметь вышеуказанные значения, но будет другого размера.

Это код, который я создал для этого:

array_list = []
for row in input_image:
    row_list = []
    for item in row:
        if np.array_equal(item, [0, 0, 0]):
            row_list.append(0)
        else:
            row_list.append(255)
    array_list.append(row_list)
output_image = np.array(array_list, np.uint8)

Приведенный выше код очень медленный для больших массивов, и я подумал, что может быть способ сделать это напрямую с numpy, хотя я не мог найти способ. Есть ли у вас какие-либо предложения, чтобы сделать это более эффективным?

Answer 1

X_new = np.mean(X, axis=2)
X_new
array([[  0., 255.],
       [255.,   0.]])