用PyTorch实现CIoU、DIoU和GIoU的NMS算法
发布时间
阅读量:
阅读量
在目标检测的后处理阶段我们需要用到NMS算法。
在NMS算法中存在一个步骤是需要计算当前score最高的框与其他框的IoU大小的
针对这一步,我们可以进行优化,改变IoU的计算方式。
现有的经典IoU计算方法主要包括GIoU、DIoU以及CIoU。对于传统的NMS算法进行优化。
CIoU NMS,DIoU NMS和GIoU NMS。
先看下传统的NMS算法的代码:
def ori_nms(self, scores, boxes, iou_thres):
if boxes.shape[0] == 0:
return torch.zeros(0,device=boxes.device).long()
x1,y1,x2,y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:,3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = torch.sort(scores, descending=True)[1] #(?,)
keep =[]
while order.numel() > 0:
if order.numel() == 1:
全部评论 (0)
还没有任何评论哟~
