pokemon-vit

Pokemon Vision Transformer Classifier

This repository contains a fine-tuned Vision Transformer for custom Pokémon image classification.

Dataset

The model was trained on the custom Pokémon dataset from week 8 with six classes:

• charizard
• charmander
• charmeleon
• ditto
• eevee
• ekans

Preprocessing

• images are converted to RGB
• images are resized to 224 x 224
• pixel values are normalized using the ViT image processor
• the training split is separated into training and validation subsets
• labels are mapped with label2id and id2label

Model

• base model: google/vit-base-patch16-224
• fine-tuning strategy: the encoder is frozen and only the classification head is trained
• task: image classification

Evaluation

The training script saves evaluation_metrics.json after training with validation and test accuracy values.

Measured results from the retraining run:

• validation accuracy: 0.4737
• test accuracy: 0.6400

Usage

from transformers import pipeline
from PIL import Image

classifier = pipeline("image-classification", model="fusin001/pokemon-vit")
image = Image.open("your_image.png")
print(classifier(image, top_k=3))

Training Script

The repository also includes train_pokemon_model.py, which can be used to retrain and export the model locally.