🚀 宝可梦类型分类模型卡

本模型利用视觉变换器（ViT）将宝可梦图像分类为 18 种不同类型。它是 CS 310 期末项目的一部分，在宝可梦图像数据集上进行训练。

🚀 快速开始

以下是如何使用该模型进行分类的步骤：

import torch
from PIL import Image
import torchvision.transforms as transforms
from transformers import ViTForImageClassification, ViTFeatureExtractor

# 加载预训练模型和特征提取器
hf_model = ViTForImageClassification.from_pretrained("NP-NP/pokemon_model")
hf_feature_extractor = ViTFeatureExtractor.from_pretrained("NP-NP/pokemon_model")

# 定义预处理转换
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=hf_feature_extractor.image_mean, std=hf_feature_extractor.image_std)
])

# 标签与索引的映射
labels_dict = {
    'Grass': 0, 'Fire': 1, 'Water': 2, 'Bug': 3, 'Normal': 4, 'Poison': 5, 'Electric': 6,
    'Ground': 7, 'Fairy': 8, 'Fighting': 9, 'Psychic': 10, 'Rock': 11, 'Ghost': 12,
    'Ice': 13, 'Dragon': 14, 'Dark': 15, 'Steel': 16, 'Flying': 17
}
idx_to_label = {v: k for k, v in labels_dict.items()}

# 加载并预处理图像
image_path = "cute-pikachu-flowers-pokemon-desktop-wallpaper.jpg"
image = Image.open(image_path).convert("RGB")
input_tensor = transform(image).unsqueeze(0)  # 形状: (1, 3, 224, 224)

# 进行预测
hf_model.eval()
with torch.no_grad():
    outputs = hf_model(input_tensor)
    logits = outputs.logits
    predicted_class_idx = torch.argmax(logits, dim=1).item()

predicted_class = idx_to_label[predicted_class_idx]
print("预测的宝可梦类型:", predicted_class)

📚 详细文档

模型详情