712 lines
27 KiB
Python
712 lines
27 KiB
Python
"""
|
||
食物分类器 - FAISS索引构建器
|
||
将训练集图片转换为特征向量并构建FAISS索引用于快速相似度检索
|
||
|
||
功能：
|
||
1. 加载训练好的ResNet50 Embedding模型
|
||
2. 遍历训练集所有图片，提取特征向量
|
||
3. 构建FAISS索引文件
|
||
4. 创建标签映射文件
|
||
5. 实现快速相似度检索功能
|
||
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import pickle
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
from tqdm import tqdm
|
||
from datetime import datetime
|
||
from typing import List, Tuple, Dict, Optional
|
||
import faiss
|
||
|
||
# 导入项目模块
|
||
from net.resnet_embedding import create_resnet50_embedding
|
||
from settings import settings
|
||
|
||
|
||
class FAISSIndexBuilder:
|
||
"""FAISS索引构建器"""
|
||
|
||
def __init__(self, model_path: str, embedding_dim: int = 512):
|
||
"""
|
||
初始化FAISS索引构建器
|
||
|
||
Args:
|
||
model_path: 训练好的模型路径
|
||
embedding_dim: 特征向量维度
|
||
"""
|
||
self.model_path = model_path
|
||
self.embedding_dim = embedding_dim
|
||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
# 初始化模型
|
||
self.model = None
|
||
self.load_model()
|
||
|
||
# 数据存储
|
||
self.embeddings = []
|
||
self.image_paths = []
|
||
self.labels = []
|
||
self.class_names = []
|
||
self.class_to_idx = {}
|
||
self.idx_to_class = {}
|
||
|
||
print(f"使用设备: {self.device}")
|
||
print(f"特征向量维度: {self.embedding_dim}")
|
||
|
||
def load_model(self):
|
||
"""加载训练好的ResNet50 Embedding模型"""
|
||
try:
|
||
print(f"正在加载模型: {self.model_path}")
|
||
|
||
# 创建模型
|
||
self.model = create_resnet50_embedding(
|
||
embedding_dim=self.embedding_dim,
|
||
pretrained=True,
|
||
use_internal_preprocess=False
|
||
)
|
||
|
||
# 加载权重
|
||
checkpoint = torch.load(self.model_path, map_location=self.device)
|
||
|
||
# 处理不同的保存格式
|
||
if isinstance(checkpoint, dict):
|
||
if 'model_state_dict' in checkpoint:
|
||
# Triplet格式：使用model_state_dict
|
||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
||
print("检测到Triplet模型格式，使用'model_state_dict'加载")
|
||
elif 'backbone_state_dict' in checkpoint:
|
||
# CosFace格式：使用backbone_state_dict（只加载backbone部分）
|
||
self.model.load_state_dict(checkpoint['backbone_state_dict'])
|
||
print("检测到CosFace模型格式，使用'backbone_state_dict'加载")
|
||
elif 'state_dict' in checkpoint:
|
||
# 通用格式：使用state_dict
|
||
self.model.load_state_dict(checkpoint['state_dict'])
|
||
print("检测到通用模型格式，使用'state_dict'加载")
|
||
else:
|
||
# 直接加载整个checkpoint作为state_dict
|
||
self.model.load_state_dict(checkpoint)
|
||
print("使用checkpoint直接加载")
|
||
else:
|
||
# checkpoint本身就是state_dict
|
||
self.model.load_state_dict(checkpoint)
|
||
print("checkpoint为state_dict格式，直接加载")
|
||
|
||
# 移动到设备并设置为评估模式
|
||
self.model.to(self.device)
|
||
self.model.eval()
|
||
|
||
print("模型加载成功!")
|
||
|
||
except Exception as e:
|
||
print(f"模型加载失败: {e}")
|
||
raise
|
||
|
||
def scan_training_data(self, train_dir: str) -> Tuple[List[str], List[str], List[int]]:
|
||
"""
|
||
扫描训练数据目录，获取所有图片路径和标签
|
||
|
||
Args:
|
||
train_dir: 训练数据目录路径
|
||
|
||
Returns:
|
||
Tuple[List[str], List[str], List[int]]: (图片路径列表, 类别名称列表, 标签索引列表)
|
||
"""
|
||
print(f"正在扫描训练数据目录: {train_dir}")
|
||
|
||
image_paths = []
|
||
class_names = []
|
||
labels = []
|
||
|
||
# 获取所有类别目录
|
||
class_dirs = [d for d in os.listdir(train_dir)
|
||
if os.path.isdir(os.path.join(train_dir, d))]
|
||
class_dirs.sort() # 确保顺序一致
|
||
|
||
print(f"发现类别: {class_dirs}")
|
||
|
||
# 创建类别映射
|
||
self.class_names = class_dirs
|
||
self.class_to_idx = {class_name: idx for idx, class_name in enumerate(class_dirs)}
|
||
self.idx_to_class = {idx: class_name for idx, class_name in enumerate(class_dirs)}
|
||
|
||
# 扫描每个类别目录
|
||
for class_name in class_dirs:
|
||
class_dir = os.path.join(train_dir, class_name)
|
||
class_idx = self.class_to_idx[class_name]
|
||
|
||
# 获取该类别下的所有图片
|
||
image_files = [f for f in os.listdir(class_dir)
|
||
if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp'))]
|
||
|
||
print(f"类别 '{class_name}': {len(image_files)} 张图片")
|
||
|
||
for image_file in image_files:
|
||
# 可以不添加那些增强的图片
|
||
if image_file.lower().startswith("img"):
|
||
image_path = os.path.join(class_dir, image_file)
|
||
image_paths.append(image_path)
|
||
class_names.append(class_name)
|
||
labels.append(class_idx)
|
||
|
||
print(f"总计扫描到 {len(image_paths)} 张图片")
|
||
return image_paths, class_names, labels
|
||
|
||
def extract_features_batch(self, image_paths: List[str], batch_size: int = 32) -> np.ndarray:
|
||
"""
|
||
批量提取图片特征向量
|
||
|
||
Args:
|
||
image_paths: 图片路径列表
|
||
batch_size: 批处理大小
|
||
|
||
Returns:
|
||
np.ndarray: 特征向量数组，形状为 (N, embedding_dim)
|
||
"""
|
||
print(f"开始提取特征向量，批大小: {batch_size}")
|
||
|
||
all_embeddings = []
|
||
|
||
with torch.no_grad():
|
||
for i in tqdm(range(0, len(image_paths), batch_size), desc="提取特征"):
|
||
batch_paths = image_paths[i:i + batch_size]
|
||
batch_images = []
|
||
valid_indices = []
|
||
|
||
# 加载并预处理批次图片
|
||
for j, img_path in enumerate(batch_paths):
|
||
try:
|
||
# 加载图片
|
||
image = Image.open(img_path).convert('RGB')
|
||
|
||
# 预处理
|
||
tensor = self.model.preprocess_image(image)
|
||
batch_images.append(tensor.squeeze(0)) # 移除batch维度
|
||
valid_indices.append(i + j)
|
||
|
||
except Exception as e:
|
||
print(f"警告: 无法处理图片 {img_path}: {e}")
|
||
continue
|
||
|
||
if not batch_images:
|
||
continue
|
||
|
||
# 堆叠成批次并移动到设备
|
||
batch_tensor = torch.stack(batch_images).to(self.device)
|
||
|
||
# 提取特征
|
||
batch_embeddings = self.model(batch_tensor, normalize=True)
|
||
|
||
# 转换为numpy并添加到结果中
|
||
all_embeddings.append(batch_embeddings.cpu().numpy())
|
||
|
||
# 合并所有批次的结果
|
||
if all_embeddings:
|
||
embeddings = np.vstack(all_embeddings)
|
||
print(f"成功提取 {embeddings.shape[0]} 个特征向量")
|
||
return embeddings
|
||
else:
|
||
raise ValueError("没有成功提取任何特征向量")
|
||
|
||
def build_faiss_index(self, embeddings: np.ndarray, index_type: str = 'flat') -> faiss.Index:
|
||
"""
|
||
构建FAISS索引
|
||
|
||
Args:
|
||
embeddings: 特征向量数组
|
||
index_type: 索引类型 ('flat', 'ivf', 'hnsw')
|
||
|
||
Returns:
|
||
faiss.Index: 构建好的FAISS索引
|
||
"""
|
||
print(f"构建FAISS索引，类型: {index_type}")
|
||
|
||
dimension = embeddings.shape[1]
|
||
n_vectors = embeddings.shape[0]
|
||
|
||
if index_type == 'flat':
|
||
# 精确搜索索引
|
||
index = faiss.IndexFlatIP(dimension) # 内积索引（适合归一化向量）
|
||
|
||
elif index_type == 'ivf':
|
||
# IVF索引（适合大规模数据）
|
||
nlist = min(100, max(1, n_vectors // 10)) # 聚类中心数量
|
||
quantizer = faiss.IndexFlatIP(dimension)
|
||
index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
|
||
|
||
# 训练索引
|
||
print("训练IVF索引...")
|
||
index.train(embeddings.astype(np.float32))
|
||
|
||
elif index_type == 'hnsw':
|
||
# HNSW索引（快速近似搜索）
|
||
M = 16 # 连接数
|
||
index = faiss.IndexHNSWFlat(dimension, M)
|
||
index.hnsw.efConstruction = 200
|
||
|
||
else:
|
||
raise ValueError(f"不支持的索引类型: {index_type}")
|
||
|
||
# 添加向量到索引
|
||
print("添加向量到索引...")
|
||
index.add(embeddings.astype(np.float32))
|
||
|
||
print(f"索引构建完成，包含 {index.ntotal} 个向量")
|
||
return index
|
||
|
||
def save_index_and_metadata(self, index: faiss.Index, output_dir: str):
|
||
"""
|
||
保存FAISS索引和相关元数据
|
||
|
||
Args:
|
||
index: FAISS索引
|
||
output_dir: 输出目录
|
||
"""
|
||
# 创建输出目录
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 保存FAISS索引
|
||
index_path = os.path.join(output_dir, 'faiss_index.bin')
|
||
faiss.write_index(index, index_path)
|
||
print(f"FAISS索引已保存到: {index_path}")
|
||
|
||
# 保存图片路径映射
|
||
paths_path = os.path.join(output_dir, 'image_paths.pkl')
|
||
with open(paths_path, 'wb') as f:
|
||
pickle.dump(self.image_paths, f)
|
||
print(f"图片路径映射已保存到: {paths_path}")
|
||
|
||
# 保存标签映射
|
||
labels_path = os.path.join(output_dir, 'labels.pkl')
|
||
labels_json_path = os.path.join(output_dir, 'labels.json')
|
||
with open(labels_path, 'wb') as f:
|
||
pickle.dump(self.labels, f)
|
||
with open(labels_json_path, "w", encoding="utf-8") as f:
|
||
json.dump(self.labels, f, ensure_ascii=False)
|
||
print(f"标签映射已保存到: {labels_path}")
|
||
|
||
# 保存类别信息
|
||
class_info = {
|
||
'class_names': self.class_names,
|
||
'class_to_idx': self.class_to_idx,
|
||
'idx_to_class': self.idx_to_class,
|
||
'embedding_dim': self.embedding_dim,
|
||
'total_images': len(self.image_paths),
|
||
'build_time': datetime.now().isoformat()
|
||
}
|
||
|
||
class_info_path = os.path.join(output_dir, 'class_info.json')
|
||
with open(class_info_path, 'w', encoding='utf-8') as f:
|
||
json.dump(class_info, f, ensure_ascii=False, indent=2)
|
||
print(f"类别信息已保存到: {class_info_path}")
|
||
|
||
# 保存特征向量（可选，用于调试）
|
||
embeddings_path = os.path.join(output_dir, 'embeddings.npy')
|
||
embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
|
||
embeddings_array = np.array(self.embeddings)
|
||
np.save(embeddings_path, embeddings_array)
|
||
|
||
# 转成 Python 列表（JSON 不支持 numpy.array 直接序列化）
|
||
embeddings_list = embeddings_array.tolist()
|
||
# 保存成 JSON 文件
|
||
with open(embeddings_json_path, "w", encoding="utf-8") as f:
|
||
json.dump(embeddings_list, f, ensure_ascii=False)
|
||
|
||
print(f"特征向量已保存到: {embeddings_path}")
|
||
|
||
def select_representative_samples(
|
||
self,
|
||
embeddings: np.ndarray,
|
||
image_paths: List[str],
|
||
labels: List[int],
|
||
max_per_class: int = 30) -> Tuple[np.ndarray, List[str], List[int]]:
|
||
"""
|
||
使用 K-Means 为每个类别筛选代表性样本，消除样本数量不均衡带来的密度偏差。
|
||
|
||
原理：
|
||
- 样本数 <= max_per_class：全部保留
|
||
- 样本数 > max_per_class：对该类向量做 K-Means(k=max_per_class)
|
||
然后为每个聚类中心找最近的真实图片，只保留这些代表性样本
|
||
|
||
Args:
|
||
embeddings : 全量嵌入向量数组 (N, dim)
|
||
image_paths : 图片路径列表
|
||
labels : 标签索引列表
|
||
max_per_class: 每类最多保留的样本数（默认 30）
|
||
|
||
Returns:
|
||
Tuple[np.ndarray, List[str], List[int]]: 筛选后的 (embeddings, paths, labels)
|
||
"""
|
||
print(f"\n{'='*60}")
|
||
print(f"K-Means 代表性样本筛选 (max_per_class={max_per_class})")
|
||
print(f"{'='*60}")
|
||
|
||
embeddings_array = embeddings if isinstance(embeddings, np.ndarray) else np.array(embeddings)
|
||
labels_array = np.array(labels)
|
||
selected_global_indices: List[int] = []
|
||
|
||
for class_idx in np.unique(labels_array):
|
||
class_name = self.idx_to_class[int(class_idx)]
|
||
global_indices = np.where(labels_array == class_idx)[0]
|
||
count = len(global_indices)
|
||
|
||
if count <= max_per_class:
|
||
# 样本数不超过上限，全部保留
|
||
selected_global_indices.extend(global_indices.tolist())
|
||
print(f" [{class_name}]: {count} 张, 全部保留")
|
||
else:
|
||
# 样本数超过上限 -> K-Means 聚类后选代表
|
||
class_embs = embeddings_array[global_indices].astype(np.float32)
|
||
dim = class_embs.shape[1]
|
||
k = max_per_class
|
||
|
||
# 1. 用 FAISS 内置 K-Means 聚类（无需 sklearn）
|
||
kmeans = faiss.Kmeans(dim, k, niter=30, verbose=False)
|
||
kmeans.train(class_embs)
|
||
centroids = kmeans.centroids # shape: (k, dim)
|
||
|
||
# 2. 为每个幽灵中心，找离它最近的真实图片（幽灵变真身）
|
||
temp_index = faiss.IndexFlatL2(dim)
|
||
temp_index.add(class_embs)
|
||
_, nearest_local = temp_index.search(centroids, 1) # (k, 1)
|
||
nearest_local = nearest_local[:, 0].tolist()
|
||
|
||
# 3. 去重（极少情况下两个中心可能最近邻相同）
|
||
unique_local = list(set(nearest_local))
|
||
|
||
# 4. 映射回全局索引
|
||
selected_global = global_indices[unique_local].tolist()
|
||
selected_global_indices.extend(selected_global)
|
||
print(f" [{class_name}]: {count} 张 -> K-Means 筛选为 {len(selected_global)} 张")
|
||
|
||
# 按索引排序，确保顺序稳定
|
||
selected_global_indices = sorted(selected_global_indices)
|
||
|
||
filtered_embeddings = embeddings_array[selected_global_indices]
|
||
filtered_paths = [image_paths[i] for i in selected_global_indices]
|
||
filtered_labels = [labels[i] for i in selected_global_indices]
|
||
|
||
total_before = len(image_paths)
|
||
total_after = len(filtered_paths)
|
||
print(f"\nK-Means 筛选完成: {total_before} -> {total_after} 个向量"
|
||
f"（减少了 {total_before - total_after} 个）")
|
||
print(f"{'='*60}\n")
|
||
|
||
return filtered_embeddings, filtered_paths, filtered_labels
|
||
|
||
def build_complete_index(self, train_dir: str, output_dir: str,
|
||
batch_size: int = 32, index_type: str = 'flat',
|
||
max_per_class: int = 30):
|
||
"""
|
||
完整的索引构建流程
|
||
|
||
Args:
|
||
train_dir: 训练数据目录
|
||
output_dir: 输出目录
|
||
batch_size: 批处理大小
|
||
index_type: 索引类型
|
||
"""
|
||
print("=" * 60)
|
||
print("\u5f00始构建 FAISS 索引")
|
||
print("=" * 60)
|
||
|
||
# 1. 扫描训练数据（全量，不做任何过滤）
|
||
image_paths, class_names, labels = self.scan_training_data(train_dir)
|
||
|
||
# 2. 提取全量特征向量（K-Means 聚类需要看到所有向量）
|
||
embeddings = self.extract_features_batch(image_paths, batch_size)
|
||
|
||
# 3. K-Means 代表性样本筛选（每类最多保留 max_per_class 个向量）
|
||
embeddings, image_paths, labels = self.select_representative_samples(
|
||
embeddings, image_paths, labels, max_per_class
|
||
)
|
||
self.image_paths = image_paths
|
||
self.labels = labels
|
||
self.embeddings = embeddings
|
||
|
||
# 4. 构建 FAISS 索引（使用筛选后的向量）
|
||
index = self.build_faiss_index(embeddings, index_type)
|
||
|
||
# 5. 保存索引和元数据
|
||
self.save_index_and_metadata(index, output_dir)
|
||
|
||
print("=" * 60)
|
||
print("FAISS 索引构建完成!")
|
||
print(f"输出目录 : {output_dir}")
|
||
print(f"入库向量数 : {len(image_paths)}（K-Means 筛选后）")
|
||
print(f"类别数 : {len(self.class_names)}")
|
||
print(f"每类上限 : {max_per_class} 张")
|
||
print(f"特征维度 : {self.embedding_dim}")
|
||
print("=" * 60)
|
||
|
||
return index
|
||
|
||
|
||
def extract_embeddings_only(model_path: str,
|
||
train_dir: str,
|
||
output_dir: str,
|
||
embedding_dim: int = 512,
|
||
batch_size: int = 16) -> Tuple[np.ndarray, List[int]]:
|
||
"""
|
||
仅提取特征向量并保存为JSON（用于可视化）
|
||
不构建完整的FAISS索引，节省时间
|
||
|
||
Args:
|
||
model_path: 模型路径
|
||
train_dir: 训练数据目录
|
||
output_dir: 输出目录
|
||
embedding_dim: 特征向量维度
|
||
batch_size: 批处理大小
|
||
|
||
Returns:
|
||
(embeddings, labels): 特征向量数组和标签列表
|
||
"""
|
||
print("=" * 60)
|
||
print("开始提取特征向量用于可视化")
|
||
print("=" * 60)
|
||
|
||
builder = FAISSIndexBuilder(model_path, embedding_dim)
|
||
|
||
# 扫描数据
|
||
image_paths, class_names, labels = builder.scan_training_data(train_dir)
|
||
builder.image_paths = image_paths
|
||
builder.labels = labels
|
||
|
||
# 提取特征
|
||
embeddings = builder.extract_features_batch(image_paths, batch_size)
|
||
builder.embeddings = embeddings
|
||
|
||
# 仅保存 embeddings.json 和 labels.json（可视化需要的）
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
|
||
labels_json_path = os.path.join(output_dir, 'labels.json')
|
||
|
||
print(f"保存特征向量到: {embeddings_json_path}")
|
||
with open(embeddings_json_path, "w", encoding="utf-8") as f:
|
||
json.dump(embeddings.tolist(), f, ensure_ascii=False)
|
||
|
||
print(f"保存标签到: {labels_json_path}")
|
||
with open(labels_json_path, "w", encoding="utf-8") as f:
|
||
json.dump(labels, f, ensure_ascii=False)
|
||
|
||
print("=" * 60)
|
||
print(f"✓ 特征向量提取完成")
|
||
print(f" 输出目录: {output_dir}")
|
||
print(f" 样本数: {len(embeddings)}")
|
||
print(f" 特征维度: {embedding_dim}")
|
||
print("=" * 60)
|
||
|
||
return embeddings, labels
|
||
|
||
|
||
class FAISSSearcher:
|
||
"""FAISS相似度检索器"""
|
||
|
||
def __init__(self, index_dir: str, model_path: str):
|
||
"""
|
||
初始化FAISS检索器
|
||
|
||
Args:
|
||
index_dir: 索引文件目录
|
||
model_path: 模型路径
|
||
"""
|
||
self.index_dir = index_dir
|
||
self.model_path = model_path
|
||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
# 加载索引和元数据
|
||
self.load_index_and_metadata()
|
||
|
||
# 加载模型
|
||
self.load_model()
|
||
|
||
def load_index_and_metadata(self):
|
||
"""加载FAISS索引和元数据"""
|
||
# 加载FAISS索引
|
||
index_path = os.path.join(self.index_dir, 'faiss_index.bin')
|
||
self.index = faiss.read_index(index_path)
|
||
print(f"FAISS索引已加载: {self.index.ntotal} 个向量")
|
||
|
||
# 加载图片路径
|
||
paths_path = os.path.join(self.index_dir, 'image_paths.pkl')
|
||
with open(paths_path, 'rb') as f:
|
||
self.image_paths = pickle.load(f)
|
||
|
||
# 加载标签
|
||
labels_path = os.path.join(self.index_dir, 'labels.pkl')
|
||
with open(labels_path, 'rb') as f:
|
||
self.labels = pickle.load(f)
|
||
|
||
# 加载类别信息
|
||
class_info_path = os.path.join(self.index_dir, 'class_info.json')
|
||
with open(class_info_path, 'r', encoding='utf-8') as f:
|
||
self.class_info = json.load(f)
|
||
|
||
self.class_names = self.class_info['class_names']
|
||
self.class_to_idx = self.class_info['class_to_idx']
|
||
self.idx_to_class = self.class_info['idx_to_class']
|
||
self.embedding_dim = self.class_info['embedding_dim']
|
||
|
||
print(f"元数据已加载: {len(self.image_paths)} 张图片, {len(self.class_names)} 个类别")
|
||
|
||
def load_model(self):
|
||
"""加载特征提取模型"""
|
||
self.model = create_resnet50_embedding(
|
||
embedding_dim=self.embedding_dim,
|
||
pretrained=True,
|
||
use_internal_preprocess=False
|
||
)
|
||
|
||
checkpoint = torch.load(self.model_path, map_location=self.device)
|
||
if isinstance(checkpoint, dict):
|
||
if 'model_state_dict' in checkpoint:
|
||
self.model.load_state_dict(checkpoint['model_state_dict'])
|
||
elif 'backbone_state_dict' in checkpoint:
|
||
# CosFace格式：使用backbone_state_dict（只加载backbone部分）
|
||
self.model.load_state_dict(checkpoint['backbone_state_dict'])
|
||
print("检测到CosFace模型格式，使用'backbone_state_dict'加载")
|
||
elif 'state_dict' in checkpoint:
|
||
self.model.load_state_dict(checkpoint['state_dict'])
|
||
else:
|
||
self.model.load_state_dict(checkpoint)
|
||
else:
|
||
self.model.load_state_dict(checkpoint)
|
||
|
||
self.model.to(self.device)
|
||
self.model.eval()
|
||
print("特征提取模型已加载")
|
||
|
||
def search_similar(self, query_image_path: str, k: int = 5) -> List[Tuple[str, str, float]]:
|
||
"""
|
||
搜索相似图片
|
||
|
||
Args:
|
||
query_image_path: 查询图片路径
|
||
k: 返回最相似的k张图片
|
||
|
||
Returns:
|
||
List[Tuple[str, str, float]]: [(图片路径, 类别名称, 相似度分数), ...]
|
||
"""
|
||
# 提取查询图片的特征
|
||
query_image = Image.open(query_image_path).convert('RGB')
|
||
query_embedding = self.model.extract_embedding(query_image, normalize=True)
|
||
query_embedding = query_embedding.reshape(1, -1).astype(np.float32)
|
||
|
||
# 在FAISS索引中搜索
|
||
scores, indices = self.index.search(query_embedding, k)
|
||
|
||
# 整理结果
|
||
results = []
|
||
for i in range(k):
|
||
if i < len(indices[0]):
|
||
idx = indices[0][i]
|
||
score = scores[0][i]
|
||
image_path = self.image_paths[idx]
|
||
class_name = self.class_names[self.labels[idx]]
|
||
results.append((image_path, class_name, float(score)))
|
||
|
||
return results
|
||
|
||
def search_by_class(self, class_name: str, k: int = 10) -> List[Tuple[str, float]]:
|
||
"""
|
||
按类别搜索图片
|
||
|
||
Args:
|
||
class_name: 类别名称
|
||
k: 返回的图片数量
|
||
|
||
Returns:
|
||
List[Tuple[str, float]]: [(图片路径, 索引), ...]
|
||
"""
|
||
if class_name not in self.class_to_idx:
|
||
raise ValueError(f"未知类别: {class_name}")
|
||
|
||
class_idx = self.class_to_idx[class_name]
|
||
class_images = []
|
||
|
||
for i, label in enumerate(self.labels):
|
||
if label == class_idx:
|
||
class_images.append((self.image_paths[i], i))
|
||
if len(class_images) >= k:
|
||
break
|
||
|
||
return class_images
|
||
|
||
|
||
def main():
|
||
"""主函数"""
|
||
# 配置参数
|
||
# MODEL_PATH = "../model/embedding_20251011_133653/best_embedding_model.pth"
|
||
# MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth"
|
||
# MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20251113_160103/best_cosface_model.pth"
|
||
#老索引
|
||
#MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth"
|
||
#新索引
|
||
#MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20260118_220318/best_cosface_model.pth"
|
||
#新索引
|
||
#MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20260121_133207/best_cosface_model.pth"
|
||
#grid_search最佳模型 s=68, m=0.42 1.22
|
||
# MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260122_112545/best_model_s68.0_m0.42.pth"
|
||
# MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260302_092600/best_model_s64.0_m0.4.pth"
|
||
MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20260407_190824/best_model_s60.0_m0.45.pth"
|
||
# MODEL_PATH = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth"
|
||
# MODEL_PATH = "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth"
|
||
# TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train"
|
||
TRAIN_DIR = "../dataset/WholeIngredientRecognition/train"
|
||
# TRAIN_DIR = "../dataset/DishClassification/train"
|
||
|
||
# OUTPUT_DIR = "ProcessedIngredientRecognition/faiss_index"
|
||
OUTPUT_DIR = "WholeIngredientRecognition/faiss_index"
|
||
# OUTPUT_DIR = "DishClassification/faiss_index"
|
||
BATCH_SIZE = 16
|
||
INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw'
|
||
EMBEDDING_DIM = 512
|
||
MAX_PER_CLASS = 30 # 每类最多保留的代表性样本数（K-Means 筛选上限）
|
||
|
||
try:
|
||
# 构建索引
|
||
builder = FAISSIndexBuilder(MODEL_PATH, EMBEDDING_DIM)
|
||
index = builder.build_complete_index(
|
||
train_dir=TRAIN_DIR,
|
||
output_dir=OUTPUT_DIR,
|
||
batch_size=BATCH_SIZE,
|
||
index_type=INDEX_TYPE,
|
||
max_per_class=MAX_PER_CLASS
|
||
)
|
||
|
||
print("\n" + "=" * 60)
|
||
print("测试相似度检索功能")
|
||
print("=" * 60)
|
||
|
||
# 测试检索功能
|
||
searcher = FAISSSearcher(OUTPUT_DIR, MODEL_PATH)
|
||
|
||
# 随机选择一张图片进行测试
|
||
test_image_path = builder.image_paths[0]
|
||
print(f"测试图片: {test_image_path}")
|
||
|
||
# 搜索相似图片
|
||
similar_images = searcher.search_similar(test_image_path, k=5)
|
||
|
||
print("\n最相似的5张图片:")
|
||
for i, (img_path, class_name, score) in enumerate(similar_images, 1):
|
||
print(f"{i}. {os.path.basename(img_path)} (类别: {class_name}, 相似度: {score:.4f})")
|
||
|
||
# 按类别搜索
|
||
print(f"\n'{builder.class_names[0]}'类别的前5张图片:")
|
||
class_images = searcher.search_by_class(builder.class_names[0], k=5)
|
||
for i, (img_path, idx) in enumerate(class_images, 1):
|
||
print(f"{i}. {os.path.basename(img_path)} (索引: {idx})")
|
||
|
||
except Exception as e:
|
||
print(f"错误: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |