---
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.