yolo voc 锚框聚类-2

0，yolo v3的锚框聚类思想

1，随机选择锚框数量，原论文数量为9
2，每个框与锚框计算IOU，最近为该类，由于iou很小，所以按照 1-iou
3，每次最近锚框计算后，修改锚框
重复1-3，直到聚类完成，并根据avg_iou判断聚类标准

方法有点，聚类结果稳定

1，聚类代码

# *_*coding:utf-8 *_*
import os
import numpy as np
import xml.etree.ElementTree as ET

# np.random.seed(0)  # 设置随机种子

def iou(box, 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):
    return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])


def translate_boxes(boxes):
    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):

    rows = boxes.shape[0]
    distances = np.empty((rows, k))
    last_clusters = np.zeros((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

def load_dataset(path):
    dataset = []
    for file in os.listdir(path):
        xml_file = os.path.join(path, file)

        root = ET.parse(xml_file)

        height = int(root.findtext("size/height"))
        width = int(root.findtext("size/width"))

        try:
            for iter in root.iter("object"):
                xmin = int(iter.findtext("bndbox/xmin"))
                ymin = int(iter.findtext("bndbox/ymin"))
                xmax = int(iter.findtext("bndbox/xmax"))
                ymax = int(iter.findtext("bndbox/ymax"))
                width_trans = (xmax - xmin) / width
                height_trans = (ymax - ymin) / height
                dataset.append([width_trans, height_trans])
        except:
            pass

    return np.array(dataset)


annotations = "./annotations"
data = load_dataset(annotations) # 拿到聚类框，width, height归一化的值
out = kmeans(data, k=5)
print("Accuracy: {:.2f}%".format(avg_iou(data, out) * 100))
print("Boxes:\n {}".format(out))

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

输出结果：

Accuracy: 61.24%
Boxes:
 [[0.1        0.17066667]
 [0.812      0.82933333]
 [0.194      0.37866667]
 [0.402      0.608     ]
 [0.042      0.07207207]]
Ratios:
[0.51, 0.58, 0.59, 0.66, 0.98]

2，相关链接
https://github.com/lars76/kmeans-anchor-boxes