基于DINOv2的图像相似性比对(含代码和模型数据)

一、概述

2023年4月份，Meta发布了DINOv2预训练模型。DINOv2属于一种自监督学习算法，相比于以往的自监督技术有了显著改进，并达到了与弱监督算法相当的性能。

DINOv2算法可以在任何图像数据集上进行训练，无需任何关联的元数据，可以视为能够从给定的所有图像中学习，而不仅仅是那些包含特定文本或标题的图像。

这次发布的DINOv2预训练模型是在没有监督的情况下对1.42亿张精心挑选的图像上进行自监督训练得到的，生成的特征无需任何微调即可使用，应用领域包括图像深度估计、语义分割、实例检索等。

考虑到DINOv2强大的图像特征提取能力，下面简单介绍下如何使用DINOv2预训练模型进行图像相似度计算。本文测试平台为Ubuntu20.04。

二、实践

2.1 安装环境

DINOv2是使用pytorch进行训练的，为了能够直接使用DINOv2预训练模型，需要先安装Pytorch。具体安装方法参考Pytorch官网。

接下来安装DINOv2相关依赖：

pip install transformers -i https://mirror.baidu.com/pypi/simple

最后下载dinov2提供的预训练模型，需要自行前往HuggingFace网站上查找dinov2模型进行下载。本文使用dinov2-base模型，对应下载下面三个文件。

下载完成后将上述三个文件放置在本地一个名为dinov2_base的文件夹中即可。

如果因网络原因无法下载，本文文末提供了完整版代码和模型链接，可以直接下载使用。

2.2 示例代码

代码如下：

from transformers import AutoImageProcessor, AutoModel
from PIL import Image
import torch.nn as nn
import torch

# 定义环境
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")

# 加载dinov2模型
model_folder = './dinov2_base'
processor = AutoImageProcessor.from_pretrained(model_folder)
model = AutoModel.from_pretrained(model_folder).to(device)

# 提取第1张图片特征
image1 = Image.open('img1.jpg')
with torch.no_grad():
    inputs1 = processor(images=image1, return_tensors="pt").to(device)
    outputs1 = model(**inputs1)
    image_features1 = outputs1.last_hidden_state
    image_features1 = image_features1.mean(dim=1)

# 提取第2张图片特征
image2 = Image.open('img2.jpg')
with torch.no_grad():
    inputs2 = processor(images=image2, return_tensors="pt").to(device)
    outputs2 = model(**inputs2)
    image_features2 = outputs2.last_hidden_state
    image_features2 = image_features2.mean(dim=1)
    
# 提取第3张图片特征
image3 = Image.open('img3.jpg')
with torch.no_grad():
    inputs3 = processor(images=image3, return_tensors="pt").to(device)
    outputs3 = model(**inputs3)
    image_features3 = outputs3.last_hidden_state
    image_features3 = image_features3.mean(dim=1)

# 计算相似度
cos = nn.CosineSimilarity(dim=0)
sim12 = cos(image_features1[0],image_features2[0]).item()
sim12 = (sim12+1)/2
print('img1和img2的相似度值: ', sim12)

sim13 = cos(image_features1[0],image_features3[0]).item()
sim13 = (sim13+1)/2
print('img1和img3的相似度值: ', sim13)

上述代码分别对三张图片使用dinov2模型进行了特征向量提取（768维），最后使用cosine函数来计算特征向量之间的相似度。三张图像如下：

img1和img2的相似度值:  0.8779885470867157
img1和img3的相似度值:  0.6644228100776672

本文完整代码+模型+测试图片的下载地址