当前位置：首页 > news >正文

K-Means聚类算法

news 来源：原创 2025/4/22 9:57:21

文章目录

前言
一、为何要使用先验框
二、K-Means聚类算法核心思想
三、代码实现
总结

前言

本文用于学习记录yolov2当中所使用到的K-Means聚类算法提取先验框。

提示：以下是本篇文章正文内容，下面案例可供参考

一、为何要使用先验框

先验框是在训练神经网络之前定义的框，用于知道yolo网络学习如何预测目标的位置和类别。yolov2中的先验框是基于训练数据集中的目标边界框而来，代表了不同目标的不同尺寸和比例。采用先验框可以使得模型快速收敛，加快训练所需时间。

二、K-Means聚类算法核心思想

初始化聚类中心，可以随机选取边界框作为初始化先验框。
计算每个数据中的边界框到每个聚类先验框的距离1-IOU(box1,box2)。
将每个边界框分配到距离最近的聚类中心。
更新聚类中心（使用分配给它的所有边界框的平均值或者中位数等）。
重复以上步骤，直到聚类中心不在发生显著变化或者达到预定的迭代次数。

三、代码实现

import numpy as npdef iou(box, clusters):"""计算单个边界框与所有聚类中心的 IoU。:param box: 单个边界框的宽高 [w, h]:param clusters: K 个聚类中心的宽高数组，形状 (K, 2):return: IoU 值数组，形状 (K,)"""x = np.minimum(clusters[:, 0], box[0])y = np.minimum(clusters[:, 1], box[1])intersection = x * ybox_area = box[0] * box[1]cluster_area = clusters[:, 0] * clusters[:, 1]return intersection / (box_area + cluster_area - intersection)def kmeans(boxes, k, max_iter=100):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():breaklast_clusters = nearest_clusters# 更新簇中心（取中位数）for cluster in range(k):clusters[cluster] = np.median(boxes[nearest_clusters == cluster], axis=0)return clustersdef convert_anchors(clusters, input_size=416, downsample=32):"""将聚类结果转换为绝对尺寸。:param clusters: 聚类中心数组，形状 (K, 2):param input_size: 输入图像尺寸（如 416）:param downsample: 下采样倍率（如 32）"""return clusters * (input_size / downsample)def avg_iou(boxes, clusters):return np.mean([np.max(iou(boxes[i], clusters)) for i in range(boxes.shape[0])])boxes = np.array([[100,80],[120,110],[150,130],[200,180],[300,250],[100,60],[60,80],[80,90],[90,100],[100,500]]
)clusters = kmeans(boxes, 2)
print(avg_iou(boxes, clusters))