增加FAISS向量数据库对所有训练图片进行向量化的过程。
This commit is contained in:
zhanghuan
@@ -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 许可证。
|
||||
Reference in New Issue
Block a user