四. 以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數特征聚類歸類圖像結束+++++++++")