思路:使用模型的特征提取层,转化成向量,然后比对向量的距离(比如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()}")