From 8d6555747b8ac8567d330ebb58e435da4feb4537 Mon Sep 17 00:00:00 2001 From: zhanghuan <1262329256@qq.com> Date: Mon, 22 Sep 2025 11:29:22 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0FAISS=E5=90=91=E9=87=8F?= =?UTF-8?q?=E6=95=B0=E6=8D=AE=E5=BA=93=E5=AF=B9=E6=89=80=E6=9C=89=E8=AE=AD?= =?UTF-8?q?=E7=BB=83=E5=9B=BE=E7=89=87=E8=BF=9B=E8=A1=8C=E5=90=91=E9=87=8F?= =?UTF-8?q?=E5=8C=96=E7=9A=84=E8=BF=87=E7=A8=8B=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- faiss_vector_db/README_FAISS.md | 228 +++++++++ faiss_vector_db/build_faiss_index.py | 517 +++++++++++++++++++++ faiss_vector_db/faiss_similarity_search.py | 318 +++++++++++++ faiss_vector_db/run_faiss_demo.py | 244 ++++++++++ 4 files changed, 1307 insertions(+) create mode 100644 faiss_vector_db/README_FAISS.md create mode 100644 faiss_vector_db/build_faiss_index.py create mode 100644 faiss_vector_db/faiss_similarity_search.py create mode 100644 faiss_vector_db/run_faiss_demo.py diff --git a/faiss_vector_db/README_FAISS.md b/faiss_vector_db/README_FAISS.md new file mode 100644 index 0000000..0c5d458 --- /dev/null +++ b/faiss_vector_db/README_FAISS.md @@ -0,0 +1,228 @@ +# FAISS相似度检索系统 + +基于ResNet50特征提取和FAISS索引的食物图片相似度检索系统。 + +## 功能特性 + +- ✅ 使用训练好的ResNet50 Embedding模型提取图片特征向量 +- ✅ 构建高效的FAISS索引用于快速相似度检索 +- ✅ 支持多种索引类型（精确搜索、近似搜索） +- ✅ 完整的标签映射和元数据管理 +- ✅ 命令行界面和可视化结果展示 +- ✅ 批量测试和性能评估 + +## 文件结构 + +``` +FoodClassifier/ +├── build_faiss_index.py # FAISS索引构建器 +├── faiss_similarity_search.py # 相似度检索演示程序 +├── README_FAISS.md # 使用说明文档 +├── net/ +│ ├── resnet_embedding.py # ResNet50特征提取网络 +│ └── food_net.py # 原始分类网络 +├── model/ +│ └── embedding_20250917_145342/ +│ └── best_embedding_model.pth # 训练好的模型 +├── dataset/ +│ └── train/ # 训练数据集 +│ ├── 回锅肉/ +│ ├── 炒细面/ +│ ├── 西红柿鸡蛋/ +│ └── 麻辣小面/ +└── faiss_vector_db/ + └── faiss_index/ # 生成的FAISS索引文件 + ├── faiss_index.bin # FAISS索引 + ├── image_paths.pkl # 图片路径映射 + ├── labels.pkl # 标签映射 + ├── class_info.json # 类别信息 + └── embeddings.npy # 特征向量数组 +``` + +## 安装依赖 + +```bash +pip install torch torchvision +pip install faiss-cpu # 或 faiss-gpu（如果有GPU） +pip install pillow numpy tqdm matplotlib +``` + +## 使用方法 + +### 1. 构建FAISS索引 + +首次使用需要构建FAISS索引： + +```bash +python build_faiss_index.py +``` + +这将： +- 加载训练好的ResNet50 Embedding模型 +- 遍历训练集所有图片，提取512维特征向量 +- 构建FAISS索引文件 +- 保存标签映射和元数据 + +### 2. 相似度检索 + +#### 查看索引信息 +```bash +python faiss_similarity_search.py --info +``` + +#### 单张图片相似度检索 +```bash +# 基本搜索 +python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg + +# 显示图片结果 +python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images + +# 保存结果图片 +python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images --save_result + +# 返回更多结果 +python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --k 10 +``` + +#### 批量测试 +```bash +# 批量测试测试集 +python faiss_similarity_search.py --batch_test dataset/test --max_tests 20 + +# 批量测试验证集 +python faiss_similarity_search.py --batch_test dataset/val --max_tests 10 +``` + +## 核心类说明 + +### FAISSIndexBuilder +负责构建FAISS索引的核心类： + +```python +from faiss_vector_db.build_faiss_index import FAISSIndexBuilder + +# 创建索引构建器 +builder = FAISSIndexBuilder( + model_path="../model/embedding_20250917_145342/best_embedding_model.pth", + embedding_dim=512 +) + +# 构建完整索引 +index = builder.build_complete_index( + train_dir="../dataset/train", + output_dir="faiss_index", + batch_size=16, + index_type='flat' # 'flat', 'ivf', 'hnsw' +) +``` + +### FAISSSearcher +负责相似度检索的核心类： + +```python +from faiss_vector_db.build_faiss_index import FAISSSearcher + +# 创建检索器 +searcher = FAISSSearcher( + index_dir="faiss_vector_db/faiss_index", + model_path="model/embedding_20250917_145342/best_embedding_model.pth" +) + +# 搜索相似图片 +results = searcher.search_similar("test_image.jpg", k=5) +# 返回: [(图片路径, 类别名称, 相似度分数), ...] + +# 按类别搜索 +class_images = searcher.search_by_class("回锅肉", k=10) +# 返回: [(图片路径, 索引), ...] +``` + +## 索引类型说明 + +### 1. Flat索引 (推荐) +- **类型**: `'flat'` +- **特点**: 精确搜索，速度快 +- **适用**: 中小规模数据集（< 100万向量） +- **优点**: 搜索结果完全准确 +- **缺点**: 内存占用较大 + +### 2. IVF索引 +- **类型**: `'ivf'` +- **特点**: 近似搜索，内存效率高 +- **适用**: 大规模数据集（> 100万向量） +- **优点**: 内存占用小，搜索速度快 +- **缺点**: 搜索结果可能不是最优 + +### 3. HNSW索引 +- **类型**: `'hnsw'` +- **特点**: 图结构索引，搜索速度极快 +- **适用**: 需要极快搜索速度的场景 +- **优点**: 搜索速度最快 +- **缺点**: 构建时间较长，内存占用中等 + +## 性能优化建议 + +### 1. 批处理大小 +- GPU内存充足：`batch_size=32` +- GPU内存有限：`batch_size=16` +- 仅CPU：`batch_size=8` + +### 2. 索引选择 +- 数据量 < 10万：使用 `'flat'` +- 数据量 10万-100万：使用 `'ivf'` +- 数据量 > 100万：使用 `'hnsw'` + +### 3. 特征向量维度 +- 默认512维提供良好的精度 +- 可调整为256维以节省存储空间 +- 1024维可提供更高精度（需重新训练） + +## 输出文件说明 + +构建完成后，`faiss_vector_db/faiss_index/` 目录包含： + +- **faiss_index.bin**: FAISS索引文件 +- **image_paths.pkl**: 图片路径列表（pickle格式） +- **labels.pkl**: 对应的标签列表（pickle格式） +- **class_info.json**: 类别信息和元数据（JSON格式） +- **embeddings.npy**: 特征向量数组（NumPy格式，可选） + +## 常见问题 + +### Q: 如何更新索引？ +A: 重新运行 `build_faiss_index.py`，会覆盖原有索引。 + +### Q: 如何添加新的图片类别？ +A: 将新类别图片添加到 `dataset/train/` 目录，然后重新构建索引。 + +### Q: 搜索速度慢怎么办？ +A: 尝试使用 `'ivf'` 或 `'hnsw'` 索引类型，或减少返回结果数量。 + +### Q: 内存不足怎么办？ +A: 减少 `batch_size`，或使用 `'ivf'` 索引类型。 + +### Q: 如何在其他项目中使用？ +A: 复制 `build_faiss_index.py` 和相关模型文件，修改路径配置即可。 + +## 扩展功能 + +系统支持以下扩展： + +1. **多模态检索**: 结合文本特征进行检索 +2. **在线更新**: 支持动态添加新图片到索引 +3. **分布式部署**: 支持多机部署和负载均衡 +4. **Web接口**: 提供REST API接口 +5. **移动端集成**: 支持移动端实时检索 + +## 技术细节 + +- **特征提取**: ResNet50 + 全连接层 → 512维向量 +- **相似度计算**: 余弦相似度（内积，向量已归一化） +- **索引结构**: FAISS多种索引类型支持 +- **数据格式**: 支持 JPG、PNG、BMP 等常见图片格式 +- **预处理**: ImageNet标准化，224x224分辨率 + +## 许可证 + +本项目遵循 MIT 许可证。 \ No newline at end of file diff --git a/faiss_vector_db/build_faiss_index.py b/faiss_vector_db/build_faiss_index.py new file mode 100644 index 0000000..47f407d --- /dev/null +++ b/faiss_vector_db/build_faiss_index.py @@ -0,0 +1,517 @@ +""" +食物分类器 - 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: + self.model.load_state_dict(checkpoint['model_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("模型加载成功!") + + 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: + 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') + with open(labels_path, 'wb') as f: + pickle.dump(self.labels, f) + 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_array = np.array(self.embeddings) + np.save(embeddings_path, embeddings_array) + print(f"特征向量已保存到: {embeddings_path}") + + def build_complete_index(self, train_dir: str, output_dir: str, + batch_size: int = 32, index_type: str = 'flat'): + """ + 完整的索引构建流程 + + Args: + train_dir: 训练数据目录 + output_dir: 输出目录 + batch_size: 批处理大小 + index_type: 索引类型 + """ + print("=" * 60) + print("开始构建FAISS索引") + print("=" * 60) + + # 1. 扫描训练数据 + image_paths, class_names, labels = self.scan_training_data(train_dir) + self.image_paths = image_paths + self.labels = labels + + # 2. 提取特征向量 + embeddings = self.extract_features_batch(image_paths, batch_size) + self.embeddings = embeddings + + # 3. 构建FAISS索引 + index = self.build_faiss_index(embeddings, index_type) + + # 4. 保存索引和元数据 + self.save_index_and_metadata(index, output_dir) + + print("=" * 60) + print("FAISS索引构建完成!") + print(f"输出目录: {output_dir}") + print(f"总图片数: {len(image_paths)}") + print(f"类别数: {len(self.class_names)}") + print(f"特征维度: {self.embedding_dim}") + print("=" * 60) + + return index + + +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 '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_20250917_145342/best_embedding_model.pth" + TRAIN_DIR = "../dataset/train" + OUTPUT_DIR = "faiss_index" + BATCH_SIZE = 16 + INDEX_TYPE = 'flat' # 'flat', 'ivf', 'hnsw' + EMBEDDING_DIM = 512 + + 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 + ) + + 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() \ No newline at end of file diff --git a/faiss_vector_db/faiss_similarity_search.py b/faiss_vector_db/faiss_similarity_search.py new file mode 100644 index 0000000..1933795 --- /dev/null +++ b/faiss_vector_db/faiss_similarity_search.py @@ -0,0 +1,318 @@ +""" +FAISS相似度检索演示程序 +提供简单的命令行界面来测试相似度检索功能 + +功能： +1. 加载预构建的FAISS索引 +2. 对指定图片进行相似度检索 +3. 显示检索结果和统计信息 +4. 支持批量检索和性能测试 + +""" + +import os +import sys +import time +import argparse +from typing import List, Tuple +from PIL import Image +import matplotlib.pyplot as plt + +# 导入构建的FAISS检索器 +from faiss_vector_db.build_faiss_index import FAISSSearcher + + +class SimilaritySearchDemo: + """相似度检索演示类""" + + def __init__(self, index_dir: str, model_path: str): + """ + 初始化演示程序 + + Args: + index_dir: FAISS索引目录 + model_path: 模型路径 + """ + self.searcher = FAISSSearcher(index_dir, model_path) + + # 设置中文字体（用于matplotlib显示） + try: + # 尝试使用系统中文字体 + plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans'] + plt.rcParams['axes.unicode_minus'] = False + except: + print("警告: 无法设置中文字体，图表中的中文可能显示异常") + + def search_and_display(self, query_image_path: str, k: int = 5, + show_images: bool = False, save_result: bool = False): + """ + 搜索相似图片并显示结果 + + Args: + query_image_path: 查询图片路径 + k: 返回最相似的k张图片 + show_images: 是否显示图片 + save_result: 是否保存结果图片 + """ + print(f"\n查询图片: {query_image_path}") + print("-" * 50) + + # 检查文件是否存在 + if not os.path.exists(query_image_path): + print(f"错误: 文件不存在 - {query_image_path}") + return + + # 开始计时 + start_time = time.time() + + try: + # 执行相似度搜索 + results = self.searcher.search_similar(query_image_path, k) + + # 计算耗时 + search_time = time.time() - start_time + + # 显示结果 + print(f"搜索完成，耗时: {search_time:.3f}秒") + print(f"找到 {len(results)} 个相似结果:\n") + + for i, (img_path, class_name, score) in enumerate(results, 1): + rel_path = os.path.relpath(img_path) + print(f"{i:2d}. {os.path.basename(img_path):20s} " + f"类别: {class_name:8s} " + f"相似度: {score:.4f}") + + # 统计类别分布 + class_counts = {} + for _, class_name, _ in results: + class_counts[class_name] = class_counts.get(class_name, 0) + 1 + + print(f"\n类别分布:") + for class_name, count in class_counts.items(): + print(f" {class_name}: {count} 张") + + # 显示图片 + if show_images: + self.display_results(query_image_path, results, save_result) + + except Exception as e: + print(f"搜索失败: {e}") + import traceback + traceback.print_exc() + + def display_results(self, query_path: str, results: List[Tuple[str, str, float]], + save_result: bool = False): + """ + 显示搜索结果图片 + + Args: + query_path: 查询图片路径 + results: 搜索结果 + save_result: 是否保存结果图片 + """ + try: + # 计算子图布局 + n_results = len(results) + n_cols = min(3, n_results + 1) # 包括查询图片 + n_rows = (n_results + 1 + n_cols - 1) // n_cols + + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4*n_cols, 4*n_rows)) + if n_rows == 1: + axes = axes.reshape(1, -1) + + # 显示查询图片 + query_img = Image.open(query_path).convert('RGB') + axes[0, 0].imshow(query_img) + axes[0, 0].set_title(f"查询图片\n{os.path.basename(query_path)}", fontsize=10) + axes[0, 0].axis('off') + + # 显示搜索结果 + for i, (img_path, class_name, score) in enumerate(results): + row = (i + 1) // n_cols + col = (i + 1) % n_cols + + if row < n_rows and col < n_cols: + try: + img = Image.open(img_path).convert('RGB') + axes[row, col].imshow(img) + axes[row, col].set_title( + f"#{i+1} {class_name}\n" + f"{os.path.basename(img_path)}\n" + f"相似度: {score:.3f}", + fontsize=9 + ) + axes[row, col].axis('off') + except Exception as e: + axes[row, col].text(0.5, 0.5, f"加载失败\n{e}", + ha='center', va='center') + axes[row, col].axis('off') + + # 隐藏多余的子图 + for i in range(n_results + 1, n_rows * n_cols): + row = i // n_cols + col = i % n_cols + if row < n_rows and col < n_cols: + axes[row, col].axis('off') + + plt.tight_layout() + + # 保存结果 + if save_result: + timestamp = time.strftime("%Y%m%d_%H%M%S") + save_path = f"similarity_search_result_{timestamp}.png" + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"结果图片已保存到: {save_path}") + + plt.show() + + except Exception as e: + print(f"显示图片失败: {e}") + + def batch_search_test(self, test_dir: str, k: int = 5, max_tests: int = 10): + """ + 批量搜索测试 + + Args: + test_dir: 测试图片目录 + k: 每次搜索返回的结果数 + max_tests: 最大测试数量 + """ + print(f"\n批量搜索测试 - 目录: {test_dir}") + print("=" * 60) + + # 获取测试图片 + test_images = [] + for root, dirs, files in os.walk(test_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')): + test_images.append(os.path.join(root, file)) + if len(test_images) >= max_tests: + break + if len(test_images) >= max_tests: + break + + if not test_images: + print("未找到测试图片") + return + + print(f"找到 {len(test_images)} 张测试图片") + + # 执行批量测试 + total_time = 0 + success_count = 0 + + for i, img_path in enumerate(test_images, 1): + print(f"\n[{i}/{len(test_images)}] 测试: {os.path.basename(img_path)}") + + try: + start_time = time.time() + results = self.searcher.search_similar(img_path, k) + search_time = time.time() - start_time + + total_time += search_time + success_count += 1 + + print(f" 搜索耗时: {search_time:.3f}秒") + print(f" 前3个结果:") + for j, (result_path, class_name, score) in enumerate(results[:3], 1): + print(f" {j}. {class_name} (相似度: {score:.3f})") + + except Exception as e: + print(f" 搜索失败: {e}") + + # 显示统计信息 + print("\n" + "=" * 60) + print("批量测试统计:") + print(f" 总测试数: {len(test_images)}") + print(f" 成功数: {success_count}") + print(f" 失败数: {len(test_images) - success_count}") + if success_count > 0: + print(f" 平均搜索时间: {total_time/success_count:.3f}秒") + print(f" 总耗时: {total_time:.3f}秒") + + def show_index_info(self): + """显示索引信息""" + print("\n" + "=" * 60) + print("FAISS索引信息") + print("=" * 60) + + print(f"索引目录: {self.searcher.index_dir}") + print(f"总向量数: {self.searcher.index.ntotal}") + print(f"特征维度: {self.searcher.embedding_dim}") + print(f"类别数: {len(self.searcher.class_names)}") + print(f"类别列表: {', '.join(self.searcher.class_names)}") + + # 显示每个类别的图片数量 + print(f"\n各类别图片数量:") + class_counts = {} + for label in self.searcher.labels: + class_name = self.searcher.class_names[label] + class_counts[class_name] = class_counts.get(class_name, 0) + 1 + + for class_name in self.searcher.class_names: + count = class_counts.get(class_name, 0) + print(f" {class_name}: {count} 张") + + print(f"\n构建时间: {self.searcher.class_info.get('build_time', '未知')}") + + +def main(): + """主函数""" + parser = argparse.ArgumentParser(description='FAISS相似度检索演示程序') + parser.add_argument('--index_dir', type=str, default='faiss_vector_db/faiss_index', + help='FAISS索引目录路径') + parser.add_argument('--model_path', type=str, + default='model/embedding_20250917_145342/best_embedding_model.pth', + help='模型文件路径') + parser.add_argument('--query', type=str, help='查询图片路径') + parser.add_argument('--k', type=int, default=5, help='返回最相似的k张图片') + parser.add_argument('--show_images', action='store_true', help='显示搜索结果图片') + parser.add_argument('--save_result', action='store_true', help='保存结果图片') + parser.add_argument('--batch_test', type=str, help='批量测试目录') + parser.add_argument('--max_tests', type=int, default=10, help='批量测试最大数量') + parser.add_argument('--info', action='store_true', help='显示索引信息') + + args = parser.parse_args() + + try: + # 检查索引目录 + if not os.path.exists(args.index_dir): + print(f"错误: 索引目录不存在 - {args.index_dir}") + print("请先运行 build_faiss_index.py 构建索引") + return + + # 检查模型文件 + if not os.path.exists(args.model_path): + print(f"错误: 模型文件不存在 - {args.model_path}") + return + + # 初始化演示程序 + demo = SimilaritySearchDemo(args.index_dir, args.model_path) + + # 显示索引信息 + if args.info: + demo.show_index_info() + + # 单张图片搜索 + if args.query: + demo.search_and_display(args.query, args.k, args.show_images, args.save_result) + + # 批量测试 + if args.batch_test: + demo.batch_search_test(args.batch_test, args.k, args.max_tests) + + # 如果没有指定任何操作，显示帮助信息 + if not any([args.query, args.batch_test, args.info]): + print("使用示例:") + print(f" python {sys.argv[0]} --info # 显示索引信息") + print(f" python {sys.argv[0]} --query dataset/test/回锅肉/aug_000.jpg --show_images") + print(f" python {sys.argv[0]} --batch_test dataset/test --max_tests 5") + + except Exception as e: + print(f"程序执行失败: {e}") + import traceback + traceback.print_exc() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/faiss_vector_db/run_faiss_demo.py b/faiss_vector_db/run_faiss_demo.py new file mode 100644 index 0000000..f6eaf5b --- /dev/null +++ b/faiss_vector_db/run_faiss_demo.py @@ -0,0 +1,244 @@ +""" +FAISS相似度检索系统快速演示脚本 +一键运行完整的索引构建和检索演示 +""" + +import os +import time + + +def check_dependencies(): + """检查依赖包""" + required_packages = [ + 'torch', 'torchvision', 'faiss', 'PIL', 'numpy', 'tqdm', 'matplotlib' + ] + + missing_packages = [] + + for package in required_packages: + try: + if package == 'PIL': + import PIL + elif package == 'faiss': + import faiss + else: + __import__(package) + except ImportError: + missing_packages.append(package) + + if missing_packages: + print("缺少以下依赖包:") + for pkg in missing_packages: + if pkg == 'faiss': + print(f" {pkg} (安装命令: pip install faiss-cpu)") + else: + print(f" {pkg}") + print("\n请先安装缺少的依赖包") + return False + + return True + +def check_files(): + """检查必要文件""" + required_files = [ + 'net/resnet_embedding.py', + 'settings/settings.py', + 'model/embedding_20250917_145342/best_embedding_model.pth', + 'dataset/train' + ] + + missing_files = [] + + for file_path in required_files: + if not os.path.exists(file_path): + missing_files.append(file_path) + + if missing_files: + print("缺少以下必要文件:") + for file_path in missing_files: + print(f" {file_path}") + return False + + return True + +def run_index_building(): + """运行索引构建""" + print("=" * 60) + print("步骤 1: 构建FAISS索引") + print("=" * 60) + + try: + from faiss_vector_db.build_faiss_index import main as build_main + build_main() + return True + except Exception as e: + print(f"索引构建失败: {e}") + import traceback + traceback.print_exc() + return False + +def run_search_demo(): + """运行检索演示""" + print("\n" + "=" * 60) + print("步骤 2: 相似度检索演示") + print("=" * 60) + + try: + from faiss_vector_db.build_faiss_index import FAISSSearcher + + # 配置参数 + index_dir = "faiss_index" + model_path = "../model/embedding_20250917_145342/best_embedding_model.pth" + + # 检查索引是否存在 + if not os.path.exists(os.path.join(index_dir, 'faiss_index.bin')): + print("错误: FAISS索引不存在，请先构建索引") + return False + + # 创建检索器 + searcher = FAISSSearcher(index_dir, model_path) + + # 显示索引信息 + print(f"索引信息:") + print(f" 总向量数: {searcher.index.ntotal}") + print(f" 特征维度: {searcher.embedding_dim}") + print(f" 类别数: {len(searcher.class_names)}") + print(f" 类别: {', '.join(searcher.class_names)}") + + # 找一些测试图片 + test_images = [] + test_dirs = ['dataset/test', 'dataset/val', 'dataset/train'] + + for test_dir in test_dirs: + if os.path.exists(test_dir): + for root, dirs, files in os.walk(test_dir): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', '.png')): + test_images.append(os.path.join(root, file)) + if len(test_images) >= 3: # 只测试3张图片 + break + if len(test_images) >= 3: + break + if len(test_images) >= 3: + break + + if not test_images: + print("未找到测试图片") + return False + + # 执行相似度检索演示 + print(f"\n开始相似度检索演示 (测试 {len(test_images)} 张图片):") + print("-" * 50) + + total_time = 0 + + for i, test_image in enumerate(test_images, 1): + print(f"\n[{i}/{len(test_images)}] 测试图片: {os.path.relpath(test_image)}") + + start_time = time.time() + results = searcher.search_similar(test_image, k=5) + search_time = time.time() - start_time + total_time += search_time + + print(f"搜索耗时: {search_time:.3f}秒") + print("最相似的5张图片:") + + for j, (img_path, class_name, score) in enumerate(results, 1): + rel_path = os.path.relpath(img_path) + print(f" {j}. {os.path.basename(img_path):20s} " + f"类别: {class_name:8s} " + f"相似度: {score:.4f}") + + # 统计类别分布 + class_counts = {} + for _, class_name, _ in results: + class_counts[class_name] = class_counts.get(class_name, 0) + 1 + + print(f" 类别分布: {dict(class_counts)}") + + print(f"\n演示完成!") + print(f"平均搜索时间: {total_time/len(test_images):.3f}秒") + + return True + + except Exception as e: + print(f"检索演示失败: {e}") + import traceback + traceback.print_exc() + return False + +def show_usage_examples(): + """显示使用示例""" + print("\n" + "=" * 60) + print("使用示例") + print("=" * 60) + + print("\n1. 查看索引信息:") + print(" python faiss_similarity_search.py --info") + + print("\n2. 单张图片相似度检索:") + print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg") + + print("\n3. 显示检索结果图片:") + print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images") + + print("\n4. 批量测试:") + print(" python faiss_similarity_search.py --batch_test dataset/test --max_tests 10") + + print("\n5. 保存结果图片:") + print(" python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images --save_result") + +def main(): + """主函数""" + print("FAISS相似度检索系统 - 快速演示") + print("=" * 60) + + # 检查依赖 + print("检查依赖包...") + if not check_dependencies(): + return + print("✓ 依赖包检查通过") + + # 检查文件 + print("检查必要文件...") + if not check_files(): + return + print("✓ 文件检查通过") + + # 检查是否已有索引 + index_exists = os.path.exists("faiss_index/faiss_index.bin") + + if index_exists: + print("✓ 发现已存在的FAISS索引") + + # 询问是否重新构建 + while True: + choice = input("\n是否重新构建索引? (y/n): ").lower().strip() + if choice in ['y', 'yes', '是']: + if not run_index_building(): + return + break + elif choice in ['n', 'no', '否']: + print("跳过索引构建，直接进行检索演示") + break + else: + print("请输入 y 或 n") + else: + print("未发现FAISS索引，开始构建...") + if not run_index_building(): + return + + # 运行检索演示 + if not run_search_demo(): + return + + # 显示使用示例 + show_usage_examples() + + print("\n" + "=" * 60) + print("演示完成! 🎉") + print("您现在可以使用 faiss_similarity_search.py 进行更多测试") + print("=" * 60) + +if __name__ == "__main__": + main() \ No newline at end of file