Refactor app.py: enhance prediction function, update title and description, and improve image handling

app.py

@@ -1,47 +1,72 @@
-# import dependencies
+import numpy as np
 import gradio as gr
 import tensorflow as tf
 import cv2
 
-# app title
-title = "Welcome on your first sketch recognition app!"
+# Load the trained MNIST model
+model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
-# app description
-head = (
-  "<center>"
-  "<img src='./mnist-classes.png' width=400>"
-  "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
-  "</center>"
-)
+# Class names (0 to 9)
+labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-# GitHub repository link
-ref = "Find the whole code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
+def predict(data):
+    # Extract the 'composite' key from the input dictionary
+    img = data["composite"]
+    img = np.array(img)
 
-# image size: 28x28
-img_size = 28
+    # Convert RGBA to RGB if needed
+    if img.shape[-1] == 4:  # RGBA
+        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
 
-# classes name (from 0 to 9)
-labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
+    # Convert RGB to Grayscale
+    if img.shape[-1] == 3:  # RGB
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
 
-# load model (trained on MNIST dataset)
-model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
+    # Resize image to 28x28
+    img = cv2.resize(img, (28, 28))
+
+    # Normalize pixel values to [0, 1]
+    img = img / 255.0
 
-# prediction function for sketch recognition
-def predict(img):
+    # Reshape to match model input (1, 28, 28, 1)
+    img = img.reshape(1, 28, 28, 1)
+
+    # Model predictions
+    preds = model.predict(img)[0]
+
+    # Get top 3 classes
+    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]
+
+    # Return top 3 predictions as a dictionary
+    return {class_names[i]: float(top_3_probs[i]) for i in range(3)}
+
+# Title and description
+title = "Welcome to your first sketch recognition app!"
+head = (
+    "<center>"
+    "<img src='./mnist-classes.png' width=400>"
+    "<p>The model is trained to classify numbers (from 0 to 9). "
+    "To test it, draw your number in the space provided (use the editing tools in the image editor).</p>"
+    "</center>"
+)
+ref = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-  # image shape: 28x28x1
-  img = cv2.resize(img, (img_size, img_size))
-  img = img.reshape(1, img_size, img_size, 1)
+with gr.Blocks(title=title) as demo:
+    # Display title and description
+    gr.Markdown(head)
+    gr.Markdown(ref)
 
-  # model predictions
-  preds = model.predict(img)[0]
+    with gr.Row():
+        # Using ImageEditor with type='numpy'
+        im = gr.ImageEditor(type="numpy", label="Draw your digit here (use brush and eraser)")
 
-  # return the probability for each classe
-  return {label: float(pred) for label, pred in zip(labels, preds)}
+        # Output label (top 3 predictions)
+        label = gr.Label(num_top_classes=3, label="Predictions")
 
-# top 3 of classes
-label = gr.outputs.Label(num_top_classes=3)
+    # Trigger prediction whenever the image changes
+    im.change(predict, inputs=im, outputs=label, show_progress="hidden")
 
-# open Gradio interface for sketch recognition
-interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
-interface.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True)