File size: 5,012 Bytes
d8098cb 85eb922 d8098cb f02678a d8098cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | ---
license: mit
base_model:
- google/vit-large-patch16-224
pipeline_tag: image-classification
datasets:
- N-o-1/pokemon-images
---
# 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. |