【卷积神经网络+欧氏距离实现对图片相似度评估】

        基于Pytorch深度学习框架卷积神经网络配合欧式距离，实现对图片之间相似度的评估，这只是一个简单demo，可根据自己需求修改特征提取器和图片与处理模块。

        

"""
@author: jialin_li
@data: 20231008
@version: v1.0

"""


import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
from scipy.spatial.distance import euclidean

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#加载resnet50预训练模型用作特征提取器
model = models.resnet50(pretrained=True)
#去掉全链接层
model = nn.Sequential(*list(model.children())[:-1])
model.eval()
model.to(device)


def extract_features(image_path):
    image = Image.open(image_path)
    #这里根据自己需求对图片预处理 我这里默认用的图片原尺寸，但需要图片尺寸一致
    preprocess = transforms.Compose([
        # transforms.Resize(256),
        # transforms.CenterCrop(1200),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    image = preprocess(image).to(device)
    #给图片增加一个维度
    image = torch.unsqueeze(image, 0)
    
    #提取特征
    with torch.no_grad():
        features = model(image)

    features = torch.flatten(features)
    
    return features.cpu().numpy()


def calculate_euclidean_distance(feature1, feature2):

    #调用欧式距离函数进行评估
    return euclidean(feature1, feature2)


if __name__ == '__main__':
    
    #图片路径
    image1_path = 'daozi2.jpg'
    image2_path = 'daozi4.jpg'


    feature1 = extract_features(image1_path)
    feature2 = extract_features(image2_path)
    distance = calculate_euclidean_distance(feature1, feature2)
    
    #距离越小那么证明相似度越高
    print(f"图像相似度距离：{distance}")