改进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

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

清风⁡⁢ㅤ⁡⁢ㅤ⁡⁢ㅤ⁡⁢ㅤ

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值