228 lines
6.5 KiB
Markdown
228 lines
6.5 KiB
Markdown
# 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 许可证。 |