增加各个脚本说明，整理项目目录。

2026-03-02 17:26:00 +08:00
parent 7b9a370054
commit 7af923f1ce
3 changed files with 132 additions and 668 deletions
@@ -0,0 +1,132 @@
 # CLAUDE.md
 This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
 ## 项目概述
 本项目采用**度量学习（Metric Learning）**的方式进行菜品识别和食材识别。核心思路是用 ResNet50 提取图像嵌入向量，配合 FAISS 向量库进行相似度检索，而非传统的端到端分类，从而支持无需重新训练即可快速添加新类别。
 **三个并行识别任务：**
 - `DishClassification` — 菜品识别
 - `WholeIngredientRecognition` — 整体食材识别（未处理）
 - `ProcessedIngredientRecognition` — 处理后食材识别（切块等）
 ## 常用命令
 **训练嵌入模型（CosFace 损失，主要方案）：**
 ```bash
 python train/train_cosface_embedding.py
 ```
 **CosFace 超参数网格搜索（scale × margin 组合）：**
 ```bash
 python train/grid_search_cosface.py
 ```
 **训练嵌入模型（Triplet + Center Loss，备用方案,目前没使用这个了，主要是未采集的图片和现有分类的相似度太高了）：**
 ```bash
 python train/train_triplet_embedding.py
 ```
 **基于训练好的模型构建 FAISS 索引：**
 ```bash
 python faiss_vector_db/build_faiss_index.py
 ```
 **启动主识别界面（嵌入向量方案）：**
 ```bash
 python classifier/embedding_food_classifier_app.py
 ```
 **启动数据管理与下载界面：**
 ```bash
 python data_management/app.py
 ```
 **训练 SegFormer 食物区域分割模型：**
 ```bash
 python SegFormer/training/train_minimal.py
 ```
 **测试 SegFormer 推理效果：**
 ```bash
 python SegFormer/inference/test_model.py
 ```
 ## 架构说明
 ### 度量学习流水线
 ```
 输入图像 → ResNet50 主干网络（ImageNet1K_V2 预训练）
         → 512 维嵌入向量 → L2 归一化
         → FAISS IndexFlatIP（余弦相似度检索）
         → 返回 Top-K 相似食物及置信度分数
 ```
 **训练损失函数：**
 - **CosFace**（`train_cosface_embedding.py`）：大间距余弦损失；scale ≈ 64，margin ≈ 0.35–0.40（因任务而异）。泛化能力更强，为首选方案。
 - **Triplet + Center Loss**（`train_triplet_embedding.py`）：在线难样本挖掘，margin = 0.5，Center Loss 权重因任务而异（0.5–20）。
 ### 关键模块
 | 路径 | 功能说明 |
 |------|---------|
 | `net/resnet_embedding.py` | ResNet50 嵌入主干网络（512 维，L2 归一化） |
 | `net/food_net.py` | 遗留的自定义 CNN，用于端到端分类（4 类） |
 | `faiss_vector_db/faiss_manager.py` | FAISS 索引的增删查存核心类 |
 | `faiss_vector_db/build_faiss_index.py` | 遍历数据集、提取嵌入、构建索引 |
 | `classifier/embedding_food_classifier_app.py` | 主 GUI（CustomTkinter，支持拖拽上传） |
 | `exp_multimodal/vlm_classifier.py` | 实验性：VLM（Ollama/Kimi）+ FAISS 开放集识别 |
 | `data_management/app.py` | 数据采集 GUI，从 MySQL 后端下载图片 |
 | `settings/settings.py` | 全局配置（NUM_CLASSES、MODEL_DIR、DEVICE 等） |
 ### FAISS 索引存储结构
 每个任务在 `faiss_vector_db/<任务名>/faiss_index/` 下独立存储：
 - `food_embeddings.index` — FAISS 二进制索引文件
 - `labels.json` / `food_labels.json` — 标签到食物名称的映射
 - `embeddings.npy` — 原始向量的 numpy 数组
 ### 数据集目录结构
 ```
 dataset/<任务名>/
 ├── train/<类别名>/  （图片）
 ├── val/<类别名>/
 └── test/<类别名>/
 ```
 ### VLM 多模态集成（实验性）
 `exp_multimodal/` 包含一套绕过封闭类别列表的开放集识别流水线：
 - `vlm_providers/ollama_provider.py` — 本地 Ollama 服务（Qwen2.5-VL、LLaVA），地址：`http://192.168.1.250:11434`
 - `vlm_providers/kimi_provider.py` — Kimi/Moonshot API（`moonshot-v1-128k-vision-preview`）
 - `build_dish_name_index.py` — 使用 `bge-large-zh-v1.5` 为菜品名称构建文本向量索引
 - VLM 配置保存在 `vlm_config.json`（提供商选择、API Key、候选菜品列表）
 ### SegFormer 食物分割
 `SegFormer/` 包含用于提取食物区域的微调 SegFormer：
 - 两阶段训练：先冻结编码器训练，再全量微调
 - 损失函数：CrossEntropy + 0.5 × Dice Loss
 - 数据以 COCO 格式存放于 `SegFormer/data/`
 ## 模型文件
 训练好的模型保存在 `model/` 下的版本子目录中：
 - `best_cosface_model.pth` — CosFace 训练的嵌入模型
 - `best_embedding_model.pth` — Triplet/Center Loss 训练的嵌入模型
 - `grid_search_<时间戳>/best_model_s<scale>_m<margin>.pth` — 网格搜索最优结果
 当前使用的模型路径通过 `settings/settings.py` 中的 `MODEL_DIR` 指定。
 ## 依赖安装
 ```bash
 pip install -r requirements.txt
 ```
 核心依赖：`torch==2.7.0`、`torchvision==0.22.0`、`faiss-cpu>=1.7.0`、`customtkinter==5.2.2`、`tkinterdnd2==0.4.3`、`opencv-python==4.11.0.86`。
 数据管理界面连接远端 MySQL 数据库（`192.168.1.250:3308`）获取图片元数据，同时使用本地 SQLite 记录下载历史。
@@ -1,149 +0,0 @@
 # CosFace超参数网格搜索使用指南
 ## 概述
 已完成CosFace超参数网格搜索功能的实施，用于优化`scale (s)`和`margin (m)`两个核心超参数。
 ## 修改文件清单
 ### 1. `requirements.txt`
 - ✅ 添加依赖：`pandas>=2.0.0` 和 `seaborn>=0.12.0`
 ### 2. `train/train_cosface_embedding.py`
 - ✅ `TaskConfig`添加`test_dir`字段
 - ✅ `TASKS`字典中所有任务添加测试集路径
 ### 3. `train/grid_search_cosface.py` (新建)
 - ✅ 完整的网格搜索脚本 (~350行)
 ## 搜索空间配置
 当前为dish任务配置的搜索空间：
 ```python
 GRID_PARAMS = {
    'dish': {
        's': [56.0, 60.0, 64.0, 68.0],    # 4个scale值
        'm': [0.32, 0.35, 0.38, 0.40],    # 4个margin值
    }
 }
 ```
 **总配置数**: 4 × 4 = **16组实验**
 ## 使用方法
 ### 1. 安装依赖
 ```bash
 pip install pandas>=2.0.0 seaborn>=0.12.0
 ```
 ### 2. 运行网格搜索
 ```bash
 # 完整搜索（16组配置）
 python train/grid_search_cosface.py --task dish
 # 测试运行（限制配置数）
 python train/grid_search_cosface.py --task dish --max_configs 2
 # 自定义参数
 python train/grid_search_cosface.py --task dish --epochs 100 --patience 10 --min_epochs 25
 ```
 ### 3. 参数说明
 - `--task`: 任务名称 (`dish` / `whole_ingredient` / `processed_ingredient`)
 - `--max_configs`: 限制最大配置数（用于测试，可选）
 - `--epochs`: 每个配置的最大训练轮数（默认100）
 - `--patience`: 早停容忍轮数（默认10）
 - `--min_epochs`: 最小训练轮数（默认25）
 ## 输出文件
 运行后会在`model/DishClassification/grid_search_YYYYMMDD_HHMMSS/`目录下生成：
 1. **`grid_search_results.csv`** - 所有配置的详细结果表格
 2. **`heatmap.png`** - 参数热力图（测试集准确率）
 3. **`summary.txt`** - 搜索总结报告（含最佳配置）
 4. **`model_s{s}_m{m}.pth`** - 每个配置的模型权重
 5. **`grid_search.log`** - 完整训练日志
 ## 评估指标
 - **主要指标**: 测试集准确率 (`test_acc`) - 用于选择最佳配置
 - **辅助指标**: 验证集准确率、训练轮数、训练时间
 ## 应用最佳配置
 网格搜索完成后：
 1. 查看`summary.txt`找到最佳配置
 2. 手动更新`train/train_cosface_embedding.py`中的`TASKS`字典：
 ```python
 'dish': TaskConfig(
    # ... 其他配置保持不变 ...
    cosface_s=64.0,  # 更新为最佳s值
    cosface_m=0.38,  # 更新为最佳m值
 )
 ```
 3. 后续训练将自动使用最佳配置
 ## 预计耗时
 基于以下假设：
 - 每个配置平均训练25-40个epoch（早停机制）
 - 每个epoch约1-2分钟
 - **单个配置**: ~30-60分钟
 - **16组配置总耗时**: ~8-16小时
 **建议**: 使用GPU运行，可在夜间或周末执行完整搜索。
 ## 注意事项
 1. **数据集要求**: 确保`dataset/DishClassification/test/`目录存在且有数据
 2. **GPU推荐**: 网格搜索计算量大，强烈建议使用GPU
 3. **磁盘空间**: 每个配置约占用500MB，16组需8GB空间
 4. **中断恢复**: 当前版本不支持断点续训，建议一次性完成
 ## 高级用法
 ### 修改搜索空间
 编辑`train/grid_search_cosface.py`中的`GRID_PARAMS`字典：
 ```python
 GRID_PARAMS = {
    'dish': {
        's': [60.0, 64.0, 68.0, 72.0],     # 自定义scale范围
        'm': [0.30, 0.35, 0.40, 0.45],     # 自定义margin范围
    }
 }
 ```
 ### 调整早停策略
 通过命令行参数调整：
 ```bash
 python train/grid_search_cosface.py --task dish --patience 15 --min_epochs 30
 ```
 ## 故障排除
 **问题1**: `ModuleNotFoundError: No module named 'pandas'`
 - 解决: `pip install pandas seaborn`
 **问题2**: 测试集目录不存在
 - 解决: 确认`dataset/DishClassification/test/`路径正确且有数据
 **问题3**: CUDA out of memory
 - 解决: 减小`batch_size`或在CPU上运行（速度较慢）
 ## 示例结果解读
 `summary.txt`示例：
 ```
 最佳配置:
  s (scale) = 64.0
  m (margin) = 0.38
  测试集准确率 = 98.50%
  验证集准确率 = 100.00%
  训练轮数 = 32
  训练时间 = 45.3 分钟
 ```
 这表示s=64.0, m=0.38是最优组合，在测试集上达到98.50%准确率。
@@ -1,519 +0,0 @@
 """
 食物分类器 - FAISS索引构建器
 将训练集图片转换为特征向量并构建FAISS索引用于快速相似度检索
 功能：
 1. 加载训练好的ResNet50 Embedding模型
 2. 遍历训练集所有图片，提取特征向量
 3. 构建FAISS索引文件
 4. 创建标签映射文件
 5. 实现快速相似度检索功能
 作者: zhangpu
 日期: 2025/9/22
 """
 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_20250930_102826/best_embedding_model.pth"
    TRAIN_DIR = "dataset/train"
    OUTPUT_DIR = "faiss_vector_db/DishClassification/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()