前言
SSD 的神經網絡結構很簡潔，可以較好的實現多尺度的目標檢測，但是對小目標物體的檢測效果並不是很好。雖然有很多 SSD 的魔改版本，比如 FSSD 和 DSSD，提高了 SSD 在小目標檢測上的表現，但是這里我們只討論怎么使用 SSD 來更好地檢測小目標，尤其是那些特征非常簡單的目標。
YOLO 的啟發
在 Yolo V3 中使用了先驗框聚類的方式來決定先驗框的尺寸，而在 SSD 的原始版本中是通過公式來決定先驗框的尺寸，最小的先驗框尺寸都有 30。如果我們的目標很小，比如只有十幾像素，那么使用這些先驗框訓練出來的 SSD 模型的表現大概率是差強人意的。所以我們可以在自己的數據集上對先驗框進行聚類，下面給出聚類的代碼：
# coding:utf-8
from pathlib import Path
from xml.etree import ElementTree as ET
import numpy as np
def iou(box: np.ndarray, boxes: np.ndarray):
""" 計算一個邊界框和其他邊界框的交並比
Parameters
----------
box: `~np.ndarray` of shape `(4, )`
邊界框
boxes: `~np.ndarray` of shape `(n, 4)`
其他邊界框
Returns
-------
iou: `~np.ndarray` of shape `(n, )`
交並比
"""
# 計算交集
xy_max = np.minimum(boxes[:, 2:], box[2:])
xy_min = np.maximum(boxes[:, :2], box[:2])
inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
inter = inter[:, 0]*inter[:, 1]
# 計算並集
area_boxes = (boxes[:, 2]-boxes[:, 0])*(boxes[:, 3]-boxes[:, 1])
area_box = (box[2]-box[0])*(box[3]-box[1])
# 計算 iou
iou = inter/(area_box+area_boxes-inter) # type: np.ndarray
return iou
class AnchorKmeans:
""" 先驗框聚類 """
def __init__(self, annotation_dir: str):
self.annotation_dir = Path(annotation_dir)
if not self.annotation_dir.exists():
raise ValueError(f'標簽文件夾 `{annotation_dir}` 不存在')
self.bbox = self.get_bbox()
def get_bbox(self) -> np.ndarray:
""" 獲取所有的邊界框 """
bbox = []
for path in self.annotation_dir.glob('*.xml'):
root = ET.parse(path).getroot()
# 圖像的寬度和高度
w = int(root.find('size/width').text)
h = int(root.find('size/height').text)
# 獲取所有邊界框
for obj in root.iter('object'):
box = obj.find('bndbox')
# 歸一化坐標
xmin = int(box.find('xmin').text)/w
ymin = int(box.find('ymin').text)/h
xmax = int(box.find('xmax').text)/w
ymax = int(box.find('ymax').text)/h
bbox.append([0, 0, xmax-xmin, ymax-ymin])
return np.array(bbox)
def get_cluster(self, n_clusters=9, metric=np.median):
""" 獲取聚類結果
Parameters
----------
n_clusters: int
聚類數
metric: callable
選取聚類中心點的方式
"""
rows = self.bbox.shape[0]
if rows < n_clusters:
raise ValueError("n_clusters 不能大於邊界框樣本數")
last_clusters = np.zeros(rows)
clusters = np.ones((n_clusters, 2))
distances = np.zeros((rows, n_clusters)) # type:np.ndarray
# 隨機選取出幾個點作為聚類中心
np.random.seed(1)
clusters = self.bbox[np.random.choice(rows, n_clusters, replace=False)]
# 開始聚類
while True:
# 計算距離
distances = 1-self.iou(clusters)
# 將每一個邊界框划到一個聚類中
nearest_clusters = distances.argmin(axis=1)
# 如果聚類中心不再變化就退出
if np.array_equal(nearest_clusters, last_clusters):
break
# 重新選取聚類中心
for i in range(n_clusters):
clusters[i] = metric(self.bbox[nearest_clusters == i], axis=0)
last_clusters = nearest_clusters
return clusters[:, 2:]
def average_iou(self, clusters: np.ndarray):
""" 計算 IOU 均值
Parameters
----------
clusters: `~np.ndarray` of shape `(n_clusters, 2)`
聚類中心
"""
clusters = np.hstack((np.zeros((clusters.shape[0], 2)), clusters))
return np.mean([np.max(iou(bbox, clusters)) for bbox in self.bbox])
def iou(self, clusters: np.ndarray):
""" 計算所有邊界框和所有聚類中心的交並比
Parameters
----------
clusters: `~np.ndarray` of shape `(n_clusters, 4)`
聚類中心
Returns
-------
iou: `~np.ndarray` of shape `(n_bbox, n_clusters)`
交並比
"""
bbox = self.bbox
A = self.bbox.shape[0]
B = clusters.shape[0]
xy_max = np.minimum(bbox[:, np.newaxis, 2:].repeat(B, axis=1),
np.broadcast_to(clusters[:, 2:], (A, B, 2)))
xy_min = np.maximum(bbox[:, np.newaxis, :2].repeat(B, axis=1),
np.broadcast_to(clusters[:, :2], (A, B, 2)))
# 計算交集面積
inter = np.clip(xy_max-xy_min, a_min=0, a_max=np.inf)
inter = inter[:, :, 0]*inter[:, :, 1]
# 計算每個矩陣的面積
area_bbox = ((bbox[:, 2]-bbox[:, 0])*(bbox[:, 3] -
bbox[:, 1]))[:, np.newaxis].repeat(B, axis=1)
area_clusters = ((clusters[:, 2] - clusters[:, 0])*(
clusters[:, 3] - clusters[:, 1]))[np.newaxis, :].repeat(A, axis=0)
return inter/(area_bbox+area_clusters-inter)
if __name__ == '__main__':
# 標簽文件夾
root = 'data/Hotspot/Annotations'
model = AnchorKmeans(root)
clusters = model.get_cluster(9)
# 將先驗框還原為原本的大小
print('聚類結果:\n', clusters*300)
print('平均 IOU:', model.average_iou(clusters))
將代碼中的先驗框尺寸參照聚類的結果進行修改，不出意外的話是可以提升 mAP 和置信度的，以上~~