Refactor prediction function: return top 3 class probabilities with corresponding class names

app.py

@@ -69,10 +69,15 @@ def predict(data):
 
     preds = preds[0]
     print(preds)
+
+    top_3_classes = np.argsort(preds)[-3:][::-1]
+    top_3_probs = preds[top_3_classes]
     
+    class_names = [labels[i] for i in top_3_classes]
+
+    print(class_names, top_3_probs, top_3_classes)
 
-    # Return the probability for each class
-    return {label: float(pred) for label, pred in zip(labels, preds)}
+    return {class_names[i]: top_3_probs[i] for i in range(3)}
 
 # Top 3 classes
 label = gr.Label(num_top_classes=3)