文章目录[隐藏]

1.NMS(非极大值抑制)介绍

在目标检测算法中，对于一个真实物体的往往会有多个候选框输出，而我们只需一个最好的检测框，多余的候选框会影响检测的精度，因此需要利用NMS(非极大值抑制)来过滤掉重叠的候选框，得到最佳的候选框。所谓非极大值抑制就是说，对于多个重叠的检测框，它不是最好的框，那么就将抑制（删除），只保留一个最好的。
在这里插入图片描述

2.算法所需的指标

检测框的得分：NMS假设一个边框的预测得分越高，这个框就要被优先考虑，其他框与该框重叠超过一定程度的边框就要被舍弃掉。
IoU：在评价两个边框重合程度时，NMS需要使用这个指标，如果两个边框的IoU超过一定的阈值时，得分低的边框就会被舍弃。

3.算法流程

算法的输入包含了所有预测框的得分、左上点坐标、右下点坐标5个预测量。已经一个设定的IoU阈值。
流程如下：

按照预测框的得分，对所有预测框进行降序排列，记录下排列的索引order,并新建一个列表keep,作为最终的筛选后的边框索引结果。
将排序后的第一个边框最为当前框，并将其保留到keep中，再求当前框与剩余所有框的IoU.
在order中，仅仅保留IoU小于设定阈值的索引，重复第2步，直到order中仅仅剩余一个边框，则将其保留到keep中，退出循环，NMS结束。

4.代码实现

# -*-coding:utf-8-*-

def nms(self,bboxes,scores,threshold=0.5):
    x1=bboxes[:,0]
    y1=bboxes[:,1]
    x2=bboxes[:,2]
    y2=bboxes[:,3]
    #计算每个box的面积
    areas=(x2-x1+1)*(y2-y1+1)
    
    #对每个框进行降序排序，order为降序排列的索引
    _,order=scores.sort(0,descending=True)
    
    #keep保留了NMS留下的边框box
    keep=[]
    while order.numel()>0:
        if order.numel()==1:
            i=order.item()
            keep.append(i)
            break
        else:
            i=order[0].item()#保留scores最大的那个表框的box[i]
            keep.append(i)

        # 巧妙利用tensor.clamp函数求取每一个框与当前框的最大值和最小值
        xx1=x1[order[1:]].clamp(min=x1[i])
        yy1=y1[order[1:]].clamp(min=y1[i])
        xx2=x2[order[1:]].clamp(max=x2[i])
        yy2=y2[order[1:]].clamp(max=y2[i])
        # 求取每一个框与当前框的重合部分面积
        inter = (xx2 - xx1).clamp(min=0) * (yy2 - yy1).clamp(min=0)
        # 计算每一个框与当前框的IoU
        iou = inter / (areas[i] + areas[order[1:]] - inter)
        # 保留IoU小于阈值的边框索引
        idx = (iou <= threshold).nonzero().squeeze()
        if idx.numel() == 0:
            break
        # 这里的+1是为了补充idx与order之间的索引差
        order = order[idx + 1]
        # 返回保留下的所有边框的索引值，类型为torch.LongTensor
        return torch.LongTensor(keep)