Refactor prediction function: change labels to a dictionary for better mapping and update return format to include top 3 class probabilities with corresponding class names

app.py

@@ -20,8 +20,18 @@ ref = "Find the complete code [here](https://github.com/ovh/ai-training-examples
 
 
 # Class names (from 0 to 9)
-labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
-
+labels = {
+    0: "zero", 
+    1: "one", 
+    2: "two",
+    3: "three",
+    4: "four",
+    5: "five",
+    6: "six",
+    7: "seven",
+    8: "eight",
+    9: "nine"    
+}
 # Load model (trained on MNIST dataset)
 model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
@@ -49,17 +59,18 @@ def predict(data):
     # Model predictions
     preds = model.predict(img)[0]
 
-    top_3_classes = np.argsort(preds)[-3:][::-1]
-    top_3_probs = preds[top_3_classes]
+    print("preds", preds)
+    values_map = {preds[i]: i for i in range(len(preds))}
     
-    class_names = [labels[i] for i in top_3_classes]
+    sorted_values = sorted(preds, reverse=True)
 
-    print("class_names, top_3_probs, top_3_classes" , class_names, top_3_probs, top_3_classes)
+    labels_map = dict()
+    for i in range(3):
+        print("sorted_values[i]", sorted_values[i], values_map[sorted_values[i]])
+        labels_map[labels[values_map[sorted_values[i]]]] = sorted_values[i]
 
-    """ return {class_names[i]: top_3_probs[i] for i in range(3)} """
-    """ # return the probability for each classe
-  return {label: float(pred) for label, pred in zip(labels, preds)} """
-    return {class_names[i]: top_3_probs[i] for i in range(3)}
+    print("labels_map", labels_map)
+    return labels_map
 
 # Top 3 classes
 label = gr.Label(num_top_classes=3)