Refactor sketch recognition app: simplify image preprocessing, update app description, and enhance prediction function

app.py

@@ -1,77 +1,55 @@
-import os
 import numpy as np
 import cv2
 import gradio as gr
 import tensorflow as tf
 
-# Disable oneDNN optimizations for consistent results
-os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+# app title
+title = "Welcome on your first sketch recognition app!"
 
-# 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>"
+# 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>"
 )
-article = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# Image size and labels
+# 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
 img_size = 28
-labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-# 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)}")
+# 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")
 
-def preprocess_image(img):
-    """
-    Convert PIL image to grayscale NumPy array, resize, normalize, and reshape.
-    """
-    # Convert PIL to NumPy array
+# prediction function for sketch recognition
+def predict(img):
+    # Convert from PIL to NumPy
     img = np.array(img)
-
-    # Ensure grayscale format
-    if len(img.shape) == 3:  # Check if it's RGB/RGBA
+    
+    # If the image is in RGB format, convert it to grayscale
+    if len(img.shape) == 3:
         img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
 
-    # Resize to 28x28
+    # Resize the image to 28x28
     img = cv2.resize(img, (img_size, img_size))
+    
+    # Reshape to the model's input shape (1,28,28,1)
+    img = img.reshape(1, img_size, img_size, 1)
+    
+    # model predictions
+    preds = model.predict(img)[0]
 
-    # Normalize pixel values to [0, 1]
-    img = img / 255.0
-
-    # Reshape for model input
-    return img.reshape(1, img_size, img_size, 1)
+    # return the probability for each class
+    return {label: float(pred) for label, pred in zip(labels, preds)}
 
+# top 3 of classes
+label = gr.Label(num_top_classes=3)
 
-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,
-)
-
-# Launch the app
+# open Gradio interface for sketch recognition
+interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
 interface.launch()