コードは以下のように表示されます:
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])
上記のコードの中心的なアイデアは、VGG を使用して画像のマルチレベルの特徴を抽出し、2 つの画像間の類似性を比較する知覚損失に似ています。違いは、MAE は一般に知覚損失で使用され、MSE は特徴の距離を比較するのに対し、ここでのコードはコサイン類似度を使用することです。
次のようなアライグマの画像 (クエリ) が与えられた場合の例は次のとおりです。
アルバム内の他のアライグマの画像を見つけたいと考えています。
上記のデータセットでは、番号 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 はすべて cihuamao であることがわかります。04 が最も高い類似度を取得する理由は、クエリのポーズに非常に類似しており、環境 (床) も類似しており、別のレベルの 2 つの画像とも類似しているためです。