Refactor sketch recognition app: simplify prediction function, update image handling, and remove OpenCV dependency

app.py

@@ -1,7 +1,6 @@
 import numpy as np
 import gradio as gr
 import tensorflow as tf
-import cv2
 
 # App title
 title = "Welcome to your first sketch recognition app!"
@@ -29,22 +28,20 @@ model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
 # Prediction function for sketch recognition
 def predict(data):
-    # Extract the 'image' key from the input dictionary
-    img = data['image']
-    # Convert to NumPy array
-    img = np.array(img)
-    # Convert to grayscale
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    # Resize image to 28x28
-    img = cv2.resize(img, (img_size, img_size))
-    # Normalize pixel values
-    img = img / 255.0
-    # Reshape image to match model input
-    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)}
+    # Reshape image to 28x28
+    img = np.reshape(data, (1, img_size, img_size, 1))
+    # Make prediction
+    pred = model.predict(img)
+    # Get top class
+    top_class = np.argmax
+    # Get top 3 classes
+    top_3_classes = np.argsort(pred[0])[-3:][::-1]
+    # Get top 3 probabilities
+    top_3_probs = pred[0][top_3_classes]
+    # Get class names
+    class_names = [labels[i] for i in top_3_classes]
+    # Return class names and probabilities
+    return {class_names[i]: top_3_probs[i] for i in range(3)}
 
 # Top 3 classes
 label = gr.Label(num_top_classes=3)
@@ -52,7 +49,7 @@ label = gr.Label(num_top_classes=3)
 # Open Gradio interface for sketch recognition
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Sketchpad(),
+    inputs=gr.Sketchpad(crop_size=(28,28), type='numpy', image_mode='L', brush=gr.Brush()),
     outputs=label,
     title=title,
     description=head,

requirements.txt

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