Non-Max Suppression

Problem

Apply non-maximum suppression to a set of bounding boxes.

Sort indices by score descending (ties broken by original index), then greedily keep a box unless its IoU with an already kept box is strictly greater than iou_threshold. Return the kept indices in selection order; boxes use the [x1, y1, x2, y2] format.

Hint

Sort by score, then compare each candidate against the boxes already kept.

Test cases

case 1

in: [[[0,0,2,2],[1,1,3,3],[10,10,12,12]],[0.9,0.8,0.7],0.5]

expected: [0,1,2]

case 2

in: [[[0,0,2,2],[1,1,3,3],[10,10,12,12]],[0.9,0.8,0.7],0.1]

expected: [0,2]

case 3

in: [[[0,0,1,1],[0,0,1,1],[0,0,1,1]],[0.5,0.9,0.7],0.5]

expected: [1]

case 4

in: [[],[],0.5]

expected: []

Starter codePython

def non_max_suppression(boxes, scores, iou_threshold):
    # Your code here
    pass

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