本文主要是介绍YOLO的Anchor聚类代码，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
代码来源于GiantPandaCV ，作者BBuf
#coding=utf-8import xml.etree.ElementTree as ET
import numpy as npdef iou(box, clusters):"""计算一个ground truth边界盒和k个先验框(Anchor)的交并比(IOU)值。参数box: 元组或者数据，代表ground truth的长宽。参数clusters: 形如(k,2)的numpy数组，其中k是聚类Anchor框的个数返回：ground truth和每个Anchor框的交并比。"""x = np.minimum(clusters[:, 0], box[0])y = np.minimum(clusters[:, 1], box[1])if np.count_nonzero(x == 0) > 0 or np.count_nonzero(y == 0) > 0:raise ValueError("Box has no area")intersection = x * ybox_area = box[0] * box[1]cluster_area = clusters[:, 0] * clusters[:, 1]iou_ = intersection / (box_area + cluster_area - intersection)return iou_def avg_iou(boxes, clusters):"""计算一个ground truth和k个Anchor的交并比的均值。"""return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])def kmeans(boxes, k, dist=np.median):"""利用IOU值进行K-means聚类参数boxes: 形状为(r, 2)的ground truth框，其中r是ground truth的个数参数k: Anchor的个数参数dist: 距离函数返回值：形状为(k, 2)的k个Anchor框"""# 即是上面提到的rrows = boxes.shape[0]# 距离数组，计算每个ground truth和k个Anchor的距离distances = np.empty((rows, k))# 上一次每个ground truth"距离"最近的Anchor索引last_clusters = np.zeros((rows,))# 设置随机数种子np.random.seed()# 初始化聚类中心，k个簇，从r个ground truth随机选k个clusters = boxes[np.random.choice(rows, k, replace=False)]# 开始聚类while True:# 计算每个ground truth和k个Anchor的距离，用1-IOU(box,anchor)来计算for row in range(rows):distances[row] = 1 - iou(boxes[row], clusters)# 对每个ground truth，选取距离最小的那个Anchor，并存下索引nearest_clusters = np.argmin(distances, axis=1)# 如果当前每个ground truth"距离"最近的Anchor索引和上一次一样，聚类结束if (last_clusters == nearest_clusters).all():break# 更新簇中心为簇里面所有的ground truth框的均值for cluster in range(k):clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)# 更新每个ground truth"距离"最近的Anchor索引last_clusters = nearest_clustersreturn clusters# 加载自己的数据集，只需要所有labelimg标注出来的xml文件即可
def load_dataset(path):dataset = []for xml_file in glob.glob("{}/*xml".format(path)):tree = ET.parse(xml_file)# 图片高度height = int(tree.findtext("./size/height"))# 图片宽度width = int(tree.findtext("./size/width"))for obj in tree.iter("object"):# 偏移量xmin = int(obj.findtext("bndbox/xmin")) / widthymin = int(obj.findtext("bndbox/ymin")) / heightxmax = int(obj.findtext("bndbox/xmax")) / widthymax = int(obj.findtext("bndbox/ymax")) / heightxmin = np.float64(xmin)ymin = np.float64(ymin)xmax = np.float64(xmax)ymax = np.float64(ymax)if xmax == xmin or ymax == ymin:print(xml_file)# 将Anchor的长宽放入dateset，运行kmeans获得Anchordataset.append([xmax - xmin, ymax - ymin])return np.array(dataset)if __name__ == '__main__':ANNOTATIONS_PATH = "F:\Annotations" #xml文件所在文件夹CLUSTERS = 9 #聚类数量，anchor数量INPUTDIM = 416 #输入网络大小data = load_dataset(ANNOTATIONS_PATH)out = kmeans(data, k=CLUSTERS)print('Boxes:')print(np.array(out)*INPUTDIM)print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100))final_anchors = np.around(out[:, 0] / out[:, 1], decimals=2).tolist()print("Before Sort Ratios:\n {}".format(final_anchors))print("After Sort Ratios:\n {}".format(sorted(final_anchors)))
这篇关于YOLO的Anchor聚类代码的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
