1. nms算法步骤如下:
nms算法的总体流程为:
for object in all objects:
(1) 获取当前目标类别下所有bbx的信息
(2) 将bbx按照confidence从高到低排序,并记录当前confidence最大的bbx
(3) 计算最大confidence对应的bbx与剩下所有的bbx的IOU,移除所有大于IOU阈值的bbx
(4) 对剩下的bbx,循环执行(2)和(3)直到所有的bbx均满足要求(即不能再移除bbx)
2. nms算法的输入
分别为:
- 左上点坐标(x1_1),(y1_1)
- 右下点坐标 (x2_1),(y2_1)
- 预测框的得分 scores1
- 给出的阈值
前面的三个信息由bboxes提供,
predicts_bboxes: {
"cup": [[x1_1, y1_1, x2_1, y2_1, scores1],
[x1_2, y1_2, x2_2, y2_2, scores2],
...],
"pen": [[x1_1, y1_1, x2_1, y2_1, scores1],
[x1_2, y1_2, x2_2, y2_2, scores2],
...]}
tip:
-
从流程上可以看到,其实分为三步,第三步的内容有点像求IOU的算法
-
如果在bboxes里取值如:bboxes [ : ,1],这个的意思是第一维度即行不管,取第一列的元素,即取x1_1,左上点x坐标,由此类推。
3. 具体代码
def nms(self, bboxes, scores, thresh=0.5):
#包含(1)(2)步骤,获得bboxes所有信息并按降序排序
x1 = bboxes[:,0]
y1 = bboxes[:,1]
x2 = bboxes[:,2]
y2 = bboxes[:,3]
areas = (x2-x1+1)*(y2-y1+1)
_, order = scores.sort(0, descending=True)
keep = []
#为步骤(3)中内容,取出最大的置信度回归框,numel()可以得到对应的元素个数
while order.numel() > 0:
if order.numel() == 1:
i = order.item()
keep.append(i)
break
else:
i = order[0].item()
keep.append(i)
#将最大回归框依次于各个回归框进行IOU操作,tensor.clamp()控制上下限,集不能小于min值,这里起到取最大最小(计算重叠面积)的作用
xx1 = x1[order[1:]].clamp(min=x1[i])
yy1 = y1[order[1:]].clamp(min=y1[i])
xx2 = x2[order[1:]].clamp(max=x2[i])
yy2 = y2[order[1:]].clamp(max=y2[i])
#求出对应重叠面积
inter = (xx2-xx1).clamp(min=0) * (yy2-yy1).clamp(min=0)
#计算iou值
iou = inter / (areas[i]+areas[order[1:]]-inter)
#当大于阈值时,该框的得分为0,故使用nonzero操作。在排除了IOU过阈值的值后,将IOU这一列删除,这里使用squeeze的操作
idx = (iou <= threshold).nonzero().squeeze()
if idx.numel() == 0:
break
order = order[idx+1]
return torch.LongTensor(keep)