增加FAISS向量数据库对所有训练图片进行向量化的过程。

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 许可证。

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()

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()

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()