Refactor sketch recognition app: update app title and description, streamline image processing, and enhance prediction function

app.py

@@ -1,73 +1,61 @@
-import os
 import numpy as np
-import cv2
 import gradio as gr
 import tensorflow as tf
-from PIL import Image
+import cv2
 
-# app title
-title = "Welcome on your first sketch recognition app!"
+# App title
+title = "Welcome to your first sketch recognition app!"
 
-# app description
+# App description
 head = (
-  "<center>"
-  "<img src='./mnist-classes.png' width=400>"
-  "<p>The robot was trained to classify numbers (0 to 9). "
-  "To test it, write your number in the space provided!</p>"
-  "</center>"
+    "<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.</p>"
+    "</center>"
 )
 
 # GitHub repository link
-ref = "Find the whole code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
+ref = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# Image size
+# Image size: 28x28
 img_size = 28
 
-# Classes
+# Class names (from 0 to 9)
 labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-# Load model
-model_path = "./sketch_recognition_numbers_model.h5"
-try:
-    model = tf.keras.models.load_model(model_path)
-except Exception as e:
-    raise FileNotFoundError(f"Model file '{model_path}' not found or failed to load. {str(e)}")
-
-def predict(img):
-    # If no image is provided, return an error message
-    if img is None:
-        return {"error": "No image provided."}
-
-    # Ensure the image is a PIL Image
-    if not isinstance(img, Image.Image):
-        img = Image.fromarray(np.uint8(img))
+# 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):
+    # Extract the 'image' key from the input dictionary
+    img = data['image']
+    # Convert to NumPy array
+    img = np.array(img)
     # Convert to grayscale
-    img = img.convert("L")
-    
-    # Convert PIL Image to a NumPy array of type uint8
-    img = np.array(img, dtype=np.uint8)
-
-    # Resize to (28x28)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    # Resize image to 28x28
     img = cv2.resize(img, (img_size, img_size))
-
-    # Reshape to match model input shape (1, 28, 28, 1)
+    # Normalize pixel values
+    img = img / 255.0
+    # Reshape image to match model input
     img = img.reshape(1, img_size, img_size, 1)
-
     # Model predictions
     preds = model.predict(img)[0]
-
-    # Return probabilities for each class
+    # Return the probability for each class
     return {label: float(pred) for label, pred in zip(labels, preds)}
 
-# Use gr.Sketchpad to ensure a PIL image is returned
+# Top 3 classes
+label = gr.Label(num_top_classes=3)
+
+# Open Gradio interface for sketch recognition
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Sketchpad(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
+    inputs=gr.Sketchpad(),
+    outputs=label,
     title=title,
     description=head,
     article=ref
 )
-
 interface.launch()

requirements.txt

@@ -1,5 +1,3 @@
 tensorflow
 opencv-python-headless
-numpy
-# PIL
-Pillow
+numpy