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

app.py

@@ -1,17 +1,17 @@
-import numpy as np
-import cv2
+# import dependencies
 import gradio as gr
 import tensorflow as tf
+import cv2
+import numpy as np
 
 # app title
-title = "Welcome on your first sketch recognition app!"
+title = "Welcome to your first sketch recognition app!"
 
 # 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>"
+    "<center>"
+    "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
+    "</center>"
 )
 
 # GitHub repository link
@@ -28,28 +28,38 @@ model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
 # prediction function for sketch recognition
 def predict(img):
-    # Convert from PIL to NumPy
-    img = np.array(img)
-    
-    # If the image is in RGB format, convert it to grayscale
-    if len(img.shape) == 3:
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-
-    # 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]
-
-    # return the probability for each class
-    return {label: float(pred) for label, pred in zip(labels, preds)}
+    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)
+
+        # model predictions
+        preds = model.predict(img)[0]
+
+        # return the probability for each class
+        return {label: float(pred) for label, pred in zip(labels, preds)}
+    return None
 
 # top 3 of classes
 label = gr.Label(num_top_classes=3)
 
 # open Gradio interface for sketch recognition
-interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
+interface = gr.Interface(
+    fn=predict, 
+    inputs=gr.Sketchpad(shape=(280, 280)), 
+    outputs=label, 
+    title=title, 
+    description=head, 
+    article=ref
+)
 interface.launch()