基于 DINOv2 模型实现图搜图相似度检索任务

一、DINOv2 模型简介及利用

DINOv2是由Meta AI开发的第二代自监督视觉变更器模型，采用 Vision Transformer (ViT) 架构 。其核心特点是在无需人工标签的情况下，通过自监督学习技术，从海量无标注图像中学习故意义的视觉特征表现，类似于 NLP 领域的自监督 Base 模型，DINOv2 已经具有了对图像的明白能力，和强大的图像特征提取能力，因此它可以作为险些全部计算机视觉任务的骨干模型。
下面是官方演示地点：
   https://dinov2.metademolab.com/demos?category=segmentation
  深度估计结果：

语义分割结果：

GitHub 开源地点：
   https://github.com/facebookresearch/dinov2
  huggingface 模型地点：
   https://huggingface.co/facebook/dinov2-base/tree/main
  

本文借助 DINOv2 强大的特征提取能力，实现图图相似度检索任务，但开始前，起首 先相识一下如何基于 DINOv2 实现图像的相似度计算：

1. from transformers import AutoImageProcessor, AutoModel
2. from PIL import Image
3. import matplotlib.pyplot as plt
4. import torch
5. plt.rcParams['font.sans-serif'] = ['SimHei']
6. # 生成图像特征
7. def gen_image_features(processor, model, device, image):
8.     with torch.no_grad():
9.         inputs = processor(images=image, return_tensors="pt").to(device)
10.         outputs = model(**inputs)
11.         image_features = outputs.last_hidden_state
12.         image_features = image_features.mean(dim=1)
13.         return image_features[0]
15. # 计算两个图像的相似度
16. def similarity_image(processor, model, device, image1, image2):
17.     features1 = gen_image_features(processor, model, device, image1)
18.     features2 = gen_image_features(processor, model, device, image2)
19.     cos_sim = torch.cosine_similarity(features1, features2, dim=0)
20.     cos_sim = (cos_sim + 1) / 2
21.     return cos_sim.item()
23. def main():
24.     model_dir = "facebook/dinov2-base"
25.     processor = AutoImageProcessor.from_pretrained(model_dir)
26.     model = AutoModel.from_pretrained(model_dir)
27.     device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
28.     model.to(device)
29.     image1 = Image.open("img/dog1.jpg")
30.     image2 = Image.open("img/dog2.jpg")
31.     similarity = similarity_image(processor, model, device, image1, image2)
32.     plt.figure()
33.     plt.axis('off')
34.     plt.title(f"相似度: {similarity}")
35.     plt.subplot(1, 2, 1)
36.     plt.imshow(image1)
37.     plt.subplot(1, 2, 2)
38.     plt.imshow(image2)
39.     plt.show()
41. if __name__ == '__main__':
42.     main()

复制代码

相似度计算的主要核心就是基于 DINOv2 天生特征向量，图图相似度检索也依赖这一点，不过需要多出一个特征向量的持久化存储端，整体实现架构如下图所示，其中特征向量存储采用 Milvus 数据库。
被检索图像数据集特征提取过程：

图像相似度检索过程：

关于 Milvus 的利用，可以参考下面这篇博客：
   Milvus 向量数据库介绍及利用
  二、图图相似度检索实现 - 图像特征持久化

起首准备图像数据集，这里我随便准备了几张猫和狗的图片：

创建 Milvus Collection，其中 DINOv2 特征向量维度为 768 维：

1. from pymilvus import MilvusClient, DataType
3. client = MilvusClient("http://192.168.0.5:19530")
5. schema = MilvusClient.create_schema(
6.     auto_id=True,
7.     enable_dynamic_field=False,
8. )
9. schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
10. schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=768)
11. schema.add_field(field_name="image_name", datatype=DataType.VARCHAR, max_length=256)
12. schema.verify()
13. index_params = client.prepare_index_params()
14. index_params.add_index(
15.     field_name="id",
16.     index_type="STL_SORT"
17. )
18. index_params.add_index(
19.     field_name="vector",
20.     index_type="IVF_FLAT",
21.     metric_type="L2",
22.     params={"nlist": 1024}
23. )
24. # 创建 collection
25. client.create_collection(
26.     collection_name="dinov2_collection",
27.     schema=schema,
28.     index_params=index_params
29. )

复制代码

图像数据集提取特征向量后，持久化到 Milvus 中：

1. from transformers import AutoImageProcessor, AutoModel
2. from PIL import Image
3. from pymilvus import MilvusClient
4. from tqdm import tqdm
5. import torch
6. import os
8. # 生成特征向量
9. def gen_image_features(processor, model, device, image):
10.     with torch.no_grad():
11.         inputs = processor(images=image, return_tensors="pt").to(device)
12.         outputs = model(**inputs)
13.         image_features = outputs.last_hidden_state
14.         image_features = image_features.mean(dim=1)
15.         return image_features[0]
17. def main():
18.     # 创建Milvus客户端
19.     client = MilvusClient("http://192.168.0.5:19530")
20.     # 加载模型
21.     model_dir = "facebook/dinov2-base"
22.     processor = AutoImageProcessor.from_pretrained(model_dir)
23.     model = AutoModel.from_pretrained(model_dir)
24.     device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
25.     model.to(device)
26.     # 读取数据集
27.     dataset_path = "./img"
28.     for image_name in tqdm(os.listdir(dataset_path)):
29.         image_path = os.path.join(dataset_path, image_name)
30.         image = Image.open(image_path)
31.         # 提取特征向量
32.         features = gen_image_features(processor, model, device, image)
33.         # 存吃至 milvus
34.         client.insert(
35.             collection_name="dinov2_collection",
36.             data={
37.                 "vector": features,
38.                 "image_name": image_name
39.             }
40.         )
42. if __name__ == '__main__':
43.     main()

复制代码

运行后，在 Milvus insight 工具中，可以看到存储的内容：

三、图图相似度检索实现 - 图像特征检索

检索就是拿着当前图像的特征，去向量数据库中检索相似的特征信息，实现过程如下：

1. from transformers import AutoImageProcessor, AutoModel
2. from PIL import Image
3. from pymilvus import MilvusClient
4. import matplotlib.pyplot as plt
5. import torch
6. import os
9. # 生成特征向量
10. def gen_image_features(processor, model, device, image):
11.     with torch.no_grad():
12.         inputs = processor(images=image, return_tensors="pt").to(device)
13.         outputs = model(**inputs)
14.         image_features = outputs.last_hidden_state
15.         image_features = image_features.mean(dim=1)
16.         return image_features[0].tolist()
19. def main():
20.     # 创建Milvus客户端
21.     client = MilvusClient("http://192.168.0.5:19530")
22.     # 加载模型
23.     model_dir = "facebook/dinov2-base"
24.     processor = AutoImageProcessor.from_pretrained(model_dir)
25.     model = AutoModel.from_pretrained(model_dir)
26.     device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
27.     model.to(device)
28.     # 检索图像, 采用不在Milvus数据集中的图像
29.     image = Image.open("./img2/dog.jpeg")
30.     # 提取特征向量
31.     features = gen_image_features(processor, model, device, image)
32.     # 特征召回
33.     results = client.search(
34.         collection_name="dinov2_collection",
35.         data=[features],
36.         limit=2,
37.         output_fields=["image_name"],
38.         search_params={
39.             "metric_type": "L2",
40.             "params": {}
41.         }
42.     )
43.     plt.figure()
44.     plt.axis('off')
45.     for i, res in enumerate(results[0]):
46.         image_name = res["entity"]["image_name"]
47.         image_path = os.path.join("./img", image_name)
48.         image = Image.open(image_path)
49.         plt.subplot(1, 2, (i+1))
50.         plt.imshow(image)
51.     plt.show()
53. if __name__ == '__main__':
54.     main()

复制代码

测试输入检索图像：

召回图像：

测试输入检索图像：

召回图像：

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。