Refactor sketch recognition app: enhance image handling, improve error messages, and update app description

app.py

@@ -1,66 +1,73 @@
-# import dependencies
+import os
+import numpy as np
+import cv2
 import gradio as gr
 import tensorflow as tf
-import cv2
-import numpy as np
+from PIL import Image
 
 # app title
-title = "Welcome to your first sketch recognition app!"
+title = "Welcome on your first sketch recognition app!"
 
 # app description
 head = (
-    "<center>"
-    "<img src='mnist-classes.png' width=400><br>"
-    "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
-    "</center>"
+  "<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>"
 )
 
 # GitHub repository link
 ref = "Find the whole code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# image size: 28x28
+# Image size
 img_size = 28
 
-# classes name (from 0 to 9)
+# Classes
 labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-# load model (trained on MNIST dataset)
-model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
+# 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)}")
 
-# prediction function for sketch recognition
 def predict(img):
-    if img is not None:
-        # Convert to numpy array if not already
-        img = np.array(img)
-        
-        # Ensure grayscale
-        if len(img.shape) == 3:
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-        
-        # Resize to required dimensions
-        img = cv2.resize(img, (img_size, img_size))
-        
-        # Normalize and reshape
-        img = img.astype('float32') / 255.0
-        img = img.reshape(1, img_size, img_size, 1)
+    # If no image is provided, return an error message
+    if img is None:
+        return {"error": "No image provided."}
 
-        # model predictions
-        preds = model.predict(img)[0]
+    # Ensure the image is a PIL Image
+    if not isinstance(img, Image.Image):
+        img = Image.fromarray(np.uint8(img))
 
-        # return the probability for each class
-        return {label: float(pred) for label, pred in zip(labels, preds)}
-    return None
+    # Convert to grayscale
+    img = img.convert("L")
+    
+    # Convert PIL Image to a NumPy array of type uint8
+    img = np.array(img, dtype=np.uint8)
 
-# top 3 of classes
-label = gr.Label(num_top_classes=3)
+    # Resize to (28x28)
+    img = cv2.resize(img, (img_size, img_size))
 
-# open Gradio interface for sketch recognition
+    # Reshape to match model input shape (1, 28, 28, 1)
+    img = img.reshape(1, img_size, img_size, 1)
+
+    # Model predictions
+    preds = model.predict(img)[0]
+
+    # Return probabilities for each class
+    return {label: float(pred) for label, pred in zip(labels, preds)}
+
+# Use gr.Sketchpad to ensure a PIL image is returned
 interface = gr.Interface(
-    fn=predict, 
-    inputs=gr.Sketchpad(height=280, width=280),  # Changed from shape to height and width
-    outputs=label, 
-    title=title, 
-    description=head, 
+    fn=predict,
+    inputs=gr.Sketchpad(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title=title,
+    description=head,
     article=ref
 )
+
 interface.launch()

requirements.txt

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