模型描述

这是一个基于 apple/mobilevitv2-1.0-imagenet1k-256 微调的模型，使用 Xenova/quickdraw-small 数据集进行素描图像识别训练。

如何使用？

from transformers import MobileViTImageProcessor, MobileViTV2ForImageClassification
from PIL import Image
import requests
import torch
import numpy as np  # 导入 NumPy

url = \"https://static.thenounproject.com/png/2024184-200.png\"
response = requests.get(url, stream=True)

# 转换为灰度图以确保单通道输入
image = Image.open(response.raw).convert('L')  # 转换为灰度图

processor = MobileViTImageProcessor.from_pretrained(\"laszlokiss27/doodle-dash2\")
model = MobileViTV2ForImageClassification.from_pretrained(\"laszlokiss27/doodle-dash2\")

# 将 PIL 图像转换为张量并添加通道维度
image_tensor = torch.unsqueeze(torch.tensor(np.array(image)), 0).float()
image_tensor = image_tensor.unsqueeze(0)  # 添加批次维度

# 检查处理器是否需要特定形式的输入
inputs = processor(images=image_tensor, return_tensors=\"pt\")

outputs = model(**inputs)
logits = outputs.logits

# 获取预测结果
predicted_class_idx = logits.argmax(-1).item()
predicted_class = model.config.id2label[predicted_class_idx]
print(\"Predicted class:\", predicted_class)