RPN NMS Top-K

Problem

Run greedy non-maximum suppression on region proposals and keep at most k.

Process boxes in descending score order (ties by index), suppress a box when its IoU with a kept box is strictly greater than iou_threshold, and stop once k boxes are kept. Return the kept indices.

Test cases

case 1

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

expected: [0,2]

case 2

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

expected: [0]

case 3

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

expected: [0,1,2]

case 4

in: [[],[],3,0.5]

expected: []

Starter codePython

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

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