在本教程中，将探讨如何利用DINOv2嵌入和C-支持向量分类(SVC)线性分类模型来对图像进行分类。教程结束时，将在MIT室内场景识别数据集上训练一个模型，该模型能够对图像中的场景进行分类。
本教程附带一个交互式笔记本，可以用来跟随指南。建议在GPU上运行DINOv2，因为GPU可以显著加速训练过程。如果没有GPU，可以在Google Colab上复制笔记本或创建一个新的。Google Colab为笔记本提供了GPU资源。
在开始构建分类模型之前，需要将一些依赖项导入到项目中。如果还没有安装numpy、opencv-python、scikit-learn、TQDM和PyTorch，请使用以下命令安装它们：
pip install torch numpy opencv-python scikit-learn
还需要安装roboflow pip包，将使用它来加载项目的数据：
pip install roboflow
接下来，让将所需的依赖项导入到项目中：
import numpy as np
import torch
import torchvision.transforms as T
from PIL import Image
import os
import cv2
import json
import glob
from tqdm.notebook import tqdm
安装了项目所需的依赖项后，可以开始下载项目的数据了。在本指南中，将使用MIT室内场景识别数据集（尽管可以在项目中使用任何想要的数据集。这个数据集的一个版本可以在Roboflow Universe上找到，这是一个在线存储库，包含超过200,000个计算机视觉数据集。需要一个免费的Roboflow账户才能从Roboflow Universe下载数据集。
创建一个新的Python文件，用于本项目，然后添加以下代码：
from roboflow import Roboflow
rf = Roboflow.login()
project = rf.workspace("popular-benchmarks").project("mit-indoor-scene-recognition")
dataset = project.version(5).download("folder")
已经以分类文件夹格式下载了数据集。在这种格式中，每个图像都在一个文件夹中，文件夹的名称等于与图像相关的标签。例如，训练数据集中的图像train/x/y.jpg的标签是x。
需要创建一个字典，将所有文件名映射到它们所在的文件夹名称，以便知道每个图像的标签。可以使用以下代码来实现：
cwd = os.getcwd()
ROOT_DIR = os.path.join(cwd, "MIT-Indoor-Scene-Recognition-5/train")
labels = {}
for folder in os.listdir(ROOT_DIR):
for file in os.listdir(os.path.join(ROOT_DIR, folder)):
if file.endswith(".jpg"):
full_name = os.path.join(ROOT_DIR, folder, file)
labels[full_name] = folder
files = list(labels.keys())
将使用线性分类模型SVC来对图像进行分类。在能够训练模型之前，需要为模型准备输入。对于这个项目，需要两个数据点：每个图像的DINOv2嵌入，以及与每个图像相关的标签。在这一步中，将为项目中的每个图像计算嵌入。
首先，让加载DINOv2模型。将使用最小的模型dino_vits14（'s'表示小）。这个模型比默认的CLIP权重小3倍以上（~84 MB vs. ~300 MB）。
dinov2_vits14 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14")
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
dinov2_vits14.to(device)
transform_image = T.Compose([T.ToTensor(), T.Resize(244), T.CenterCrop(224), T.Normalize([0.5], [0.5])])
在上述代码的最后一行，定义了一个函数，将图像转换为DINOv2接受的格式。接下来，可以编写函数来加载图像，并为图像列表中的每个图像计算嵌入：
def load_image(img: str) -> torch.Tensor:
"""
加载图像并返回一个张量，该张量可以作为DINOv2的输入。
"""
img = Image.open(img)
transformed_img = transform_image(img)[:3].unsqueeze(0)
return transformed_img
def compute_embeddings(files: list) -> dict:
"""
创建一个索引，其中包含指定文件列表中的所有图像。
"""
all_embeddings = {}
with torch.no_grad():
for i, file in enumerate(tqdm(files)):
embeddings = dinov2_vits14(load_image(file).to(device))
all_embeddings[file] = np.array(embeddings[0].cpu().numpy()).reshape(1, -1).tolist()
with open("all_embeddings.json", "w") as f:
f.write(json.dumps(all_embeddings))
return all_embeddings
有了这些函数，可以开始为训练数据集中的图像计算嵌入了。为此，可以通过在教程中定义的文件列表通过compute_embeddings()函数。
embeddings = compute_embeddings(files)
这段代码可能需要几分钟才能运行完毕，具体取决于数据集大小。
已经准备好使用嵌入和标签来拟合分类模型了：
from sklearn import svm
clf = svm.SVC(gamma='scale')
y = [labels[file] for file in files]
print(len(embeddings.values()))
embedding_list = list(embeddings.values())
clf.fit(np.array(embedding_list).reshape(-1, 384), y)
运行此代码后，得到了一个SVC模型，可以在该模型上运行推理。要使用该模型，需要：
对于测试，可以使用数据集中的测试集或验证集中的图像。以下示例中，将对一张电梯图像进行推理：
input_file = "MIT-Indoor-Scene-Recognition-5/test/elevator/elevator_google_0053_jpg.rf.41487c3b9c1690a5de26ee0218452627.jpg"
new_image = load_image(input_file)
with torch.no_grad():
embedding = dinov2_vits14(new_image.to(device))
prediction = clf.predict(np.array(embedding[0].cpu()).reshape(1, -1))
print("Predicted class: " + prediction[0])
此代码将返回以下输出：
Predicted class: elevator
代码已成功对图像进行了分类！Roboflow团队在MIT室内场景识别上对SVC性能进行了基准测试，达到了0.884的准确率（与CLIP的0.883准确率相比）。
在本指南中，已经展示了如何使用DINOv2嵌入构建分类模型。从Roboflow Universe加载了MIT室内场景识别数据集，为数据集中的每个图像计算了嵌入，然后在嵌入上训练了一个SVC模型。