Pytorch比较图片间的相似度。

import torch
from torchvision import models, transforms
from PIL import Image
import numpy as np

# 加载预训练的模型
model = models.resnet50(pretrained=True)
model.eval()

# 图片预处理和特征提取函数
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 存储特征向量的字典
feature_vectors = {}

# 遍历每个文件夹
for folder_name in folder_names:
    folder_path = 'path_to_folder/' + folder_name  # 文件夹路径
    images = os.listdir(folder_path)  # 文件夹下的所有图片文件名
    for image_name in images:
        # 读取图片并进行预处理
        image_path = os.path.join(folder_path, image_name)
        img = Image.open(image_path)
        img = preprocess(img)
        img = img.unsqueeze(0)

        # 提取特征向量
        with torch.no_grad():
            features = model(img)
        
        # 存储特征向量
        feature_vectors[image_name] = features.numpy()

# 计算图片之间的相似度
similar_images = []
for image1_name in feature_vectors:
    for image2_name in feature_vectors:
        if image1_name != image2_name:
            # 计算相似度得分，可以使用不同的度量方法
            similarity_score = cosine_similarity(feature_vectors[image1_name], feature_vectors[image2_name])
            # 如果得分高于阈值，将相似的图片添加到列表中
            if similarity_score > threshold:
                similar_images.append((image1_name, image2_name, similarity_score))

# 根据相似度得分排序
similar_images.sort(key=lambda x: x[2], reverse=True)

# 输出相似度高的图片对
for image_pair in similar_images:
    image1_name, image2_name, similarity_score = image_pair
    print(f"Image pair: {image1_name}, {image2_name}, Similarity score: {similarity_score}")