四. 以Annoy算法建树的方式聚类清洗图像数据集，一次建树，无限次聚类搜索，提升聚类搜索效率。(附完整代码)

BB_CC_DD2025-04-26 11:14

文章内容结构：

一. 先介绍什么是Annoy算法。
二. 用Annoy算法建树的完整代码。
三. 用Annoy建树后的树特征匹配聚类归类图像。

一. 先介绍什么是Annoy算法

下面的文章链接将Annoy算法讲解的很详细，这里就不再做过多原理的分析了，想详细了解的可以看看这篇文章内容。

https://zhuanlan.zhihu.com/p/148819536

总的来说：

（1）通过多次递归迭代，建立一个二叉树，以二叉树的方式，提升数据聚类和搜索速度，但会损失一些精度。

（2）建树过程相对比较耗时，但建树只需要一次，部署到线上或者其他设备上，能无数次聚类搜索。(类似于人脸识别的人脸底库)

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

二. 用Annoy算法建树的完整代码

对底库聚类建树，生成Annoy树特征文件。

下面参数说明：

python 复制代码 
最佳聚类类别数量， 是根据《三.以聚类的方式清洗图像数据集，找到最佳聚类类别数 (图像特征提取+Kmeans聚类)》获取得到
BEST_NUM_CLUSTERS = 2501


图像特征提取后的向量维度，是pt或者onnx模型输出的类别数
FEATURE_DIM = 190


推断图像尺寸，是根据训练pt模型时，输入的图像尺寸大小
CLASSIFY_SIZE = 224

以下是正式的代码：

python 复制代码 
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) :

                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)



# 卸载 onnxruntime
# 安装  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__":

    root_path =  "/home/xxx/Download"
    # ONNX模型路径
    MODEL_PATH = os.path.join(root_path, "08以图搜图_找相似度/98_weights/classify_modified_model_224.onnx")
    # 图像文件夹路径
    IMAGE_DIR = os.path.join(root_path, "08以图搜图_找相似度/99_test_datasets/8_bcd已验收/8")
    # 分类结果输出路径
    OUTPUT_DIR = os.path.join(root_path, "08以图搜图_找相似度/99_test_datasets/8_bcd已验收/8_kmeans_besk_k_classify")
    # 保存ann建树文件路径
    ANNOY_PATH = "08以图搜图_找相似度/01kmeans和DBscan/kmeans/annoy_cls.ann"
    # 最佳聚类类别数量(用kmeans和inner找到的)
    BEST_NUM_CLUSTERS = 2501
    # 图像特征提取后的向量维度
    FEATURE_DIM = 190  # 根据自己的模型输出维度修改
    # 推断图像尺寸
    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

    from annoy import AnnoyIndex
    t = AnnoyIndex(FEATURE_DIM, metric="angular")  # FEATURE_DIM是图像特征提取后的向量维度

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

    print("====开始对所有图像推理, 提取特征====")
    for index, filename in tqdm(enumerate(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,  190维度
                features.append(feature.reshape(-1))
                image_paths.append(path)

                # 增加到Annoy树
                t.add_item(index, feature.reshape(-1))
               
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}")

    t.build(BEST_NUM_CLUSTERS)    # 根据kmeans聚类找到最佳的聚类类别数量
    t.save(ANNOY_PATH)
    print("+++++提取特征结束+++++")
    print("+++++Annoy建树结束+++++++++")

生成建树annoy_cls.ann文件。

三. 用Annoy建树后的树特征匹配聚类归类图像

使用流程：

（1）加载ann建树文件

（2）提取单张A图像特征

（3）单张A图像特征与ann建树文件的特征进行比对，找到ann建树文件里面的与A图像特征相似的TOP_K的底库图像，拷贝走或者移动走。

python 复制代码 
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) :

                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)





# todo
# 卸载 onnxruntime
# 安装  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__":

    root_path =  "/home/xxx/Download"
    # ONNX模型路径
    MODEL_PATH = os.path.join(root_path, "08以图搜图_找相似度/98_weights/classify_modified_model_224.onnx")
    # 图像文件夹路径
    IMAGE_DIR = os.path.join(root_path, "08以图搜图_找相似度/99_test_datasets/8_bcd已验收/8")
    # 分类结果输出路径
    OUTPUT_DIR = os.path.join(root_path, "08以图搜图_找相似度/99_test_datasets/8_bcd已验收/8_kmeans_besk_k_classify")
    # 保存annoy建树路径
    ANNOY_PATH = os.path.join(root_path, "08以图搜图_找相似度/01kmeans和DBscan/kmeans/annoy_cls.ann")
    # 最佳聚类类别数量
    BEST_NUM_CLUSTERS = 2501
    # 图像特征提取后的向量维度
    FEATURE_DIM = 190
    # 推断图像尺寸
    CLASSIFY_SIZE = 224
    # 取top10
    TOP_K = 10
    # 手动划分分类数量
    # 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

    from annoy import AnnoyIndex
    Annoy_ = AnnoyIndex(FEATURE_DIM, metric="angular")  # FEATURE_DIM是图像特征提取后的向量维度
    Annoy_.load(ANNOY_PATH) 

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

    # 获取所有图像路径
    for _, filename in tqdm(enumerate(os.listdir(IMAGE_DIR))):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            path = os.path.join(IMAGE_DIR, filename)
            image_paths.append(path)


    print("====开始对所有图像推理, 提取特征, 根据创建的树进行聚类====")
    for _, filename in tqdm(enumerate(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,  190维度
                features.append(feature.reshape(-1))
                # image_paths.append(path)
                # 取top10的相似图像
                similar_img_indices, similar_img_distances=Annoy_.get_nns_by_vector(feature.reshape(-1), TOP_K, include_distances=True)
                print("similar_img_index:", similar_img_indices)
                print("similar_img_distance:", similar_img_distances)

                shutil.copy(path, os.path.join(OUTPUT_DIR,"11"))
                #  移动相似图像到输出目录
                for idx in similar_img_indices:
                    similar_image_path = image_paths[idx]
                    # shutil.move(similar_image_path, OUTPUT_DIR)
                    shutil.copy(similar_image_path, OUTPUT_DIR)

               
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}")

    print("+++++提取特征结束+++++")
    print("+++++根据Annoy数特征聚类归类图像结束+++++++++")
相关推荐
林泽毅9 分钟前
PyTRIO快速入门实战篇（二）：用 GRPO 提升 GSM8K 数学推理准确率
人工智能·深度学习·机器学习·llm·强化学习
AKA__Zas13 分钟前
芝士算法（位运算）
java·开发语言·算法·leetcode·学习方法
wabs66631 分钟前
关于图论【Kruskal算法|卡码网53.寻宝的思考】
数据结构·算法·图论·卡码网
hans汉斯32 分钟前
人工智能与机器人研究|面向无标签数据的三维场景语义理解方法研究
人工智能·神经网络·算法·信息可视化·cnn·机器人
小程故事多_8011 小时前
从A2C、TRPO、PPO到GRPO，强化学习策略梯度算法完整演进与大模型落地实战解析
人工智能·算法
AI人工智能+12 小时前
一种基于深度学习技术的高精度医疗机构执业许可证识别系统，构建了一套基于深度神经网络的端到端智能识别系统，为医疗行业提
深度学习·ocr·医疗机构执业许可证识别
硅谷秋水12 小时前
EgoSteer：一种基于第一人称视角视频、实现可控灵巧操作的全栈系统
深度学习·机器学习·语言模型·机器人·音视频
2601_9561219712 小时前
二分算法(知识点+题目）
c++·算法
AI街潜水的八角13 小时前
基于YOLO26交通标志检测系统1：交通标志检测数据集说明(含下载链接)
深度学习·神经网络
aiblog14 小时前
深度学习中“Transformer”怎么翻译为中文？
人工智能·深度学习·transformer
热门推荐
01如何新建文件夹？ 电脑新建文件夹的4种方法02GitHub 镜像站点03国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）04AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路05Agnes AI 免费 API 接入指南：文本、生图、生视频，一套接口全免费06微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新072026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？082026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片09历年考研数学一、数学二、数学三真题试卷及答案PDF1000 Debian字符界面如何支持中文