深度学习模型判断图片相似度

思路:使用模型的特征提取层,转化成向量,然后比对向量的距离(比如cos)

import torch
import torchvision.models as models
import torchvision.transforms as transforms

from PIL import Image

# load model & feature only
model = models.mobilenet_v3_small(pretrained=True)
model = model.features
model.eval()

# pre process
preprocess = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def image_to_feature_vector(image_path):
    img = Image.open(image_path)
    img_t = preprocess(img)
    batch_t = torch.unsqueeze(img_t, 0)
    with torch.no_grad():
        # get feature vector
        output = model(batch_t)
    # flatten to 1 dims
    return torch.flatten(output, 1)

def compare_images(img_path1, img_path2):
    vector1 = image_to_feature_vector(img_path1)
    vector2 = image_to_feature_vector(img_path2)
    
    cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
    cos_sim = cos(vector1, vector2)
    print(f"Cosine Similarity: {cos_sim.item()}")

 

Leave a Reply

Your email address will not be published. Required fields are marked *