NMS

Девяносто пять кандидатов превращаются в шестнадцать рамок циклом на два десятка строк. Напишем этот цикл.

Задание

Реализуйте функцию nms(boxes, scores, iou_threshold), которая принимает рамки одного класса и возвращает индексы тех, что пережили NMS.

Правило то же, что разобрано в уроке: взять самого уверенного кандидата, перенести его в ответ, выбросить всех оставшихся, у кого перекрытие с ним строго больше порога, и повторять, пока список не опустеет.

• boxes — список рамок в формате [x1, y1, x2, y2], всегда x1 < x2 и y1 < y2;
• scores — список уверенностей той же длины, все значения различны;
• iou_threshold — порог перекрытия, число от 0 до 1.

Индексы в ответе — позиции рамок в исходном списке boxes, а порядок ответа — по убыванию уверенности. На пустом входе возвращается пустой список.

Пример

boxes = [[100, 100, 200, 200], [105, 105, 205, 205], [300, 300, 400, 400]]
scores = [0.9, 0.8, 0.7]

nms(boxes, scores, 0.5)   # [0, 2]

Вторая рамка перекрывается с первой на 0.82 и выбрасывается, третья стоит отдельно и остаётся.

Примечание

Слово «строго» в правиле не украшение: рамка с перекрытием ровно 0.5 при пороге 0.5 остаётся в ответе. Так устроена torchvision.ops.nms, и ожидаемые ответы тестов сняты с неё.

Сторонние библиотеки не нужны — хватит списков и арифметики. Вызывать функцию самостоятельно не требуется, проверка запустит её сама.

Что проверяется

Код проверок скрыт, так что вот полный список того, что в них разбирается:

• формат ответа: список индексов по убыванию уверенности;
• пример из условия;
• рамки не пересекаются, а уверенности идут вразнобой — порядок ответа задаётся ими, а не списком;
• пустой вход и одна рамка;
• перекрытие ровно на пороге рамку оставляет, а порог ниже — убирает;
• цепочка из трёх рамок при трёх разных порогах: подавление идёт от отобранных, а не от всех более уверенных подряд;
• семь рамок тремя парами при порогах 0.3, 0.5 и 0.7;
• порог, равный единице, не убирает никого.