【NMS】non max suppression非最大抑制及其实现

黑色的是被检测的人，蓝色是检测框，蓝色数字是每个框的概率分数。
当我们检测一个物体的时候。
会有多个检测框，此时，我们需要用NMS 来删除多余的框。
我们首先选出得分最高的框–在这里就是得分0.8的框。
此时我们需要判断得分最高的框和另一个框的iou值，
假设我们设定iou的标准阈值是0.5，
如果2个框的iou算出来大于我们设定iou的标准阈值是0.5，那么我们就删掉另一个框，保留得分最高的框。
如果小于，就不删。

注意，这个操作必须是针对同一个物体的。如果是这样的情况：

下面我们来实现NMS：

import torch
from iou import intersection_over_union
def non_max_suppression(
    bboxes,
    iou_threshold,
    prob_threshold,
    box_format="corners"
):
    # bboxes的结构假设是这样的【类别，框的概率，x1,y1,x2,y2】
    assert type(bboxes) == list
    bboxes = [box for box in bboxes if box[1] > prob_threshold]
    bboxes = sorted(bboxes,key= lambda x:x[1],reverse=True)
    bboxes_after_nms = []

    while bboxes:
        chosen_box = bboxes.pop(0)

        bboxes = [
            box for box in bboxes
            if box[0] != chosen_box[0]
            or intersection_over_union(
                torch.tensor(chosen_box[2:]),
                torch.tensor(box[2:])
            )
            < iou_threshold
        ]
        bboxes_after_nms.append(chosen_box)
    return bboxes_after_nms

iou就是我目标检测专栏的iou