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

文章内容结构：

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

一. 先介绍什么是Annoy算法

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

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

总的来说：

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

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

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

--------

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

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

下面参数说明：

最佳聚类类别数量， 是根据《三.以聚类的方式清洗图像数据集，找到最佳聚类类别数 (图像特征提取+Kmeans聚类)》获取得到
BEST_NUM_CLUSTERS = 2501图像特征提取后的向量维度，是pt或者onnx模型输出的类别数
FEATURE_DIM = 190推断图像尺寸，是根据训练pt模型时，输入的图像尺寸大小
CLASSIFY_SIZE = 224

以下是正式的代码：

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_SIZEnew_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]) // 2x_offset = (CLASSIFY_SIZE - roi_frame.shape[1]) // 2gray_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)# 将图像转为 rgbgray_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)/stdreturn np.expand_dims(img_np, axis=0)# 卸载 onnxruntime
# 安装  pip install onnxruntime-gpu
def get_onnx_providers():# 检查是否安装了GPU版本的ONNX Runtimeall_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 providerselse: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].nameoutput_name = ort_session.get_outputs()[0].namefrom annoy import AnnoyIndext = 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的底库图像，拷贝走或者移动走。

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_SIZEnew_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]) // 2x_offset = (CLASSIFY_SIZE - roi_frame.shape[1]) // 2gray_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)# 将图像转为 rgbgray_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)/stdreturn np.expand_dims(img_np, axis=0)# todo
# 卸载 onnxruntime
# 安装  pip install onnxruntime-gpu
def get_onnx_providers():# 检查是否安装了GPU版本的ONNX Runtimeall_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 providerselse: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# 取top10TOP_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].nameoutput_name = ort_session.get_outputs()[0].namefrom annoy import AnnoyIndexAnnoy_ = 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数特征聚类归类图像结束+++++++++")