Argmax Class Map

Problem

Given an H x W x C map of class scores, return an H x W map of the argmax class index per pixel.

Each pixel is a list of C scores; ties keep the smallest index.

Hint

Compare with strict greater-than so equal scores keep the earlier class.

case 1

in: [[[[0.1,0.9,0.5],[0.7,0.2,0.1]],[[0.3,0.3,0.3],[0.6,0.4,0.5]]]]

expected: [[1,0],[0,0]]

case 2

in: [[[[1,2]]]]

expected: [[1]]

case 3

in: [[[[0.2,0.2,0.1]]]]

expected: [[0]]

Starter codePython

def argmax_class(scores):
    # Your code here
    pass

3 test cases · run your code against them in the editor