6.5 KiB
6.5 KiB
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 # 特征向量数组
安装依赖
pip install torch torchvision
pip install faiss-cpu # 或 faiss-gpu（如果有GPU）
pip install pillow numpy tqdm matplotlib
使用方法
1. 构建FAISS索引
首次使用需要构建FAISS索引：
python build_faiss_index.py
这将：
- 加载训练好的ResNet50 Embedding模型
- 遍历训练集所有图片，提取512维特征向量
- 构建FAISS索引文件
- 保存标签映射和元数据
2. 相似度检索
查看索引信息
python faiss_similarity_search.py --info
单张图片相似度检索
# 基本搜索
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
批量测试
# 批量测试测试集
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索引的核心类：
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
负责相似度检索的核心类：
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 和相关模型文件，修改路径配置即可。
扩展功能
系统支持以下扩展：
- 多模态检索: 结合文本特征进行检索
- 在线更新: 支持动态添加新图片到索引
- 分布式部署: 支持多机部署和负载均衡
- Web接口: 提供REST API接口
- 移动端集成: 支持移动端实时检索
技术细节
- 特征提取: ResNet50 + 全连接层 → 512维向量
- 相似度计算: 余弦相似度（内积，向量已归一化）
- 索引结构: FAISS多种索引类型支持
- 数据格式: 支持 JPG、PNG、BMP 等常见图片格式
- 预处理: ImageNet标准化，224x224分辨率
许可证
本项目遵循 MIT 许可证。