【深度学习笔记】NMS代码学习

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)