k-mean聚类获取anchor的先验大小

最新推荐文章于 2023-11-26 18:19:05 发布

原创最新推荐文章于 2023-11-26 18:19:05 发布 · 5.7k 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#yolo

yolo 专栏收录该内容

4 篇文章

订阅专栏

本文探讨了聚类算法在目标检测任务中的应用，特别是如何利用k-means算法生成YOLOv2和YOLOv3所需的anchor boxes。通过自定义距离度量公式，确保聚类结果能提高检测框与真实框的IOU值，从而提升模型的检测精度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

方法1：欧式距离

#!/usr/bin/env python
# -*- coding: utf8 -*-
import sys
from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement
from lxml import etree
import numpy as np
import os
import sys
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

XML_EXT = '.xml'
ENCODE_METHOD = 'utf-8'

#pascalVocReader readers the voc xml files parse it
class PascalVocReader:
    """
    this class will be used to get transfered width and height from voc xml files
    """
    def __init__(self, filepath,width,height):
        # shapes type:
        # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult]
        self.shapes = []
        self.filepath = filepath
        self.verified = False
        self.width=width
        self.height=height

        try:
            self.parseXML()
        except:
            pass

    def getShapes(self):
        return self.shapes

    def addShape(self, bndbox, width,height):
        xmin = int(bndbox.find('xmin').text)
        ymin = int(bndbox.find('ymin').text)
        xmax = int(bndbox.find('xmax').text)
        ymax = int(bndbox.find('ymax').text)
        width_trans = (xmax - xmin)/width*self.width
        height_trans = (ymax-ymin)/height *self.height
        points = [width_trans,height_trans]
        self.shapes.append((points))

    def parseXML(self):
        assert self.filepath.endswith(XML_EXT), "Unsupport file format"
        parser = etree.XMLParser(encoding=ENCODE_METHOD)
        xmltree = ElementTree.parse(self.filepath, parser=parser).getroot()
        pic_size = xmltree.find('size')
        size = (int(pic_size.find('width').text),int(pic_size.find('height').text))
        for object_iter in xmltree.findall('object'):
            bndbox = object_iter.find("bndbox")
            self.addShape(bndbox, *size)
        return True

class create_w_h_txt:
    def __init__(self,vocxml_path,txt_path):
        self.voc_path = vocxml_path
        self.txt_path = txt_path
    def _gether_w_h(self):
        pass
    def _write_to_txt(self):
        pass
    def process_file(self):
        file_w = open(self.txt_path,'a')
       # print (self.txt_path)
        for file in os.listdir(self.voc_path):
            file_path = os.path.join(self.voc_path, file)
            xml_parse = PascalVocReader(file_path,1024,1024)
            data = xml_parse.getShapes()
            for w,h in data :
                txtstr = str(w)+' '+str(h)+'\n'
                #print (txtstr)
                file_w.write(txtstr)
        file_w.close()

class kMean_parse:
    def __init__(self,path_txt):
        self.path = path_txt
        self.km = KMeans(n_clusters=5,init="k-means++",n_init=10,max_iter=3000000,tol=1e-3,random_state=0)
        self._load_data()

    def _load_data (self):
        self.data = np.loadtxt(self.path)

    def parse_data (self):
        self.y_k = self.km.fit_predict(self.data)
        print(self.km.cluster_centers_)

    def plot_data (self):
        plt.scatter(self.data[self.y_k == 0, 0], self.data[self.y_k == 0, 1], s=50, c="orange", marker="o", label="cluster 1")
        plt.scatter(self.data[self.y_k == 1, 0], self.data[self.y_k == 1, 1], s=50, c="green", marker="s", label="cluster 2")
        plt.scatter(self.data[self.y_k == 2, 0], self.data[self.y_k == 2, 1], s=50, c="blue", marker="^", label="cluster 3")
        plt.scatter(self.data[self.y_k == 3, 0], self.data[self.y_k == 3, 1], s=50, c="gray", marker="*",label="cluster 4")
        plt.scatter(self.data[self.y_k == 4, 0], self.data[self.y_k == 4, 1], s=50, c="yellow", marker="d",label="cluster 5")
       # draw the centers
        plt.scatter(self.km.cluster_centers_[:, 0], self.km.cluster_centers_[:, 1], s=250, marker="*", c="red", label="cluster center")
        plt.legend()
        plt.grid()
        plt.show()




if __name__ == '__main__':
     whtxt = create_w_h_txt("C:\\Users\\wytwh\\Desktop\\newlabel\\Annotations","./data1.txt") #指定为voc标注路径，以及存放生成文件路径
     whtxt.process_file()
     kmean_parse = kMean_parse("./data1.txt")#路径和生成文件相同。
     kmean_parse.parse_data()
     kmean_parse.plot_data() #绘图部分只支持五个簇，要增加，需要自家改代码即可

方法2：yolov2、yolov3所使用的聚类

k-means需要有数据，中心点个数是需要人为指定的，位置可以随机初始化，但是还需要度量到聚类中心的距离。这里怎么度量这个距离是很关键的。
距离度量如果是方法1使用标准的欧氏距离，大盒子会比小盒子产生更多的错误。例 (100-95)^2=25, (5-2.5)^2=6.25 。因此这里使用其他的距离度量公式。聚类的目的是anchor boxes和临近的ground truth有更大的IOU值，这和anchor box的尺寸没有直接关系。自定义的距离度量公式：

d(box,centroid)=1-IOU(box,centroid)

到聚类中心的距离越小越好，但IOU值是越大越好，所以使用 1 - IOU，这样就保证距离越小，IOU值越大。

具体步骤：

使用的聚类原始数据是只有标注框的检测数据集，YOLOv2、v3都会生成一个包含标注框位置和类别的TXT文件，其中每行都包含，即ground truth boxes相对于原图的坐标，是框的中心点，是框的宽和高，N是所有标注框的个数；
首先给定k个聚类中心点，这里的是anchor boxes的宽和高尺寸，由于anchor boxes位置不固定，所以没有(x,y)的坐标，只有宽和高；
计算每个标注框和每个聚类中心点的距离 d=1-IOU(标注框,聚类中心)，计算时每个标注框的中心点都与聚类中心重合，这样才能计算IOU值，即。将标注框分配给“距离”最近的聚类中心；
所有标注框分配完毕以后，对每个簇重新计算聚类中心点，计算方式为，是第i个簇的标注框个数，就是求该簇中所有标注框的宽和高的平均值。
重复第3、4步，直到聚类中心改变量很小。

原工程：https://github.com/lars76/kmeans-anchor-boxes

import glob
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import numpy as np

from kmeans import kmeans, avg_iou



def plot_data ():
    for i in range(CLUSTERS):
        color = ['orange','green','blue','gray','yellow','purple','pink','black','brown']
        mark= ['o','s','^','*','d','+','x','p','|']
        lab = 'claster' + str(i+1)
        plt.scatter(data[index == i, 0], data[index == i, 1], s=10, c=color[i], marker=mark[i], label=lab)
       # draw the centers
    plt.scatter(out[:, 0], out[:, 1], s=250, marker="*", c="red", label="cluster center")
    plt.legend()
    plt.grid()
    plt.show()
        
        
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

			dataset.append([xmax - xmin, ymax - ymin])

	return np.array(dataset)


if __name__ == "__main__":
    ANNOTATIONS_PATH = "C:\\Users\\wytwh\\Desktop\\newlabel\\Annotations"
    CLUSTERS = 9
    data = load_dataset(ANNOTATIONS_PATH)
    out, index = kmeans(data, k=CLUSTERS)
    print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100))
    print("Boxes:\n {}".format(out * 1024))
    
    ratios = np.around(out[:, 0] / out[:, 1], decimals=2).tolist()
    print("Ratios:\n {}".format(sorted(ratios)))
    plot_data ()