Advertisement

用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)

还没有任何评论哟~