【深度學習筆記】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)