如何处理 SSD 神经网络在小目标检测数据集上 mAP 和置信度较低的问题
前言
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 和置信度的，以上~~
