DINOv2 基于自定义数据集进行图片分类预测

idakelly

已于 2024-03-28 12:59:35 修改

阅读量5.3k

点赞数 7

分类专栏：图像处理文章标签：分类数据挖掘人工智能

于 2024-03-28 12:57:57 首次发布

本文链接：https://blog.csdn.net/sinat_34770838/article/details/137021023

版权

本文详细介绍了如何在Python中利用DINOv2模型（如dinov2_vits14）下载预训练权重并提取图片的embedding特征，包括了两种方法：一是直接使用HuggingFace库，二是通过修改DINOV2分类模型进行操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

目标：获取图片的embedding

一、方法一：下载DINOv2模型

（前提是安装好pytorch）

1. pip安装

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

2. 下载dinov2提供的预训练模型，需要自行前往HuggingFace网站上查找dinov2模型进行下载。模型需要下载三个文件：

config.json

preprocessor_config.json

pytorch_model.bin

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

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)

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