【算法】NMS非极大值抑制

IOU

IoU 的全称为交并比（Intersection over Union），IoU 计算方法有两种
1、 “预测的边框” 和 “真实的边框” 的交集和并集的比值。
2、 “预测的边框” 和 “真实的边框” 的交集和较小的预测框的比值

一般常用第一种计算 但是当一个框在另一个框内部 IOU很小时 可以次啊用第二种

求两个框的交集：
两个框box1=[x1,y1,x2,y2]，box2=[x1,y1,x2,y2] 分别是框左上角的坐标和右下角的坐标，交集部分的坐标就是以两个框左上角x、y坐标的最大值和右下角x、y坐标的最小值构成的矩形

IOU实现

import torch

# 计算某个边框和其余边框的IOU
# box=[x1,y1,x2,y2],boxes=[box1,box2...]
# ismin是否采用最小值
def iou(box,boxes,ismin=False):
    # 计算该框的面积
    box_area=(box[2]-box[0])*(box[3]-box[1])
    # 计算其他框的面积 可以用for实现
    boxex_area=(boxes[:,2]-boxes[:,0])*(boxes[:,3]-boxes[:,1])

    # 求交集矩形的坐标
    xx1=torch.maximum(box[0],boxes[:,0])
    yy1=torch.maximum(box[1],boxes[:,1])
    xx2=torch.minimum(box[2],boxes[:,2])
    yy2=torch.minimum(box[3],boxes[:,3])

    # 交集矩形的宽和高 因为可能没有交集
    w,h = torch.maximum(torch.tensor([0]),xx2-xx1),torch.maximum(torch.tensor([0]),yy2-yy1)
    over_area=w*h

    if ismin:
        return over_area/torch.minimum(box_area,boxex_area)
    return over_area/(box_area+boxex_area-over_area)

if __name__ == '__main__':
    box=torch.tensor([0,0,4,4])
    boxes=torch.tensor([
        [4,4,5,5],
        [1,1,5,5]
    ])
    print(iou(box,boxes))

NMS

（1）将所有框的得分排序，选中最高分及其对应的框
（2）遍历其余的框，如果和当前最高分框的重叠面积(IOU)大于一定阈值，我们就将框删除。
（3）从未处理的框中继续选一个得分最高的，重复上述过程。

NMS实现

# box=[score,x1,y1,x2,y2],
def NMS(boxes,threshold=0.3):
    # 按照得分降序排序
    new_boxes=boxes[boxes[:,0].argsort(descending=True)]
    # 存储需要保留下来的框 刚开始肯定得分最大的保留下来
    keep_boxes=[]
    while len(new_boxes)>1:
        _box=new_boxes[0]
        _boxex=new_boxes[1:]
        keep_boxes.append(_box)

        # 接着使用这个得分最大的框和剩下的所有框进行IOU的计算 留下那些IOU小于阈值的框
        # torch.where 返回满足条件的索引
        new_boxes=_boxex[torch.where(iou(_box,_boxex)<threshold)]
    # 出循环的时候里面还剩一个
    keep_boxes.append(new_boxes[0])
    # stack函数将 [tonsor[],tensor[]] -> tensor[[],[]]
    return torch.stack(keep_boxes)