BiernyVR/pokemon-classifier-mobilenetv3

⚡ Real-World Pokémon Classifier (DenseNet121 — Generation 1)

A high-accuracy, lightweight DenseNet121 computer vision model specifically trained to identify Generation 1 Pokémon in real-world environments (physical trading cards, 3D figurines, plush toys, merchandise, anime screenshots, and screens from various camera angles).

Achieves 97.46% Top-1 validation accuracy and 99.42% Top-5 accuracy across all 151 original Pokémon (from Bulbasaur #001 to Mew #151), based on research and methodology from teekaytai/pokedex.

Designed for Edge AI and mobile deployment (30–40ms CPU inference on smartphones, ~28.8 MB file size, 100% offline). Originally built as the native vision intelligence for the Android Pokédex AI application.

---

🌟 Model Highlights

• Real-World Robustness: Trained on over 12,000 real-world images (Kaggle dataset lantian773030/pokemonclassification + HuggingFace JJMack + official 3D Pokémon Home renders) with heavy 3D perspective distortion, random affine rotation (±35°), photometric color jitter, shadows, and coarse dropout (to handle fingers holding cards/toys).
• DenseNet121 Feature Reuse: Employs dense connectivity to preserve fine morphological details (ears, tails, eye shapes, markings) alongside global body structure.
• Dual Formats Provided:
  • pokemon_classifier.onnx — Fully optimized self-contained ONNX model with embedded weights for ONNX Runtime (Android, iOS, Raspberry Pi, Web, C++, Python).
  • pokemon_densenet121_best.pth — PyTorch checkpoint with trained weights.
  • pokemon_labels.json — 151 labeled entries mapping class index to National Pokédex ID and name.
• Background Noise Rejection: Genuine Pokémon produce high logits (> 3.0) and large margins (> 2.0). Empty backgrounds (desks, blank walls, random clutter) produce flat logits (< 1.5, margin < 0.8), enabling clean non-Pokémon rejection.

---

🚀 Quick Start (Inference)

Option 1: Using ONNX Runtime (Recommended, No PyTorch Needed)

pip install onnxruntime pillow numpy huggingface_hub

import json
import numpy as np
from PIL import Image
import onnxruntime as ort
from huggingface_hub import hf_hub_download

# 1. Download model and labels from Hugging Face Hub
model_path = hf_hub_download(repo_id="BiernyVR/pokemon-classifier-mobilenetv3", filename="pokemon_classifier.onnx")
labels_path = hf_hub_download(repo_id="BiernyVR/pokemon-classifier-mobilenetv3", filename="pokemon_labels.json")

with open(labels_path, "r", encoding="utf-8") as f:
    labels = json.load(f)

# 2. Preprocess input image (224x224 RGB, standard ImageNet normalisation)
img = Image.open("your_pokemon_image.jpg").convert("RGB").resize((224, 224), Image.Resampling.BILINEAR)
arr = (np.array(img, dtype=np.float32) / 255.0 - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
tensor = np.expand_dims(np.transpose(arr, (2, 0, 1)), axis=0).astype(np.float32)

# 3. Run Inference
session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
logits = session.run(None, {"input": tensor})[0][0]

# 4. Compute probabilities & Top 5
probs = np.exp((logits - np.max(logits)) / 0.25)
probs /= probs.sum()
top5 = np.argsort(logits)[::-1][:5]

print(f"Top prediction: {labels[top5[0]]['name'].title()} ({probs[top5[0]]*100:.1f}%)")
for rank, idx in enumerate(top5, 1):
    poke = labels[idx]
    print(f"#{rank}: {poke['name'].title()} (ID: #{poke['id']}) - {probs[idx]*100:.1f}%")

Option 2: Standalone CLI

You can also use the included infer.py:

# Using ONNX
python infer.py --image sample_pikachu.png --topk 5

# Using PyTorch
python infer.py --image sample_charizard.png --backend pytorch

---

📊 Training Specifications

• Dataset: 12,009 images (85% Train: 10,275 images / 15% Val: 1,734 images).
• Architecture: DenseNet121 (timm.create_model('densenet121', pretrained=True, num_classes=151, drop_rate=0.25)).
• Augmentation Pipeline:
  • RandomResizedCrop(size=(224, 224), scale=(0.7, 1.0))
  • Affine(scale=(0.8, 1.2), rotate=(-35, 35), translate_percent=(-0.1, 0.1))
  • Perspective(scale=(0.05, 0.15))
  • ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1)
  • GaussianBlur(blur_limit=(3, 5))
  • CoarseDropout(num_holes_range=(1, 8), hole_height_range=(8, 32), hole_width_range=(8, 32))
• Hardware: NVIDIA GeForce RTX 5080 GPU with Automatic Mixed Precision (fp16).
• Validation Metrics:
  • Top-1 Accuracy: 97.46%
  • Top-5 Accuracy: 99.42%

---

⚖️ Disclaimer & Attribution

• Pokémon and Pokémon character names are trademarks and copyright of Nintendo, Creatures Inc., and GAME FREAK Inc.
• This model is a fan-made, non-commercial educational project designed for on-device computer vision research.