改进YOLOv5、YOLOv7系列:修改NMS为:DIoU-NMS,SIoU-NMS,EIoU-NMS,CIoU-NMS,GIoU-NMS

1.在 general.py 文件中加入NMS方法

def NMS(boxes, scores, iou_thres, class_nms='CIoU'):
    # class_nms=class_nms
    GIoU=CIoU=DIoU=EIoU=SIoU=False
    if class_nms == 'CIoU':
        CIoU=True
    elif class_nms == 'DIoU':
        DIoU=True
    elif class_nms == 'GIoU':
        GIoU=True
    elif class_nms == 'EIoU':
        EIoU=True
    else :
        SIoU=True
    B = torch.argsort(scores, dim=-1, descending=True)
    keep = []
    while B.numel() > 0:
        index = B[0]
        keep.append(index)
        if B.numel() == 1: break
        iou = bbox_iou(boxes[index, :], boxes[B[1:], :], GIoU=GIoU, DIoU=DIoU, CIoU=CIoU, EIoU=EIoU, SIoU=SIoU)
        inds = torch.nonzero(iou <= iou_thres).reshape(-1)
        B = B[inds + 1]
    return torch.tensor(keep)
 

2.使用:i = NMS(boxes, scores, iou_thres, class_nms='xxx')
替换non_max_suppression 方法中的:i = torchvision.ops.nms(boxes, scores, iou_thres)

3.改NMS更改:

在i = NMS(boxes, scores, iou_thres, class_nms='DIoU') 中将class_nms 设置为= DloU时,则开启DloU-NMS

在i = NMS(boxes, scores, iou_thres, class_nms='DIoU') 中将class_nms 设置为= SloU时,则开启SloU-NMS

在i = NMS(boxes, scores, iou_thres, class_nms='DIoU') 中将class_nms 设置为= EloU时,则开启EloU-NMS

在i = NMS(boxes, scores, iou_thres, class_nms='DIoU') 中将class_nms 设置为= GloU时,则开启GloU-NMS

你可能感兴趣的:(YOLOv5,YOLOv7系列,深度学习,python,pytorch)