EfficientIR - 基于 EfficientNet 的图片检索工具
基于内容的图片检索方法有很多，即便是传统算法也能达到不错的效果。 本项目起源于要求使用传统算法（SIFT）完成的课程作业，但实现过程中遇到了一些影响实际使用的问题（例如建立索引慢）。 经过简单调研（指咨询群友和 Google）后发现 PC 上并没有什么好用的本地图片检索工具，于是干脆在作业外另外弄了这个小工具。
使用到的主要工具是 EfficientNet 和 Hnswlib，使用前者在 ImageNet 上的预训练模型进行特征抽取，使用后者进行特征索引及检索。
使用前请注意，如果有任何使用上的问题请尽量自行解决。
比起花费时间维护，咱更希望像 digiKam/Eagle 这样的专业图像管理软件加入特性持续维护。
如果您有任何新点子，欢迎 fork 实现后 pull request。
功能特性
- 以图搜图
- 图片相似度计算
- 重复图片查找
依赖
- Python 3.6+
- numpy
- onnxruntime
- hnswlib
- pillow
- tqdm
- onnx (若需更换模型则必要)
- PyQt5
使用
GUI 使用
当前 GUI 版本打包仅支持 Windows x86_64 环境使用
在 Release 页面下载最新版程序打包，并解压到不包含中日韩文字的路径下。
初次使用需要建立索引：
- 双击
start.bat打开 - 单击
设置-添加索引目录添加需要索引的图库目录 - 单击
更新索引目录建立索引
搜索图片：
- 单击
...选择图片 - 设置返回结果数量
- 单击
开始搜索 - 双击返回结果的文件路径打开图片
图片查重：
- 单击
查重切换到查重功能 - 设置过滤结果使用到的相似度阈值（最低为 70%）
- 单击
开始查重，并坐和放宽 - 双击返回结果的文件路径打开图片
后续索引更新：
- 单击
设置-更新索引目录更新索引
更换模型
当前本 repo 已包含以下模型：
- EfficientNet-B2
- Once For All (flops@[email protected]_finetune@75)
关于将 PyTorch 或其他机器学习框架模型导出为 ONNX 模型的方法此处不再赘述。需要注意的是，导出的 ONNX 模型必须经过优化过程，可以使用 onnx-simplifier （推荐）或本项目包含的 opti.py。
更换模型前可以使用 Netron 检查模型的输入矩阵形状是否为 1x3xNxN,输出向量是否为 1xN。其中输入矩阵的形状对应 efficient_ir.py 中的常量 img_size，输出向量的形状对应 init_index() 方法的初始化维度 以及 get_fv() 方法的返回值。
若只是希望更换模型为更大的 EfficientNet 模型，那么只需要确认并修改 efficient_ir.py 中的 img_size 和 model_path。
但如果需要更换为 Once For All 模型，虽然其输入与 EfficientNet-B2 相同，但输出是 N 并不是 1xN，故除修改 model_path 外，还需将 get_fv() 中返回所在行做出如下修改:
修改前：
return self.session.run([], {self.model_input: norm_img_data})[0][0]
修改后
return self.session.run([], {self.model_input: norm_img_data})[0]
注意：更换模型后一定要重新建立索引
GPU 加速
当前选择的模型均对性能进行了权衡，在支持 AVX 指令集的 CPU 上索引速度略高于甜品级 GPU。
若期待通过 GPU 加速获得更好的性能及加速比，可以将模型换成更大规模的或增加索引时同时处理的图片数量。
具体的操作请自行阅读并修改代码实现。
如果是 NVIDIA 显卡，切换 GPU 推理的步骤：
- 安装
onnxruntime-gpu； - 取消
efficient_ir.py第 61 行的注释； - 将 provider 需要换成
GPUExecutionProvider。
支持 DX12 Compute 的任意显卡（包括集成显卡），切换 GPU 推理的步骤：
- 安装
onnxruntime-dml； - 取消
efficient_ir.py第 61 行的注释。
Q&A
Q：可承载最大索引数量是多少？如何修改？
A：目前是 1000000。可以在efficient_ir.py中修改，数值将在下一次加载时生效。
Q：检索效果不佳怎么解决？
A：当前代码中使用 EfficientNet-b2 模型是经过权衡后决定的，若追求更佳检索效果请自行更换更大规模的 EfficientNet 模型或其他的 SOTA 模型。本项目将持续关注 SOTA 模型的发展，并在 Wiki 中更新相关测试结果。
TODO
- 实现 GUI 工具
- 移植到 C++ 使用
- 替换 digiKam 内的模糊搜索