Pytorch版SSD目标检测算法后处理DIoUNMS实现

最近在看经典目标检测算法SSD的Pytorch的代码，顺便把后处理的非极大值抑制NMS改了一下，改成了基于DIoU的非极大值抑制DIoUNMS。
由于NMS是基于IoU进行评价的，而IoU的做法对目标框尺度和距离的影响不同，请查看具体论文DIOU。

图一 当IoU相同时，如上图所示，当相邻框的中心点越靠近当前最大得分框的中心点，则可认为其更有可能是冗余框。第一种相比于第三种更不太可能是冗余框。因此，研究者使用所提出的DIoU替代IoU作为NMS的评判准则，公式如下： 图二 DIoUNMS DIoU定义为DIoU=IoU-d²/c²，其中c和d的定义如下图所示： 图三 下面是我放在github上的pytorch版SSD代码中nms部分的代码

• 原始nms pytorch实现⬇⬇⬇

def nms(bboxes, scores, threshold=0.2, top_k=200):  #bboxes维度为[N,4],scores维度为[N,],均为tensor
    x1 = bboxes[:, 0]   #获得每一个框的左上角和右下角坐标
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    areas=(x2-x1)*(y2-y1)  #获得每个框的面积
    _,order=scores.sort(0,descending=True)  #按降序排列
    order=order[:top_k]     #取前top_k个
    keep=[]
    count=0
    while order.numel()>0:
        if order.numel()==1:
            break
        count += 1
        # print(order)
        i=order[0]
        keep.append(i)

        xx1=x1[order[1:]].clamp(min=x1[i].item())   #[N-1,]
        yy1=y1[order[1:]].clamp(min=y1[i].item())
        xx2=x2[order[1:]].clamp(max=x2[i].item())
        yy2=y2[order[1:]].clamp(max=y2[i].item())

        w=(xx2-xx1).clamp(min=0)
        h=(yy2-yy1).clamp(min=0)
        inter=w*h                        #相交的面积  [N-1,]

        overlap=inter/(areas[i]+areas[order[1:]]-inter)  #计算IOU   [N-1,]
        ids=(overlap<=threshold).nonzero().squeeze()   #返回一个包含输入 input 中非零元素索引的张量.输出张量中的每行包含 input 中非零元素的索引
        if ids.numel()==0:
            break
        order=order[ids+1]           #ids中索引为0的值在order中实际为1，后面所有的元素也一样，新的order是经过了一轮计算后留下来的bbox的索引
    return torch.tensor(keep,dtype=torch.long),count

• DIOUnms pytorch实现⬇⬇⬇

def DIOUnms(bboxes, scores, threshold=0.2, top_k=200):  #bboxes维度为[N,4],scores维度为[N,],均为tensor
    x1 = bboxes[:, 0]   #获得每一个框的左上角和右下角坐标
    y1 = bboxes[:, 1]
    x2 = bboxes[:, 2]
    y2 = bboxes[:, 3]

    center_x=x2-x1/2.0
    center_y=y2-y1/2.0

    areas=(x2-x1)*(y2-y1)  #获得每个框的面积
    _,order=scores.sort(0,descending=True)  #按降序排列
    order=order[:top_k]     #取前top_k个
    keep=[]
    count=0
    while order.numel()>0:
        if order.numel()==1:
            break
        count += 1
        # print(order)
        i=order[0]
        keep.append(i)

        xx1=x1[order[1:]].clamp(min=x1[i].item())   #[N-1,]
        yy1=y1[order[1:]].clamp(min=y1[i].item())
        xx2=x2[order[1:]].clamp(max=x2[i].item())
        yy2=y2[order[1:]].clamp(max=y2[i].item())

        w=(xx2-xx1).clamp(min=0)
        h=(yy2-yy1).clamp(min=0)
        inter=w*h                        #相交的面积  [N-1,]

        overlap=inter/(areas[i]+areas[order[1:]]-inter)  #计算IOU   [N-1,]

        # DIOU计算
        xxx1=list()
        xxx2=list()
        yyy1=list()
        yyy2=list()
        for j in range(len(np.array(xx1))):
            xxx1.append(min(x1[order[j+1]].item(), x1[i].item()))
            xxx2.append(min(x2[order[j+1]].item(), x2[i].item()))
            yyy1.append(max(y1[order[j+1]].item(), y1[i].item()))
            yyy2.append(max(y2[order[j+1]].item(), y2[i].item()))

        xxx1 = torch.Tensor(xxx1).clamp(min=0)
        xxx2 = torch.Tensor(xxx2).clamp(min=0)
        yyy1 = torch.Tensor(yyy1)
        yyy2 = torch.Tensor(yyy2)

        CDistance=torch.pow(xxx2-xxx1,2)+torch.pow(yyy2-yyy1,2)
        DDistance=torch.pow(center_x[i]-center_x[order[1:]],2)+torch.pow(center_y[i]-center_y[order[1:]],2)
        overlap=overlap-DDistance/CDistance

        ids=(overlap<=threshold).nonzero().squeeze()   #返回一个包含输入 input 中非零元素索引的张量.输出张量中的每行包含 input 中非零元素的索引
        if ids.numel()==0:
            break
        order=order[ids+1]           #ids中索引为0的值在order中实际为1，后面所有的元素也一样，新的order是经过了一轮计算后留下来的bbox的索引
    return torch.tensor(keep,dtype=torch.long),count

下面是基于YOLOV3的NMS与DIoUNMS效果图对比： NMS VS DIoUNMS