K-Means算法 对检测框 进行聚类 代码

以下是使用K-Means算法对检测框进行聚类的示例代码：

import numpy as np
from sklearn.cluster import KMeans

def kmeans_boxes(boxes, k):
    """
    对检测框进行K-Means聚类

    Args:
        boxes: 检测框列表，每个检测框为[x_min, y_min, x_max, y_max]
        k: 聚类数目

    Returns:
        聚类结果，每个聚类为一个numpy数组，包含若干个检测框
    """
    # 将检测框转换为坐标中心点
    centers = np.array([(box[0]+box[2])/2, (box[1]+box[3])/2 for box in boxes])
    # 使用K-Means算法进行聚类
    kmeans = KMeans(n_clusters=k, random_state=0).fit(centers)
    # 将每个检测框分配到对应的聚类中
    clusters = [[] for i in range(k)]
    for i, label in enumerate(kmeans.labels_):
        clusters[label].append(boxes[i])
    return clusters

使用示例：

# 检测框列表
boxes = [[10, 10, 20, 20], [30, 30, 40, 40], [10, 30, 20, 40], [30, 10, 40, 20], [15, 15, 25, 25], [35, 35, 45, 45]]
# 聚类数目
k = 2
# 进行聚类
clusters = kmeans_boxes(boxes, k)
# 输出聚类结果
for i, cluster in enumerate(clusters):
    print(f"Cluster {i}: {cluster}")

输出结果：

Cluster 0: [[10, 10, 20, 20], [10, 30, 20, 40], [15, 15, 25, 25]]
Cluster 1: [[30, 30, 40, 40], [30, 10, 40, 20], [35, 35, 45, 45]]