【图像搜索】基于pytorch官方预训练模型的简易相似图片搜索

简易的相似图像搜索算法

图片数据库

查询结果

本文主要方法

流程

1. 预训练模型 + 单张图像的特征 （ 逐个保存，形成数据库）+ 特征压缩（选做）

2. 获取查询图像的特征向量

3. 将查询的特征向量与数据库保存的所有特征进行余弦距离计算

4. 返回结果

实际

1. 编写自定义图片数据集读取代码
2. pytorch SWAV预训练模型 （paper Unsupervised Learning of Visual Features by Contrasting Cluster Assignments.）
3. 抽取数据集每一张图片，模型推理得到，4096维度的特征向量 ，保存每一个特征向量
4. 得到将需要查询的图片的4096维度的向量
5. 计算查询向量与所有其他图片的余弦距离，并返回距离最近的topk个图片，完成查询

主要参考：

基于论文复杂结构_搜索算法（牛津数据集）

• PyTorch+flask演示

• End-to-end Learning of Deep Visual Representations for Image Retrieval

• 分别学习相似图像与不相似图像的特征：

• https://github.com/keshik6/deep-image-retrieval#pytorch-source-code

基于自编码机(AE)_搜索算法（以cafir-10为例）：

• PyTorch

• 有效果图

• https://blog.csdn.net/weixin_43786143/article/details/116137867

基于成对相似度数据_搜索算法（以cafir-10为例）

• PaddlePaddle(百度飞浆)
• https://www.paddlepaddle.org.cn/documentation/docs/zh/tutorial/cv_case/image_search/image_search.html

基于预训练模型_搜索算法（任一小数据集）

• Keras
• 基于vgg16预训练模型
• github 433 star:https://github.com/willard-yuan/flask-keras-cnn-image-retrieval
• 问题：图片较多时，无法直接使用该项目

Hash图像_特征的获取(多种hash算法)

• PyTorch
• 提取图像传统特征，并转换为hash编码
• https://github.com/JohannesBuchner/imagehash

主要代码

• https://gitee.com/zengxy2020/csdn_image_retreve

测试数据集导入

python test_read_img.py


保存数据集中的图片到文件夹

修改文件夹路径，以及图片后缀，运行：

python save_fearures_2_npy.py


搜索相同、相似图片

python torch_pretrain_swav_search_one_image.py