一、DINOv2 模型简介及使用DINOv2是由Meta AI开发的第二代自监督视觉变换器模型采用Vision Transformer (ViT)架构 。其核心特点是在无需人工标签的情况下通过自监督学习技术从海量无标注图像中学习有意义的视觉特征表示类似于NLP领域的自监督Base模型DINOv2已经具有了对图像的理解能力和强大的图像特征提取能力因此它可以作为几乎所有计算机视觉任务的骨干模型。下面是官方演示地址https://dinov2.metademolab.com/demos?categorysegmentation深度估计效果语义分割效果GitHub开源地址https://github.com/facebookresearch/dinov2huggingface模型地址https://huggingface.co/facebook/dinov2-base/tree/main本文借助DINOv2强大的特征提取能力实现图图相似度检索任务但开始前首先 先了解一下如何基于DINOv2实现图像的相似度计算fromtransformersimportAutoImageProcessor,AutoModelfromPILimportImageimportmatplotlib.pyplotaspltimporttorch plt.rcParams[font.sans-serif][SimHei]# 生成图像特征defgen_image_features(processor,model,device,image):withtorch.no_grad():inputsprocessor(imagesimage,return_tensorspt).to(device)outputsmodel(**inputs)image_featuresoutputs.last_hidden_state image_featuresimage_features.mean(dim1)returnimage_features[0]# 计算两个图像的相似度defsimilarity_image(processor,model,device,image1,image2):features1gen_image_features(processor,model,device,image1)features2gen_image_features(processor,model,device,image2)cos_simtorch.cosine_similarity(features1,features2,dim0)cos_sim(cos_sim1)/2returncos_sim.item()defmain():model_dirfacebook/dinov2-baseprocessorAutoImageProcessor.from_pretrained(model_dir)modelAutoModel.from_pretrained(model_dir)devicetorch.device(cuda:0iftorch.cuda.is_available()elsecpu)model.to(device)image1Image.open(img/dog1.jpg)image2Image.open(img/dog2.jpg)similaritysimilarity_image(processor,model,device,image1,image2)plt.figure()plt.axis(off)plt.title(f相似度:{similarity})plt.subplot(1,2,1)plt.imshow(image1)plt.subplot(1,2,2)plt.imshow(image2)plt.show()if__name____main__:main()相似度计算的主要核心就是基于 DINOv2 生成特征向量图图相似度检索也依赖这一点不过需要多出一个特征向量的持久化存储端整体实现架构如下图所示其中特征向量存储采用Milvus数据库。被检索图像数据集特征提取过程图像相似度检索过程关于Milvus的使用可以参考下面这篇博客Milvus 向量数据库介绍及使用二、图图相似度检索实现 - 图像特征持久化首先准备图像数据集这里我随便准备了几张猫和狗的图片创建Milvus Collection其中DINOv2特征向量维度为768维frompymilvusimportMilvusClient,DataType clientMilvusClient(http://192.168.0.5:19530)schemaMilvusClient.create_schema(auto_idTrue,enable_dynamic_fieldFalse,)schema.add_field(field_nameid,datatypeDataType.INT64,is_primaryTrue)schema.add_field(field_namevector,datatypeDataType.FLOAT_VECTOR,dim768)schema.add_field(field_nameimage_name,datatypeDataType.VARCHAR,max_length256)schema.verify()index_paramsclient.prepare_index_params()index_params.add_index(field_nameid,index_typeSTL_SORT)index_params.add_index(field_namevector,index_typeIVF_FLAT,metric_typeL2,params{nlist:1024})# 创建 collectionclient.create_collection(collection_namedinov2_collection,schemaschema,index_paramsindex_params)图像数据集提取特征向量后持久化到Milvus中fromtransformersimportAutoImageProcessor,AutoModelfromPILimportImagefrompymilvusimportMilvusClientfromtqdmimporttqdmimporttorchimportos# 生成特征向量defgen_image_features(processor,model,device,image):withtorch.no_grad():inputsprocessor(imagesimage,return_tensorspt).to(device)outputsmodel(**inputs)image_featuresoutputs.last_hidden_state image_featuresimage_features.mean(dim1)returnimage_features[0]defmain():# 创建Milvus客户端clientMilvusClient(http://192.168.0.5:19530)# 加载模型model_dirfacebook/dinov2-baseprocessorAutoImageProcessor.from_pretrained(model_dir)modelAutoModel.from_pretrained(model_dir)devicetorch.device(cuda:0iftorch.cuda.is_available()elsecpu)model.to(device)# 读取数据集dataset_path./imgforimage_nameintqdm(os.listdir(dataset_path)):image_pathos.path.join(dataset_path,image_name)imageImage.open(image_path)# 提取特征向量featuresgen_image_features(processor,model,device,image)# 存吃至 milvusclient.insert(collection_namedinov2_collection,data{vector:features,image_name:image_name})if__name____main__:main()运行后在Milvus insight工具中可以看到存储的内容三、图图相似度检索实现 - 图像特征检索检索就是拿着当前图像的特征去向量数据库中检索相似的特征信息实现过程如下fromtransformersimportAutoImageProcessor,AutoModelfromPILimportImagefrompymilvusimportMilvusClientimportmatplotlib.pyplotaspltimporttorchimportos# 生成特征向量defgen_image_features(processor,model,device,image):withtorch.no_grad():inputsprocessor(imagesimage,return_tensorspt).to(device)outputsmodel(**inputs)image_featuresoutputs.last_hidden_state image_featuresimage_features.mean(dim1)returnimage_features[0].tolist()defmain():# 创建Milvus客户端clientMilvusClient(http://192.168.0.5:19530)# 加载模型model_dirfacebook/dinov2-baseprocessorAutoImageProcessor.from_pretrained(model_dir)modelAutoModel.from_pretrained(model_dir)devicetorch.device(cuda:0iftorch.cuda.is_available()elsecpu)model.to(device)# 检索图像, 采用不在Milvus数据集中的图像imageImage.open(./img2/dog.jpeg)# 提取特征向量featuresgen_image_features(processor,model,device,image)# 特征召回resultsclient.search(collection_namedinov2_collection,data[features],limit2,output_fields[image_name],search_params{metric_type:L2,params:{}})plt.figure()plt.axis(off)fori,resinenumerate(results[0]):image_nameres[entity][image_name]image_pathos.path.join(./img,image_name)imageImage.open(image_path)plt.subplot(1,2,(i1))plt.imshow(image)plt.show()if__name____main__:main()测试输入检索图像召回图像测试输入检索图像召回图像