1 問題現象
使用yolo v3 等目標檢測模型訓練自己數據集,預測圖片時出現問題: 兩框重疊,如下圖所示:對於同樣一輛汽車,模型反覆的標記。
2 解決辦法
解決辦法就是:非極大值抑制(Non-Maximum Suppression)
3 Non-Maximum Suppression 原理
3.1 什麼是非極大值抑制
非極大值抑制,簡稱爲NMS算法,英文爲Non-Maximum Suppression。其思想是搜素局部最大值,抑制極大值。NMS算法在不同應用中的具體實現不太一樣,但思想是一樣的。非極大值抑制,在計算機視覺任務中得到了廣泛的應用,例如邊緣檢測、人臉檢測、目標檢測(DPM,YOLO,SSD,Faster R-CNN)等。
3.2 爲什麼要用非極大值抑制
以目標檢測爲例:目標檢測的過程中在同一目標的位置上會產生大量的候選框,這些候選框相互之間可能會有重疊,此時我們需要利用非極大值抑制找到最佳的目標邊界框,消除冗餘的邊界框。Demo如下圖:
左圖是人臉檢測的候選框結果,每個邊界框有一個置信度得分(confidence score),如果不使用非極大值抑制,就會有多個候選框出現。右圖是使用非極大值抑制之後的結果,符合我們人臉檢測的預期結果。
3.3 如何使用非極大值抑制
前提:目標邊界框列表及其對應的置信度得分列表,設定閾值,閾值用來刪除重疊較大的邊界框。
IoU:intersection-over-union,即兩個邊界框的交集部分除以它們的並集。
非極大值抑制的流程如下:
- 根據置信度得分進行排序
- 選擇置信度最高的比邊界框添加到最終輸出列表中,將其從邊界框列表中刪除
- 計算所有邊界框的面積
- 計算置信度最高的邊界框與其它候選框的IoU。
- 刪除IoU大於閾值的邊界框
- 重複上述過程,直至邊界框列表爲空。
Python代碼如下:
#!/usr/bin/env python
# _*_ coding: utf-8 _*_
import cv2
import numpy as np
"""
Non-max Suppression Algorithm
@param list Object candidate bounding boxes
@param list Confidence score of bounding boxes
@param float IoU threshold
@return Rest boxes after nms operation
"""
def nms(bounding_boxes, confidence_score, threshold):
# If no bounding boxes, return empty list
if len(bounding_boxes) == 0:
return [], []
# Bounding boxes
boxes = np.array(bounding_boxes)
# coordinates of bounding boxes
start_x = boxes[:, 0]
start_y = boxes[:, 1]
end_x = boxes[:, 2]
end_y = boxes[:, 3]
# Confidence scores of bounding boxes
score = np.array(confidence_score)
# Picked bounding boxes
picked_boxes = []
picked_score = []
# Compute areas of bounding boxes
areas = (end_x - start_x + 1) * (end_y - start_y + 1)
# Sort by confidence score of bounding boxes
order = np.argsort(score)
# Iterate bounding boxes
while order.size > 0:
# The index of largest confidence score
index = order[-1]
# Pick the bounding box with largest confidence score
picked_boxes.append(bounding_boxes[index])
picked_score.append(confidence_score[index])
# Compute ordinates of intersection-over-union(IOU)
x1 = np.maximum(start_x[index], start_x[order[:-1]])
x2 = np.minimum(end_x[index], end_x[order[:-1]])
y1 = np.maximum(start_y[index], start_y[order[:-1]])
y2 = np.minimum(end_y[index], end_y[order[:-1]])
# Compute areas of intersection-over-union
w = np.maximum(0.0, x2 - x1 + 1)
h = np.maximum(0.0, y2 - y1 + 1)
intersection = w * h
# Compute the ratio between intersection and union
ratio = intersection / (areas[index] + areas[order[:-1]] - intersection)
left = np.where(ratio < threshold)
order = order[left]
return picked_boxes, picked_score
# Image name
image_name = 'nms.jpg'
# Bounding boxes
bounding_boxes = [(187, 82, 337, 317), (150, 67, 305, 282), (246, 121, 368, 304)]
confidence_score = [0.9, 0.75, 0.8]
# Read image
image = cv2.imread(image_name)
# Copy image as original
org = image.copy()
# Draw parameters
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 1
thickness = 2
# IoU threshold
threshold = 0.4
# Draw bounding boxes and confidence score
for (start_x, start_y, end_x, end_y), confidence in zip(bounding_boxes, confidence_score):
(w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
cv2.rectangle(org, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
cv2.rectangle(org, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
cv2.putText(org, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)
# Run non-max suppression algorithm
picked_boxes, picked_score = nms(bounding_boxes, confidence_score, threshold)
# Draw bounding boxes and confidence score after non-maximum supression
for (start_x, start_y, end_x, end_y), confidence in zip(picked_boxes, picked_score):
(w, h), baseline = cv2.getTextSize(str(confidence), font, font_scale, thickness)
cv2.rectangle(image, (start_x, start_y - (2 * baseline + 5)), (start_x + w, start_y), (0, 255, 255), -1)
cv2.rectangle(image, (start_x, start_y), (end_x, end_y), (0, 255, 255), 2)
cv2.putText(image, str(confidence), (start_x, start_y), font, font_scale, (0, 0, 0), thickness)
# Show image
cv2.imshow('Original', org)
cv2.imshow('NMS', image)
cv2.waitKey(0)
3.4 效果
具體解決辦法就是減小 IoU threshold (IoU 閾值)
IoU閾值爲0.6的時候:
IoU閾值爲0.4的時候:
4 參考資料:
- 《非極大值抑制(Non-Maximum Suppression)》:https://zhuanlan.zhihu.com/p/37489043
- 《訓練自己數據集,預測圖片時出現問題》https://github.com/qqwweee/keras-yolo3/issues/354