二.以聚类和搜图方式清洗图像数据集，解决图像冗余和不均衡问题 (图像特征提取+Kmeans聚类篇)（附案例代码）

文章内容结构：

一. 简单Kmeans聚类例子, 样本是针对向量数据

二. 提取图像特征，Kmeans对图像特征进行聚类

三. 解决图像数据样冗余和不均衡问题

（注： 这里全部是个人经验，能提升样本标注和清洗效率，不是标准的数据处理方式，希望对您有帮助。）

--------

一. 简单Kmeans聚类例子, 样本是针对向量数据

（1）X为样本特征，Y为样本簇类别，创建1000个样本，每个样本2个特征，对应x和y轴，共4个簇，簇中心在[-1,-1], [0,0],[1,1], [2,2]， 簇方差分别为[0.4, 0.2, 0.2]。

# -*- encoding:utf-8 -*-
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans
# X为样本特征，Y为样本簇类别，共1000个样本，每个样本2个特征，对应x和y轴，共4个簇，
# 簇中心在[-1,-1], [0,0],[1,1], [2,2]， 簇方差分别为[0.4, 0.2, 0.2]
X, y = make_blobs(n_samples=1000, n_features=2, centers=[[-1, -1], [0, 0], [1, 1], [2, 2]],
                  cluster_std=[0.4, 0.2, 0.2, 0.2])
model=KMeans(n_clusters=4)
model.fit(X)
y_pred = model.predict(X)
print(y_pred)

下面打印的结果： 每个相同的数字是聚类成了一个结果。

--------

二. 提取图像特征，Kmeans对图像特征进行聚类

（1）模型转onnx

将一个预训练的分类模型pt转为onnx（我的预训练模型是yolo分类，共190类别）

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    # 生成onnx模型在同目录下，名字为best_190labels_cls.onnx
    model = YOLO('/home/xxx/Downloads/best_190labels_cls.pt') 
    model.export(imgsz=224, 
                 format='onnx', 
                 simplify=True,  # 简化模型
                 dynamic=True,   # 若要多batch推理，置为True
                 )

（2）将onnx模型裁剪掉最后的Softmax层

import onnx
import os
from onnx import helper, shape_inference

def remove_final_softmax(input_path, output_path):
    model = onnx.load(input_path)
    graph = model.graph

    # 查找最后一个Softmax节点
    softmax_node = None
    for node in reversed(graph.node):
        if node.op_type == 'Softmax':
            softmax_node = node
            break
    if not softmax_node:
        raise ValueError("未找到Softmax节点")

    # 获取Gemm层的输出名称,也就是softmax_node 输入的信息
    print("softmax_node.input[0]-->", softmax_node.input[0])

    gemm_output_name = softmax_node.input[0]

    # 修改模型输出指向Gemm层
    for output in graph.output:
        if output.name == softmax_node.output[0]:
            print("softmax_node.output[0]-->",softmax_node.output[0])
            # 从value_info中获取Gemm输出的类型信息
            gemm_value_info = next((vi for vi in graph.value_info if vi.name == gemm_output_name), None)
            if gemm_value_info:
                output.CopyFrom(gemm_value_info)
            else:
                # 如果value_info中没有，则创建新的Tensor类型
                output.CopyFrom(helper.make_tensor_value_info(
                    gemm_output_name,
                    onnx.TensorProto.FLOAT,
                    None
                ))
            break

    # 删除Softmax节点
    graph.node.remove(softmax_node)

    # 进行形状推断以更新类型信息
    model = shape_inference.infer_shapes(model)

    # 验证并保存模型
    onnx.checker.check_model(model)
    onnx.save(model, output_path)
    print(f"新模型已保存至: {output_path}")


if __name__ =="__main__":
    root_path = "/home/xxx/Download"
    ori_onnx_path  = os.path.join(root_path, "best_190labels_cls.onnx")
    save_onnx_path = os.path.join(root_path, "best_190labels_cls_modified.onnx")

    # 移除softmax层
    remove_final_softmax(ori_onnx_path,save_onnx_path)

（3）onnx模型提取图像特征，Kmeans对图像特征进行聚类。

如果有显卡GPU，建议安装使用onnxruntime-gpu进行推理，速度快。若没有GPU，也可以用CPU进行推理。

安装  pip install onnxruntime-gpu

import os
import cv2
import numpy as np
from PIL import Image
import onnxruntime as ort
import shutil
from sklearn.cluster import KMeans
from sklearn.preprocessing import Normalizer
from  tqdm import tqdm
import math
import matplotlib.pyplot as plt


# 图像预处理函数
def preprocess_image(image_path):

    roi_frame= cv2.imread(image_path)
    width = roi_frame.shape[1]
    height = roi_frame.shape[0]

    if (width != CLASSIFY_SIZE) or (height != CLASSIFY_SIZE) :
                # 图像瓶子矫正，长边缩放到224，短边同比例缩放，再一起放到224*224的灰边图上
                if width > height:
                    # 将图像逆时针旋转90度
                    roi_frame = cv2.rotate(roi_frame, cv2.ROTATE_90_COUNTERCLOCKWISE)

                new_height = CLASSIFY_SIZE
                new_width = int(roi_frame.shape[1] * (CLASSIFY_SIZE / roi_frame.shape[0]))
                roi_frame = cv2.resize(roi_frame, (new_width, new_height))

                # 计算上下左右偏移量
                y_offset = (CLASSIFY_SIZE - roi_frame.shape[0]) // 2
                x_offset = (CLASSIFY_SIZE - roi_frame.shape[1]) // 2

                gray_image = np.full((CLASSIFY_SIZE, CLASSIFY_SIZE, 3), 128, dtype=np.uint8)
                # 将调整大小后的目标图像放置到灰度图上
                gray_image[y_offset:y_offset + roi_frame.shape[0], x_offset:x_offset + roi_frame.shape[1]] = roi_frame

                # # 显示结果
                # cv2.imshow("gray_image", gray_image)
                # cv2.waitKey(1)

                # 将图像转为 rgb
                gray_image = cv2.cvtColor(gray_image, cv2.COLOR_BGR2RGB)

    else:
        gray_image = cv2.cvtColor(roi_frame, cv2.COLOR_BGR2RGB)

    img_np = np.array(gray_image).transpose(2, 0, 1).astype(np.float32)

    # 假设模型需要[0,1]归一化
    img_np = img_np / 255.0

    # 均值，方差
    mean = np.array([0.485, 0.456, 0.406],dtype=np.float32).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225],dtype=np.float32).reshape(3, 1, 1)
    img_np= (img_np - mean)/std

    return np.expand_dims(img_np, axis=0)




# 安装  pip install onnxruntime-gpu
def get_onnx_providers():
    # 检查是否安装了GPU版本的ONNX Runtime
    all_provider = ort.get_available_providers()

    if "CUDAExecutionProvider" in all_provider:
        providers = [
            ("CUDAExecutionProvider", {
                "device_id": 0,
                "arena_extend_strategy": "kNextPowerOfTwo",
                "gpu_mem_limit": 6 * 1024 * 1024 * 1024,  # 限制GPU内存使用为2GB
                "cudnn_conv_algo_search": "EXHAUSTIVE",
                "do_copy_in_default_stream": True,
            }),
            "CPUExecutionProvider"
        ]

        print("检测到NVIDIA GPU，使用CUDA加速")
        return providers
    else:
        print("未检测到NVIDIA GPU，使用CPU")
        return ["CPUExecutionProvider"]



if __name__ =="__main__":

    # ONNX模型路径
    MODEL_PATH = "/home/xxx/Downloads/best_190labels_cls_modified.onnx"
    # 原图像文件夹路径
    IMAGE_DIR = "08以图搜图_找相似度/99_test_datasets/xxx/8"
    # 分类结果输出路径
    OUTPUT_DIR = "08以图搜图_找相似度/99_test_datasets/xxx/8_kmeans_besk_k_classify"
    # 推断图像尺寸
    CLASSIFY_SIZE = 224
    # 手动划分分类数量
    NUM_CLUSTERS = 3000

    # 创建输出文件夹
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("ONNX Runtime版本:", ort.__version__)
    print("可用执行器:", ort.get_available_providers())

    # 可用执行器: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'AzureExecutionProvider', 'CPUExecutionProvider']

    # 加载ONNX模型（动态获取输入/输出名称）
    ort_session = ort.InferenceSession(
        MODEL_PATH,
        providers=get_onnx_providers()
    )


    # 确保输出名称正确
    input_name = ort_session.get_inputs()[0].name
    output_name = ort_session.get_outputs()[0].name


    # 提取特征向量
    features = []
    image_paths = []

    print("====开始对所有图像推理, 提取特征====")
    for filename in tqdm(os.listdir(IMAGE_DIR)):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            path = os.path.join(IMAGE_DIR, filename)
            try:
                # 前处理
                input_tensor = preprocess_image(path)
                # 推断
                feature = ort_session.run([output_name], {input_name: input_tensor})[0]
                # 确保特征展平为1D
                features.append(feature.reshape(-1))
                image_paths.append(path)
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}")
    print("+++++提取特征结束+++++")


    # L2归一化（使欧氏距离等价于余弦相似度）
    normalizer = Normalizer(norm="l2")
    normalized_features = normalizer.fit_transform(features)

    
    print("==== 开始: 正式聚类  ====")                                                                                                                
    # K-means聚类
    kmeans = KMeans(n_clusters=NUM_CLUSTERS, random_state=42)
    clusters = kmeans.fit_predict(normalized_features)
    print("+++++ 结束: 正式聚类 +++++")


    print("==== 开始: 移动分类小图  ====")                                                                                                                
    # 保存结果
    for cluster_id in range(best_num_clusters):   # NUM_CLUSTERS
        cluster_dir = os.path.join(OUTPUT_DIR, f"cluster_{cluster_id}")
        os.makedirs(cluster_dir, exist_ok=True)

    for path, cluster_id in zip(image_paths, clusters):
        dest_dir = os.path.join(OUTPUT_DIR, f"cluster_{cluster_id}")
        shutil.copy(path, dest_dir)

    print(f"++++ 分类完成！结果保存在 {OUTPUT_DIR} +++++")

以下是聚类结果：

--------

三. 解决图像数据样本冗余和不均衡问题

分类小图数据：

拿到聚类结果以后，依次遍历每个文件夹里面取1张或2张小图，这样能解决分类小图冗余、类别不均衡问题。

检测大图数据：

拿到聚类结果以后，依次遍历每个文件夹里面取1张小图，再根据取到的小图名找到对应的大图，优先根据大图名找其他文件夹中同名的其他小图，这样能缓解检测大图冗余、类别不均衡问题。