Pytorch版SSD目标检测算法后处理DIoUNMS实现
2022-06-09 本文已影响0人
教训小磊
最近在看经典目标检测算法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
NMS VS DIoUNMS
