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.

Downloads last month
1
Safetensors
Model size
85.8M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Space using fusin001/pokemon-vit 1