前言
SSD 的神经网络结构很简洁,可以较好的实现多尺度的目标检测,但是对小目标物体的检测效果并不是很好。虽然有很多 SSD 的魔改版本,比如 FSSD 和 DSSD,提高了 SSD 在小目标检测上的表现,但是这里我们只讨论怎么使用 SSD 来更好地检测小目标,尤其是那些特征非常简单的目标。
YOLO 的启发
在 Yolo V3 中使用了先验框聚类的方式来决定先验框的尺寸,而在 SSD 的原始版本中是通过公式来决定先验框的尺寸,最小的先验框尺寸都有 30。如果我们的目标很小,比如只有十几像素,那么使用这些先验框训练出来的 SSD 模型的表现大概率是差强人意的。所以我们可以在自己的数据集上对先验框进行聚类,下面给出聚类的代码:
# coding:utf-8from pathlib import Pathfrom xml.etree import ElementTree as ETimport numpy as npdef 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])# 计算 iouiou = inter/(area_box+area_boxes-inter) # type: np.ndarrayreturn iouclass 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)/wymin = int(box.find(\'ymin\').text)/hxmax = int(box.find(\'xmax\').text)/wymax = int(box.find(\'ymax\').text)/hbbox.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_clustersreturn 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.bboxA = 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 和置信度的,以上~~