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
文章出处登录后可见!