Как быстро декодировать матрицу NumPy с горячим кодированием с помощью векторизации?

Question

Дана матрица изображения формы (height, width) со значениями в диапазоне uint8, которая была закодирована в горячем режиме (преобразована в категориальную) в форму (height, width, n), где n - количество возможных категорий, 3 в этот экземпляр приводит к форме (height, width, 3), я хотел бы отменить категориальное преобразование и получить исходную форму (height, width). Следующее решение работает, но его можно было бы сделать намного быстрее:

def decode(image):
    image = image

    height = image.shape[0]
    width = image.shape[1]

    decoded_image = numpy.ndarray(shape=(height, width), dtype=numpy.uint8)

    for i in range(0, height):
        for j in range(0, width):
            decoded_image[i][j] = numpy.argmax(image[i][j])

    return decoded_image

Мне нужно решение, использующее NumPy vectorization, без необходимости в более медленном Python for loop.

Спасибо за любые предложения.

Answer 1

Похоже, вы хотите уменьшить последнее измерение вашего массива, в частности numpy.argmax. К счастью, эта функция numpy принимает ключевое слово axis, так что вы сможете сделать то же самое всего за один вызов:

decoded_image = numpy.argmax(image, axis=2)