Refactor prediction function: streamline prediction logic and update return format to include probabilities for all classes

app.py

@@ -25,22 +25,6 @@ labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight"
 # Load model (trained on MNIST dataset)
 model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
-""" # Prediction function for sketch recognition
-def predict(data):
-    print(data['composite'].shape)
-    # Reshape image to 28x28
-    img = np.reshape(data['composite'], (1, img_size, img_size, 1))
-    # Make prediction
-    pred = model.predict(img)
-    # Get top class
-    top_3_classes = np.argsort(pred[0])[-3:][::-1]
-    # Get top 3 probabilities
-    top_3_probs = pred[0][top_3_classes]
-    # Get class names
-    class_names = [labels[i] for i in top_3_classes]
-    # Return class names and probabilities
-    return {class_names[i]: top_3_probs[i] for i in range(3)} """
-
 def predict(data):
     # Extract the 'composite' key from the input dictionary
     img = data['composite']
@@ -63,20 +47,18 @@ def predict(data):
     img = img.reshape(1, 28, 28, 1)
 
     # Model predictions
-    preds = model.predict(img)
-
-    print(preds)
-
-    preds = preds[0]
-    print(preds)
+    preds = model.predict(img)[0]
 
     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)
+    print("class_names, top_3_probs, top_3_classes" , class_names, top_3_probs, top_3_classes)
 
+    """ 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)}
 
 # Top 3 classes