如何使用矢量化快速解码 one-hot 编码的 NumPy 矩阵？

Question

Given an image matrix of shape (height, width) with values in the uint8 range, which was one-hot encoded (converted to categorical) to a shape of (height, width, n) where n is the number of possible categories, 3 in this instance resulting in a shape of (height, width, 3) , I would like to undo the categorical conversion and get the original shape of (height, width) . The following solution works, but could be made much faster:

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

I would like a solution, using NumPy vectorization , without the need for a slower Python for loop .

Thank you for any suggestions.

Answer 1

Looks like you want to do a reduction over the last dimension of your array, in particular a numpy.argmax . Fortunately, this numpy function accepts an axis keyword, so you should be able to do the same in just one call:

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