使用pytorch中的resnet预训练模型进行特征提取，以及查找相似图像

本文链接：https://blog.csdn.net/qq_30125323/article/details/126020275

想通过缩略图找原图？之前p过的图像想找原图？汇报时使用压缩过的图像现在想找原图？如何从大量图像文件中快速找到与目标图像相似的那个？pytorch 只需要几行代码就可以搞定。

模型的选取

一般进行特征提取使用图像分类网络即可。参看上一篇使用pytorch中的resnet预训练模型进行快速图像分类。

代码如下

提取查询图像和候选图像的特征，计算二者的余弦相似度，相似度越大则图像越相似。输出图像的路径，将相似的图像保存到指定目录下。

# load model
import torch
import torchvision
model = torchvision.models.resnet101(pretrained=True)
# or any of these variants
# resnet18, resnet34, resnet50, resnet101, resnet152
model.eval()


from PIL import Image
from torchvision import transforms
from tqdm import tqdm

# args
path_to_query = 'target.jpg' # given one query image
path_to_data = '/path/to/data/' # gallery images
BATCH_SIZE = 256
target_dir = 'targetdir' # we can save similar images in target directory
threshold = 0.8 # pick out if the cosine similarity > threshold.

# build dataloader
preprocess = transforms.Compose([
    transforms