Enhance sketch recognition app: improve image processing in prediction function and add NumPy dependency

app.py

@@ -4,6 +4,7 @@ os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"  # Disable oneDNN optimizations
 import gradio as gr
 import tensorflow as tf
 import cv2
+import numpy as np
 
 # app title
 title = "Welcome on your first sketch recognition app!"
@@ -28,12 +29,34 @@ labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight"
 # load model (trained on MNIST dataset)
 model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
-# prediction function for sketch recognition
+# Prediction function for sketch recognition
 def predict(img):
-    img = cv2.resize(img, (img_size, img_size))
-    img = img.reshape(1, img_size, img_size, 1)
-    preds = model.predict(img)[0]
-    return {label: float(pred) for label, pred in zip(labels, preds)}
+
+    try:
+        # Convert PIL image to NumPy array
+        img = np.array(img)
+
+        # Ensure grayscale format (convert from RGB if necessary)
+        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))
+
+        # Normalize pixel values to [0, 1]
+        img = img / 255.0
+
+        # Reshape to match the model input shape
+        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)}
+
+    except Exception as e:
+        return {"error": f"Image processing failed: {str(e)}"}
 
 # top 3 of classes
 label = gr.Label(num_top_classes=3)

requirements.txt

@@ -1,2 +1,3 @@
 tensorflow
 opencv-python
+numpy