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

app.py

@@ -1,47 +1,77 @@
-# import dependencies
+import os
+import numpy as np
+import cv2
 import gradio as gr
 import tensorflow as tf
-import cv2
 
-# app title
-title = "Welcome on your first sketch recognition app!"
+# Disable oneDNN optimizations for consistent results
+os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-# 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>"
+# App configuration
+title = "Welcome to your first sketch recognition app!"
+description = (
+    "<center>"
+    "<img src='mnist-classes.png' width=400>"
+    "<p>The robot was trained to classify numbers (from 0 to 9). "
+    "To test it, write your number in the space provided!</p>"
+    "</center>"
 )
+article = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# 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 and labels
 img_size = 28
-
-# classes name (from 0 to 9)
 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 the trained MNIST 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):
 
-  # image shape: 28x28x1
-  img = cv2.resize(img, (img_size, img_size))
-  img = img.reshape(1, img_size, img_size, 1)
+def preprocess_image(img):
+    """
+    Convert PIL image to grayscale NumPy array, resize, normalize, and reshape.
+    """
+    # Convert PIL to NumPy array
+    img = np.array(img)
+
+    # Ensure grayscale format
+    if len(img.shape) == 3:  # Check if it's RGB/RGBA
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
 
-  # model predictions
-  preds = model.predict(img)[0]
+    # Resize to 28x28
+    img = cv2.resize(img, (img_size, img_size))
 
-  # return the probability for each classe
-  return {label: float(pred) for label, pred in zip(labels, preds)}
+    # Normalize pixel values to [0, 1]
+    img = img / 255.0
 
-# top 3 of classes
-label = gr.Label(num_top_classes=3)
+    # Reshape for model input
+    return img.reshape(1, img_size, img_size, 1)
+
+
+def predict(img):
+    """
+    Predict the digit class probabilities from the input sketch image.
+    """
+    try:
+        processed_img = preprocess_image(img)
+        predictions = model.predict(processed_img)[0]
+        return {label: float(pred) for label, pred in zip(labels, predictions)}
+    except Exception as e:
+        return {"error": f"Prediction failed: {str(e)}"}
+
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs="sketchpad",
+    outputs=gr.Label(num_top_classes=3),
+    title=title,
+    description=description,
+    article=article,
+)
 
-# open Gradio interface for sketch recognition
-interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
+# Launch the app
 interface.launch()