Update README.md
Browse files# Pokemon Team Classification with Vision Transformer
A fine-tuned Vision Transformer (ViT) model for classifying 6 specific Pokemon from a competitive team setup. This model can identify Arceus, Marshadow, Sandy Shocks, Slaking, Reshiram, and Magearna with high accuracy.
## Model Details
- **Base Model**: `google/vit-base-patch16-224`
- **Model Type**: Vision Transformer for Image Classification
- **Classes**: 6 Pokemon (Arceus, Marshadow, Sandy Shocks, Slaking, Reshiram, Magearna)
- **Input Size**: 224x224 RGB images
- **Framework**: PyTorch + Transformers
## Training Details
### Dataset
- **Arceus**: 644 images
- **Marshadow**: 101 images
- **Sandy Shocks**: 75 images
- **Slaking**: 152 images
- **Reshiram**: 118 images
- **Magearna**: 200 images
### Training Strategy
- **Balanced Sampling**: Each epoch uses exactly 75 samples per class to prevent overfitting on Arceus
- **Data Augmentation**: Random horizontal flip, rotation (±15°), color jitter, and resized crop
- **Transfer Learning**: Froze early ViT layers, fine-tuned classifier and later transformer layers
- **Early Stopping**: Training stopped when validation loss plateaued (patience=3 epochs)
### Hyperparameters
- **Learning Rate**: 2e-5
- **Batch Size**: 16
- **Weight Decay**: 0.01
- **Optimizer**: AdamW
- **Epochs**: ~18 (early stopped from max 1000)
## Performance
The model achieves excellent classification performance with balanced accuracy across all 6 Pokemon classes despite the imbalanced training dataset.
## Usage
### Basic Classification
```python
from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import torch
# Load model and processor
model = ViTForImageClassification.from_pretrained("your-username/pokemon-team-vit")
processor = ViTImageProcessor.from_pretrained("your-username/pokemon-team-vit")
# Load and process image
image = Image.open("pokemon_image.jpg")
inputs = processor(images=image, return_tensors="pt")
# Get predictions
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Get results
pokemon_names = ["arceus", "marshadow", "sandy-shocks", "slaking", "reshiram", "magearna"]
predicted_class = predictions.argmax().item()
confidence = predictions.max().item()
print(f"Predicted: {pokemon_names[predicted_class]} (confidence: {confidence:.2%})")
```
### Detailed Probabilities
```python
# Get all class probabilities
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
results = {}
for idx, pokemon in enumerate(pokemon_names):
results[pokemon] = float(probabilities[idx])
# Sort by probability
sorted_results = sorted(results.items(), key=lambda x: x[1], reverse=True)
for pokemon, prob in sorted_results:
print(f"{pokemon}: {prob:.1%}")
```
## Applications
- **Pokemon Recognition**: Identify specific Pokemon in images, artwork, or screenshots
- **Competitive Team Analysis**: Analyze team compositions in competitive Pokemon content
- **Content Moderation**: Filter or categorize Pokemon-related content
- **Educational Tools**: Pokemon identification for learning applications
## Limitations
- **Specific Pokemon Only**: Only recognizes the 6 trained Pokemon classes
- **Image Quality**: Performance may vary with very low resolution or heavily distorted images
- **Artistic Variations**: May struggle with highly stylized or non-canonical Pokemon representations
- **Background Complexity**: Performance may decrease with very cluttered backgrounds
## Model Architecture
The model uses the Vision Transformer (ViT) architecture:
- **Patch Size**: 16x16
- **Hidden Size**: 768
- **Attention Heads**: 12
- **Layers**: 12
- **Parameters**: ~86M (base model) + classification head
## Training Infrastructure
- **Hardware**: AMD GPU with ROCm support
- **Framework**: PyTorch with Transformers library
- **Duration**: ~2 minutes per epoch, early stopped at epoch 18
- **Memory**: Optimized for consumer-grade GPU memory
## Citation
If you use this model, please cite:
```bibtex
@misc {pokemon-team-vit,
title={Pokemon Team Classification with Vision Transformer},
author={Steven Van Ingelgem},
year={2025},
url={https://huggingface.co/your-username/pokemon-team-vit}
}
```
## License
This model is released under the MIT License. The training data consists of Pokemon images which are © The Pokémon Company/Nintendo. This model is for research and educational purposes.
## Acknowledgments
- Base model: Google's Vision Transformer (ViT)
- Training framework: Hugging Face Transformers
- Pokemon images: Various sources for competitive team analysis
---
**Note**: This model is specifically trained for a competitive Pokemon team setup and may not generalize to other Pokemon or use cases. For broader Pokemon classification, consider training on a more comprehensive dataset.