手写NMS代码

对视频的学习NMS非极大值抑制原理及代码实现_

import torch
#坐标原点在左上角，横轴x,纵轴y,y向下画
def iou(box,boxes,isMin = False):#一个框与一堆框计算,最小面积
    #先计算box的面积[x1,y1,x2,y2]
    box_area = (box[2] -box[0]) * (box[3] -box[1])
    boxes_area =(boxes[:,2]- boxes[:,0]) * (boxes[:,3] -boxes[:,1])#把张量的某一列全取出来
    #求交集,求左上角x，y最大；右下角x,y最小
    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])  #
    #判断是否有交集,与0比较
    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.min(box_area,boxes_area)
    else:
        return over_area / (box_area +boxes_area-over_area)#并集等于两面积相加-交集

#定义nms
def nms(boxes,thresh = 0.3,isMin = False):#阈值0.3
    #根据boxes的置信度排序，假设置信度在第一列
    new_boxes =boxes[boxes[:,0].argsort(descending=True)]#设置默认参数，降序排序,用boxes【】索引
    #取出置信度最大的框
    keep_boxes = []#定义保留最大框数组
    while len(new_boxes) > 0:
        max_box = new_boxes[0]
        keep_boxes.append(max_box)
        if len(new_boxes)>1:
        #存其他框，除第一个
            other_boxes = new_boxes[1:]
            #将这些框进行iou计算
            # 小于我们设定的阈值的框才保留，说明框的不是同一个物体
            #torch.where()返回满足条件的索引
            #用other_boxes接收索引，重新赋给new，循环操作
            new_boxes = other_boxes[torch.where(iou(max_box[1:],other_boxes[:,1:],isMin)<thresh)]#切掉第一个置信度值，再进行iou运算
        else:
            break
    return torch.stack(keep_boxes)

if __name__ == '__main__':
    boxes = torch.Tensor([[0.5,1,1,10,10],[0.9,1,2,11,11],[0.4,8,8,12,12]])
    print(nms(boxes))