# 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="DishClassification/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 许可证。