trying to add comparison between model and user class

app.py

@@ -11,11 +11,6 @@ feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch1
 model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")
 
 
-title="ImageNet Roulette"
-description="Try guessing the category of each image displayed, from the options provided below.\
-            After 10 guesses, we will show you your accuracy!\
-            "
-
 classdict = OrderedDict()
 for line in open('LOC_synset_mapping.txt', 'r').readlines():
     try:
@@ -37,15 +32,14 @@ def model_classify(radio, im):
         outputs = model(**inputs)
         logits = outputs.logits
         predicted_class_idx = logits.argmax(-1).item()
-        return model.config.id2label[predicted_class_idx], True
+        return model.config.id2label[predicted_class_idx], predicted_class_idx, True
     else:
-        return None, False
+        return None, None, False
 
 def random_image():
     imname = random.choice(images)
     im = Image.open('images/'+ imname +'.jpg')
     label = str(imagedict[imname])
-    print(label)
     labels.remove(label)
     options = sample(labels,3)
     options.append(label)
@@ -55,9 +49,9 @@ def random_image():
 
 def check_score(pred, truth, current_score, total_score, has_guessed):
     if not(has_guessed):
-        if pred == "lion":
+        if pred == classes[int(truth)]:
             total_score +=1
-            return current_score + 1, f"Your score is {current_score+1} out of {total_score}", total_score
+            return current_score + 1, f"Your score is {current_score+1} out of {total_score}!", total_score
         else:
             if pred is not None:
                 total_score +=1
@@ -67,34 +61,42 @@ def check_score(pred, truth, current_score, total_score, has_guessed):
 
 
 
-def compare_score(userclass, prediction):
-    if userclass == str(prediction).split(',')[0]:
-        return "Great! You and the model agree on the category"
-    return "You and the model disagree"
+def compare_score(userclass, model_class, truth, has_guessed):
+    print(model_class)
+    prediction= classes[int(model_class)]
+    if userclass == classes[int(truth)] == prediction:
+        return  "Great! You and the model both got the correct answer"
+    elif userclass == classes[int(truth)]:
+        return "Great! You guessed it right"
+    elif prediction == classes[int(truth)]:
+        return "The AI model got it right this time, try again!"
 
 with gr.Blocks() as demo:
     user_score = gr.State(0)
     model_score = gr.State(0)
     image_label = gr.State()
-    prediction = gr.State()
+    model_class = gr.State()
     total_score = gr.State(0)
     has_guessed = gr.State(False)
 
+    gr.Markdown("# ImageNet Quiz")
+    gr.Markdown("Try guessing the category of each image displayed, from the options provided below.")
     with gr.Row():
+
         with gr.Column(min_width= 900):
             image = gr.Image(shape=(600, 600))
             radio = gr.Radio(["option1", "option2", "option3"], label="Pick a category", interactive=True)
         with gr.Column():
-            prediction = gr.Label(label="The AI predicts:")
+            prediction = gr.Label(label="The AI model predicts:")
             score = gr.Label(label="Your Score")
-            #message = gr.Text()
+            message = gr.Text(label="Who guessed it right?")
 
     btn = gr.Button("Next image")
 
     demo.load(random_image, None, [image, image_label, radio, prediction])
-    radio.change(model_classify, [radio, image], [prediction, has_guessed])
+    radio.change(model_classify, [radio, image], [prediction, model_class, has_guessed])
     radio.change(check_score, [radio, image_label, user_score, total_score, has_guessed], [user_score, score, total_score])
-    #radio.change(compare_score, [radio, prediction], message)
+    radio.change(compare_score, [radio, prediction, image_label, has_guessed], message)
     btn.click(random_image, None, [image, image_label, radio, prediction])
     btn.click(lambda :False, None, has_guessed)