new version!

app.py

@@ -5,8 +5,9 @@ from collections import OrderedDict
 from random import sample
 import csv
 from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import random
 
-extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
+feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
 model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
 
 
@@ -18,34 +19,70 @@ description="Try guessing the category of each image displayed, from the options
 classdict = OrderedDict()
 for line in open('LOC_synset_mapping.txt', 'r').readlines():
     try:
-        classdict[line.split(' ')[0]]= ' '.join(line.split(' ')[1:]).replace('\n','')
+        classdict[line.split(' ')[0]]= ' '.join(line.split(' ')[1:]).replace('\n','').split(',')[0]
     except:
         continue
-
+classes = list(classdict.values())
 imagedict={}
 with open('image_labels.csv', 'r') as csv_file:
     reader = csv.DictReader(csv_file)
     for row in reader:
         imagedict[row['image_name']] = row['image_label']
+images= list(imagedict.keys())
+labels = list(set(imagedict.values()))
 
 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])
+    return model.config.id2label[predicted_class_idx]
+
 
-def check_answer(im):
+def random_image():
+    imname = random.choice(images)
+    im = Image.open('images/'+ imname +'.jpg')
+    label = str(imagedict[imname])
+    labels.remove(label)
+    options = sample(labels,3)
+    options.append(label)
+    random.shuffle(options)
+    options = [classes[int(i)] for i in options]
+    return im, label, gr.Radio.update(choices=options), None
 
-    return {'cat': 0.3, 'dog': 0.7}
+def check_score(pred, truth, current_score):
+    if pred == classes[int(truth)]:
+        return current_score + 1, f"Your score is {current_score}"
+    return current_score, f"Your score is {current_score}"
 
+def compare_score(userclass, prediction):
+    print(userclass)
+    print(prediction)
+    if userclass == str(prediction).split(',')[0]:
+        return "Great! You and the model agree on the category"
+    return "You and the model disagree"
 
 with gr.Blocks() as demo:
+    user_score = gr.State(0)
+    model_score = gr.State(0)
+    image_label = gr.State()
+    prediction = gr.State()
+
     with gr.Row():
         with gr.Column():
-            im = Image.open('images/'+sample(imagedict.keys(),1)[0]+'.jpg')
-            image = gr.Image(im,shape=(224, 224))
-            radio = gr.Radio(["option1", "option2", "option3"], label="Pick a category")
+            image = gr.Image(shape=(448, 448))
+            radio = gr.Radio(["option1", "option2", "option3"], label="Pick a category", interactive=True)
+        with gr.Column():
+            prediction = gr.Label(label="Model Prediction")
+            score = gr.Label(label="Your Score")
+            message = gr.Text()
+
+    btn = gr.Button("Next image")
 
+    demo.load(random_image, None, [image, image_label, radio, prediction])
+    radio.change(model_classify, image, prediction)
+    radio.change(check_score, [radio, image_label, user_score], [user_score, score])
+    radio.change(compare_score, [radio, prediction], message)
+    btn.click(random_image, None, [image, image_label, radio, prediction])
 
 demo.launch()