yolov2/v3 k-means计算anchor源码

YOLOv2开始增加了anchor机制，在v3中增加到9个anchor。例如yolov3-voc.cfg中这组anchor，anchors = 10,13,  16,30,  33,23,  30,61,  62,45,  59,119,  116,90,  156,198,  373,326，这些anchor size 是由 Redmon大神在VOC的20类跟COCO的80类中聚类的结果，适用较广但如果自己拿来train实际类别的时候，由于数据集特点不同所以这些anchor并不能完全很好的覆盖数据集中的anchor导致检测效率降低，这里贴出如何通过k-means计算自己训练集中的anchor box。提供VOC和COCO两种类型的计算。代码如下：

k-means.py:

 

import numpy as np


def iou(box, clusters):
    """
    Calculates the Intersection over Union (IoU) between a box and k clusters.
    :param box: tuple or array, shifted to the origin (i. e. width and height)
    :param clusters: numpy array of shape (k, 2) where k is the number of clusters
    :return: numpy array of shape (k, 0) where k is the number of clusters
    """
    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 * y
    box_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):
    """
    Calculates the average Intersection over Union (IoU) between a numpy array of boxes and k clusters.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param clusters: numpy array of shape (k, 2) where k is the number of clusters
    :return: average IoU as a single float
    """
    return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])


def translate_boxes(boxes):
    """
    Translates all the boxes to the origin.
    :param boxes: numpy array of shape (r, 4)
    :return: numpy array of shape (r, 2)
    """
    new_boxes = boxes.copy()
    for row in range(new_boxes.shape[0]):
        new_boxes[row][2] = np.abs(new_boxes[row][2] - new_boxes[row][0])
        new_boxes[row][3] = np.abs(new_boxes[row][3] - new_boxes[row][1])
    return np.delete(new_boxes, [0, 1], axis=1)


def kmeans(boxes, k, dist=np.median):
    """
    Calculates k-means clustering with the Intersection over Union (IoU) metric.
    :param boxes: numpy array of shape (r, 2), where r is the number of rows
    :param k: number of clusters
    :param dist: distance function
    :return: numpy array of shape (k, 2)
    """
    rows = boxes.shape[0]

    distances = np.empty((rows, k))
    last_clusters = np.zeros((rows,))

    np.random.seed()

    # the Forgy method will fail if the whole array contains the same rows
    clusters = boxes[np.random.choice(rows, k, replace=False)]

    while True:
        for row in range(rows):
            distances[row] = 1 - iou(boxes[row], clusters)

        nearest_clusters = np.argmin(distances, axis=1)

        if (last_clusters == nearest_clusters).all():
            break

        for cluster in range(k):
            clusters[cluster] = dist(boxes[nearest_clusters == cluster], axis=0)

        last_clusters = nearest_clusters

    return clusters

get_Result.py:

import glob
import xml.etree.ElementTree as ET

import numpy as np
import json

from kmeans import kmeans, avg_iou

ANNOTATIONS_PATH = "Annotations"
JSON_PATH1 = ""
JSON_PATH2 = ""
JSON_PATH3 = ""
CLUSTERS = 9


def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]

    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def load_voc_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)

def load_coco_dataset(json_path):
    dataset = []
    data = json.load(open(json_path,'r'))

    for img in data['images']:
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = ann["bbox"]
                w = box[2]
                h = box[3]
                dw = 1. / (img_width)
                dh = 1. / (img_height)
                w = w*dw
                h = h*dh
                dataset.append([w,h])
    return np.array(dataset)




if __name__ == '__main__':
    # print(__file__)
    data = load_voc_dataset(ANNOTATIONS_PATH)
    out = kmeans(data, k=CLUSTERS)
    # clusters = [[10,13],[16,30],[33,23],[30,61],[62,45],[59,119],[116,90],[156,198],[373,326]]
    # out= np.array(clusters)/416.0
    print(out)
    print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100))
    print("Boxes:\n {}-{}".format(out[:, 0] * 416, out[:, 1] * 416))

    ratios = np.around(out[:, 0] / out[:, 1], decimals=2).tolist()
    print("Ratios:\n {}".format(sorted(ratios)))

在 get_result.py中放入自己训练集的.xml/.json的位置 然后执行即可。 图片做了归一化处理，最后默认的anchor值是相对yolov3配置文件中的416的。可根据自己需要修改。

Kmeans因为初始点敏感，所以每次运行得到的anchor值不一样，但是对应的avg iou稳定。用于训练的话就需要统计多组anchor，针对固定的测试集比较了。

最终结果如下，把他按照area大小写入yolov3的配置文件即可。