ViT inference.py

from transformers import ViTFeatureExtractor, ViTForImageClassification
import requests
from io import BytesIO
from PIL import Image

model_name_or_path = 'vit-base-avengers-v1'
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path)
model = ViTForImageClassification.from_pretrained(model_name_or_path)


url = "https://static1.srcdn.com/wordpress/wp-content/uploads/2022/06/Iron-Man-Change-After-The-Avengers-Is-Deeper-Than-You-Thought.jpg"
response = requests.get(url)
img = Image.open(BytesIO(response.content))
display(img) 

inputs = feature_extractor(images=img, return_tensors="pt")
outputs = model(**inputs)

logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
print("Predicted class:", model.config.id2label[predicted_class_idx])