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