四. 以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_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的底库图像，拷贝走或者移动走。
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数特征聚类归类图像结束+++++++++")
