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

这篇具有很好参考价值的文章主要介绍了Pytorch基于VGG cosine similarity实现简单的以图搜图(图像检索)。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

代码如下:

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)如下:
Pytorch基于VGG cosine similarity实现简单的以图搜图(图像检索),划水
我们希望找到相册中其他狸花的图像:
Pytorch基于VGG cosine similarity实现简单的以图搜图(图像检索),划水
上述数据集中,编号01到10的为奶牛猫,编号11到20的则为狸花猫。运行代码,结果如下:

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的姿势十分相似,且环境也差不多(地板),这也是另一种层面上的两图像相似。文章来源地址https://www.toymoban.com/news/detail-647422.html

到了这里,关于Pytorch基于VGG cosine similarity实现简单的以图搜图(图像检索)的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

原文地址:https://blog.csdn.net/qq_40714949/article/details/132212179

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包