Views
No views yet
google/vit-base-patch16-2241from transformers import ViTImageProcessor, ViTForImageClassification
2from PIL import Image
3import torch
4
5# Load model and processor
6model = ViTForImageClassification.from_pretrained("your-username/pokemon-team-vit")
7processor = ViTImageProcessor.from_pretrained("your-username/pokemon-team-vit")
8
9# Load and process image
10image = Image.open("pokemon_image.jpg")
11inputs = processor(images=image, return_tensors="pt")
12
13# Get predictions
14with torch.no_grad():
15 outputs = model(**inputs)
16 predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
17
18# Get results
19pokemon_names = ["arceus", "marshadow", "sandy-shocks", "slaking", "reshiram", "magearna"]
20predicted_class = predictions.argmax().item()
21confidence = predictions.max().item()
22
23print(f"Predicted: {pokemon_names[predicted_class]} (confidence: {confidence:.2%})")1# Get all class probabilities
2probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
3
4results = {}
5for idx, pokemon in enumerate(pokemon_names):
6 results[pokemon] = float(probabilities[idx])
7
8# Sort by probability
9sorted_results = sorted(results.items(), key=lambda x: x[1], reverse=True)
10for pokemon, prob in sorted_results:
11 print(f"{pokemon}: {prob:.1%}")1@misc{pokemon-team-vit,
2 title={Pokemon Team Classification with Vision Transformer},
3 author={Steven Van Ingelgem},
4 year={2025},
5 url={https://huggingface.co/your-username/pokemon-team-vit}
6}