如何将3D数组转换为argmax的2D数组？

Question

I have an Numpy array of size (W, H, C), where 'C' is a number of classes for a semantic segmentation task. What I need is a Numpy array of size (H, W), where each element is the index of the class that is appropriate for that pixel.

I've found a way to do it that runs VERY slowly.

masks = {list of 2d binary masks}
output_mask = np.zeros(width * height)
output_mask = output_mask.reshape(width, height)

for i in range(width):
    for j in range(height):
        class_id = 0
        for mask in masks:
            class_id += 1
            if mask[i, j] == 1:
                output_mask[i, j] = class_id

I was hoping there might be a better way. Can anyone help me?

Answer 1

import numpy as np

arr = np.random.rand(10, 10, 3)
max_val = np.argmax(arr, axis=-1)

print(max_val.shape)