Views
No views yet
transformers library:1from transformers import ViTForImageClassification, ViTImageProcessor
2from PIL import Image
3import torch
4
5# Define the device
6device = "cuda" if torch.cuda.is_available() else "cpu"
7
8# Load the model and image processor
9model_id = "skshmjn/Pokemon-classifier-gen9-1025"
10model = ViTForImageClassification.from_pretrained(model_id).to(device)
11image_processor = ViTImageProcessor.from_pretrained(model_id)
12
13# Load and process an image
14img = Image.open('test.jpg').convert("RGB")
15inputs = image_processor(images=img, return_tensors='pt').to(device)
16
17# Make predictions
18outputs = model(**inputs)
19predicted_id = outputs.logits.argmax(-1).item()
20predicted_pokemon = model.config.id2label[predicted_id]
21
22# Print predicted class
23print(f"Predicted Pokémon Pokédex number: {predicted_id+1}")
24print(f"Predicted Pokémon: {predicted_pokemon}")