YOLOv3使用自己数据集——Kmeans聚类计算anchor boxes

使用自己数据集聚类得到anchors。
相比于作者使用VOC数据集的精度更高。

# kmeans 聚类计算anchor boxes

import glob
import xml.etree.ElementTree as ET

import numpy as np

from kmeans import kmeans, avg_iou
 
# 标签文件地址
ANNOTATIONS_PATH = "/home/peter/桌面/项目文档/相机-雷达识别/724DATA/label_png"
CLUSTERS = 9


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")) / width
            ymin = int(obj.findtext("bndbox/ymin")) / height
            xmax = int(obj.findtext("bndbox/xmax")) / width
            ymax = int(obj.findtext("bndbox/ymax")) / height
            
            xmin = np.float64(xmin)
            ymin = np.float64(ymin)
            xmax = np.float64(xmax)
            ymax = np.float64(ymax)
            if xmax == xmin or ymax == ymin:
                print(xml_file)
            dataset.append([xmax - xmin, ymax - ymin])
    return np.array(dataset)
 
if __name__ == '__main__':
    #print(__file__)
    ''''
    #使用10次循环去均值，但是所得精度82%反而不如一次的88%高
    for i in range(10):
        data = load_dataset(ANNOTATIONS_PATH)
        out = kmeans(data, k=CLUSTERS)
        if i == 0:
            out1 = out
        else:
            out1 = out1 + out
        print(i)
    out1 = out1 / 10
    #clusters = [[10,13],[16,30],[33,23],[30,61],[62,45],[59,119],[116,90],[156,198],[373,326]]
    #out= np.array(clusters)/608.0
    out = out1
    '''
    data = load_dataset(ANNOTATIONS_PATH)
    out = kmeans(data, k = CLUSTERS)
    # 输入图片大小为width = 608, height = 608
    print('out*608:')
    out1 = out * 608
    out1.astype(np.int32)
    print(out1)
    print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100))
    print("Boxes:\n {}-{}".format(out[:, 0] * 608, out[:, 1] * 608))
   
    ratios = np.around(out[:, 0] / out[:, 1], decimals = 2).tolist()
    print("Ratios:\n {}".format(sorted(ratios)))

输出结果：

out*608:
[[19.37481032 10.93853073]
 [13.83915023 10.93853073]
 [18.90633295 13.67316342]
 [24.91047041 10.93853073]
 [21.21198331  9.11544228]
 [16.37011756 12.30584708]
 [14.99241275  8.65967016]
 [22.60394537 12.76161919]
 [16.83124763 10.02698651]]
Accuracy: 88.91%
Boxes:
 [19.37481032 13.83915023 18.90633295 24.91047041 21.21198331 16.37011756
 14.99241275 22.60394537 16.83124763]-[10.93853073 10.93853073 13.67316342 10.93853073  9.11544228 12.30584708
  8.65967016 12.76161919 10.02698651]
Ratios:

最终的anchros及输出结果：

out.astype(np.int32)

array([[14, 10],
       [14,  9],
       [25, 10],
       [20,  9],
       [17, 10],
       [17, 12],
       [17,  8],
       [20, 14],
       [22, 11]], dtype=int32)

将上面int32的内容替换yolov3.cfg中的anchros保存即可。