NMS的原理解析

YOLOv5中NMS原理:

1、首先去掉置信度小于设定阈值的框；

2、根据执行度从大到小排列

3、计算第一个框(置信度最高的框)和后面所有的框的IOU，得到的值小于设定IOU的阈值留下，大于IOU 的值说明相交太多，舍去。

4、经过一轮后，计算第二个框和后面所有框的IOU，依此类推，最后一直遍历，得到所有的框。

import torch

def iou(box,boxes,isMin=False):
  box_area = (box[2]-box[0])*(box[3]-box[1])
  boxes_areas = (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)

  ovr_area = w*h
  if isMin:
    return ovr_area/torch.min(box_area,boxes_areas)
  else:
    return ovr_area/(box_area+boxes_areas-ovr_area)
    
 
def nms(boxes,thresh=0.3,isMin=False):
  # 对所有框根据置信度排序
  new_boxes = boxes[boxes[:,0].argsort(descending=True)]
  keep_boxes = []
  while len(new_boxes)>0:
    _box = new_boxes[0]
    keep_boxes.append(_box)
    if len(new_boxes)>1:
      _boxes = new_boxes[1:]
      new_boxes = _boxes[torch.where(iou(_box,_boxes,isMin)<thresh)]
    else:
      break
  return torch.stack(keep_boxes)



if __name__=="__main__":
    boxes = torch.tensor([[0.5,1,1,10,10],
                          [0.9,1,1,11,11],
                          [0.4,8,8,12,12]])
    
    print(nms(boxes,0.3))