adding model

app.py

@@ -4,6 +4,10 @@ from PIL import Image
 from collections import OrderedDict
 from random import sample
 import csv
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+
+extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
 
 
 title="ImageNet Roulette"
@@ -24,7 +28,12 @@ with open('image_labels.csv', 'r') as csv_file:
     for row in reader:
         imagedict[row['image_name']] = row['image_label']
 
-
+def model_classify(im):
+    inputs = feature_extractor(images=im, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    return("Predicted class:", model.config.id2label[predicted_class_idx])
 
 def check_answer(im):