Refactor sketch recognition app: update OpenCV dependency to headless version and simplify prediction function

app.py

@@ -1,10 +1,7 @@
-import os
-os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"  # Disable oneDNN optimizations
-
+# 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!"
@@ -12,7 +9,7 @@ title = "Welcome on your first sketch recognition app!"
 # app description
 head = (
   "<center>"
-  "<img src='file/mnist-classes.png' width=400>"
+  "<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>"
 )
@@ -29,34 +26,18 @@ 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):
 
-    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]
+  # image shape: 28x28x1
+  img = cv2.resize(img, (img_size, img_size))
+  img = img.reshape(1, img_size, img_size, 1)
 
-        # Return probabilities for each class
-        return {label: float(pred) for label, pred in zip(labels, preds)}
+  # model predictions
+  preds = model.predict(img)[0]
 
-    except Exception as e:
-        return {"error": f"Image processing failed: {str(e)}"}
+  # return the probability for each classe
+  return {label: float(pred) for label, pred in zip(labels, preds)}
 
 # top 3 of classes
 label = gr.Label(num_top_classes=3)

requirements.txt

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