Pytorch基于VGG cosine similarity实现简单的以图搜图(图像检索)

代码如下：

python 复制代码

from PIL import Image
from torchvision import transforms
import os
import torch
import torchvision
import torch.nn.functional as F

class VGGSim(torch.nn.Module):
    def __init__(self):
        super(VGGSim, self).__init__()
        blocks = []
        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())
        for bl in blocks:
            for p in bl:
                p.requires_grad = False
        self.blocks = torch.nn.ModuleList(blocks)
        self.transform = torch.nn.functional.interpolate
        self.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1))
        self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1))

    def forward(self, input, target):
        if input.shape[1] != 3:
            input = input.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        input = (input-self.mean) / self.std
        target = (target-self.mean) / self.std
        x = input
        y = target

        res = []
        for block in self.blocks:
            x = block(x)
            y = block(y)
            x_flat = torch.flatten(x, start_dim=1)
            y_flat = torch.flatten(y, start_dim=1)
            similarity = torch.nn.functional.cosine_similarity(x_flat, y_flat)
            res.append(similarity.cpu().item())
        # 仅利用VGG最后一层的全局(分类)特征计算余弦相似度
        # return res[-1]
        # 或者，利用VGG各Block的特征计算余弦相似度
        return sum(res)

def load_image(path):
    image = Image.open(path).convert('RGB')
    image = transforms.Resize([224,224])(image)
    image = transforms.ToTensor()(image)
    image = image.unsqueeze(0)
    return image.cuda()

query_image_path = "query.jpeg"  # 想要查找的图像
query_image = load_image(query_image_path) 
target_image_dir = "cat_images/" # 待搜索的相册
target_images = [os.path.join(target_image_dir, name) for name in os.listdir(target_image_dir)]
vgg_sim = VGGSim().cuda()
scores = []
for path in target_images:
    target_image = load_image(path)
    score = vgg_sim(query_image, target_image)
    scores.append([path, score])
scores.sort(key=lambda x: -x[1])
for i in range(5):
    print("Top", (i + 1), "similiar =>", scores[i][0].split("/")[-1])

上述代码的核心思想类似于感知损失(Perceptual Loss)，利用VGG提取图像的多级特征，从而比较两张图像之间的相似性。区别在于Perceptual Loss中一般使用MAE，MSE比较特征的距离，而这里的代码使用余弦相似度。

一个例子如下，给定一张狸花的图像(query)如下：

我们希望找到相册中其他狸花的图像：

上述数据集中，编号01到10的为奶牛猫，编号11到20的则为狸花猫。运行代码，结果如下：

python 复制代码

Top 1 similiar => 04.jpeg
Top 2 similiar => 20.jpeg
Top 3 similiar => 14.jpeg
Top 4 similiar => 12.jpeg
Top 5 similiar => 15.jpeg

可以看到，检索基本是正确的，20，14，12，15均为狸花猫。04得到最高相似度的原因是其与query的姿势十分相似，且环境也差不多(地板)，这也是另一种层面上的两图像相似。