一、技術背景
在目標檢測任務中,模型通常會對同一目標生成多個重疊的候選框(如錨框或預測框)。非極大值抑制(Non-Maximum Suppression, NMS) 是一種關鍵的后處理技術,用于去除冗余的檢測結果,保留置信度最高且位置最優的邊界框。本文將通過一段Python代碼解析NMS的核心實現邏輯,并演示其在OpenCV環境中的實際效果。
二、算法核心思想
NMS的核心是通過以下步驟篩選邊界框:
- 按置信度排序:優先處理置信度最高的預測框。
- 計算交并比(IoU):與當前框重疊度高的候選框將被抑制。
- 迭代篩選:重復上述過程直至處理完所有候選框。
三、代碼實現解析
1. 輸入數據結構
輸入為字典類型 predicts_dict
,鍵為類別名稱,值為該類別對應的邊界框列表。每個邊界框格式為 [x1, y1, x2, y2, score]
,表示左上角和右下角坐標及置信度。
predicts_dict = {'black1': [[83,54,165,163,0.8], [67,48,118,132,0.5], ...]}
2. 核心函數 non_max_suppress
def non_max_suppress(predicts_dict, threshold):for object_name, bbox in predicts_dict.items():bbox_array = np.array(bbox, dtype=float)# 提取坐標和置信度x1, y1, x2, y2, score = bbox_array[:,0], bbox_array[:,1], bbox_array[:,2], bbox_array[:,3], bbox_array[:,4]# 按置信度降序排序order = score.argsort()[::-1]area = (x2 - x1 + 1) * (y2 - y1 + 1)keep = [] # 保留的索引列表while order.size > 0:i = order[0] # 當前最高分框keep.append(i)# 計算IoUxx1 = np.maximum(x1[i], x1[order[1:]])yy1 = np.maximum(y1[i], y1[order[1:]])xx2 = np.minimum(x2[i], x2[order[1:]])yy2 = np.minimum(y2[i], y2[order[1:]])inter = np.maximum(0.0, xx2 - xx1 + 1) * np.maximum(0.0, yy2 - yy1 + 1)iou = inter / (area[i] + area[order[1:]] - inter)# 保留IoU低于閾值的框inds = np.where(iou <= threshold)[0]order = order[inds + 1]# 更新篩選后的結果predicts_dict[object_name] = bbox_array[keep].tolist()return predicts_dict
關鍵步驟說明:
- 坐標提取與排序:將邊界框轉換為NumPy數組后,按置信度降序排列。
- IoU計算:通過最大-最小值法計算交集區域,公式為:
IoU = Intersection Union ? Intersection \text{IoU} = \frac{\text{Intersection}}{\text{Union} - \text{Intersection}} IoU=Union?IntersectionIntersection? - 動態索引更新:通過
order = order[inds + 1]
跳過被抑制的框,逐步縮小處理范圍。
3. 可視化測試代碼
- 繪制原始預測框:在全黑圖像上繪制未經過NMS處理的邊界框及置信度。
- NMS處理與對比:調用
non_max_suppress
后,在另一窗口展示抑制后的結果。
# 繪制原始框
for box in bbox:cv2.rectangle(img, (x1, y1), (x2, y2), (255,255,255), 2)
# 處理并繪制NMS后的框
predicts_dict_nms = non_max_suppress(predicts_dict, 0.1)
for box in bbox_nms:cv2.rectangle(img_cp, (x1, y1), (x2, y2), (255,255,255), 2)
四、優化與注意事項
- 閾值選擇:閾值過小可能導致漏檢,過大則冗余框增多(通常目標檢測任務中閾值設為0.5)。
- 多類別處理:代碼支持同時對多個類別獨立進行NMS,如輸入
black1
和black2
兩個類別的預測結果。 - 坐標修正:代碼中
+1
的操作是為了避免零寬度/高度,確保面積計算正確。
import cv2
import random
import numpy as npdef non_max_suppress(predicts_dict, threshold):for object_name, bbox in predicts_dict.items(): # 對每一個類別分別進行NMS;一次讀取一對鍵值(即某個類別的所有框)bbox_array = np.array(bbox, dtype=np.float)print(bbox_array)# 下面分別獲取框的左上角坐標(x1,y1),右下角坐標(x2,y2)及此框的置信度;這里需要注意的是圖像左上角可以看做坐標點(0,0),右下角可以看做坐標點(1,1),也就是說從左往右x值增大,從上往下y值增大x1 = bbox_array[:, 0]y1 = bbox_array[:, 1]x2 = bbox_array[:, 2]y2 = bbox_array[:, 3]scores = bbox_array[:, 4] # class confidence, ndarrayprint(scores, type(scores)) order = scores.argsort()[::-1] # argsort函數返回的是數組值從小到大的索引值,[::-1]表示取反。即這里返回的是數組值從大到小的索引值areas = (x2 - x1 + 1) * (y2 - y1 + 1) # 當前類所有框的面積(python會自動使用廣播機制,相當于MATLAB中的.*即兩矩陣對應元素相乘);x1=3,x2=5,習慣上計算x方向長度就是x=3、4、5這三個像素,即5-3+1=3,而不是5-3=2,所以需要加1print(areas, type(areas)) keep = []# 按confidence從高到低遍歷bbx,移除所有與該矩形框的IoU值大于threshold的矩形框while order.size > 0:i = order[0]keep.append(i) # 保留當前最大confidence對應的bbx索引# 獲取所有與當前bbx的交集對應的左上角和右下角坐標,并計算IoU(注意這里是同時計算一個bbx與其他所有bbx的IoU)xx1 = np.maximum(x1[i], x1[order[1:]]) # 最大置信度的左上角坐標分別與剩余所有的框的左上角坐標進行比較,分別保存較大值;因此這里的xx1的維數應該是當前類的框的個數減1print("xx1:", xx1)yy1 = np.maximum(y1[i], y1[order[1:]])xx2 = np.minimum(x2[i], x2[order[1:]])yy2 = np.minimum(y2[i], y2[order[1:]])inter = np.maximum(0.0, xx2-xx1+1) * np.maximum(0.0, yy2-yy1+1)iou = inter / (areas[i] + areas[order[1:]] - inter) # 注意這里都是采用廣播機制,同時計算了置信度最高的框與其余框的IoUprint(iou, type(iou))print(np.where(iou <= threshold))inds = np.where(iou <= threshold)[0] # 保留iou小于等于闕值的框的索引值print('inds:', inds)order = order[inds + 1] # 將order中的第inds+1處的值重新賦值給order;即更新保留下來的索引,加1是因為因為沒有計算與自身的IOU,所以索引相差1,需要加上bbox = bbox_array[keep]predicts_dict[object_name] = bbox.tolist()return predicts_dict# 下面在一張全黑圖片上測試非極大值抑制的效果
img = np.zeros((600,600), np.uint8)
predicts_dict = {'black1': [[83, 54, 165, 163, 0.8], [67, 48, 118, 132, 0.5], [91, 38, 192, 171, 0.6]]}
# predicts_dict = {'black1': [[83, 54, 165, 163, 0.8], [67, 48, 118, 132, 0.5], [91, 38, 192, 171, 0.6]], 'black2': [[59, 120, 137, 368, 0.12], [54, 154, 148, 382, 0.13]] }
"""
# 在全黑的圖像上畫出設定的幾個框
for object_name, bbox in predicts_dict.items():for box in bbox:x1, y1, x2, y2, score = box[0], box[1], box[2], box[3], box[-1]y_text = int(random.uniform(y1, y2)) # uniform()是不能直接訪問的,需要導入 random 模塊,然后通過 random 靜態對象調用該方法。uniform() 方法將隨機生成下一個實數,它在 [x, y) 范圍內cv2.rectangle(img, (x1, y1), (x2, y2), (255, 255, 255), 2)cv2.putText(img, str(score), (x2 - 30, y_text), 2, 1, (255, 255, 0))cv2.namedWindow("black1_roi") # 創建一個顯示圖像的窗口cv2.imshow("black1_roi", img) # 在窗口中顯示圖像;注意這里的窗口名字如果不是剛剛創建的窗口的名字則會自動創建一個新的窗口并將圖像顯示在這個窗口cv2.waitKey(0) # 如果不添這一句,在IDLE中執行窗口直接無響應。在命令行中執行的話,則是一閃而過。
cv2.destroyAllWindows() # 最后釋放窗口是個好習慣!
"""
# 在全黑圖片上畫出經過非極大值抑制后的框
img_cp = np.zeros((600,600), np.uint8)
predicts_dict_nms = non_max_suppress(predicts_dict, 0.1)
for object_name, bbox in predicts_dict_nms.items():for box in bbox:x1, y1, x2, y2, score = int(box[0]), int(box[1]), int(box[2]), int(box[3]), box[-1]y_text = int(random.uniform(y1, y2)) # uniform()是不能直接訪問的,需要導入 random 模塊,然后通過 random 靜態對象調用該方法。uniform() 方法將隨機生成下一個實數,它在 [x, y) 范圍內cv2.rectangle(img_cp, (x1, y1), (x2, y2), (255, 255, 255), 2)cv2.putText(img_cp, str(score), (x2 - 30, y_text), 2, 1, (255, 255, 0))cv2.namedWindow("black1_nms") # 創建一個顯示圖像的窗口cv2.imshow("black1_nms", img_cp) # 在窗口中顯示圖像;注意這里的窗口名字如果不是剛剛創建的窗口的名字則會自動創建一個新的窗口并將圖像顯示在這個窗口cv2.waitKey(0) # 如果不添這一句,在IDLE中執行窗口直接無響應。在命令行中執行的話,則是一閃而過。
cv2.destroyAllWindows() # 最后釋放窗口是個好習慣!