Refactor prediction function: enhance image preprocessing, convert to grayscale, and update model input handling

app.py

@@ -1,6 +1,7 @@
 import numpy as np
 import gradio as gr
 import tensorflow as tf
+import cv2
 
 # App title
 title = "Welcome to your first sketch recognition app!"
@@ -17,8 +18,6 @@ head = (
 # GitHub repository link
 ref = "Find the complete code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# Image size: 28x28
-img_size = 28
 
 # Class names (from 0 to 9)
 labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
@@ -26,7 +25,7 @@ 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(data):
     print(data['composite'].shape)
     # Reshape image to 28x28
@@ -40,7 +39,34 @@ def predict(data):
     # 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)}
+    return {class_names[i]: top_3_probs[i] for i in range(3)} """
+
+def predict(data):
+    # Extract the 'image' key from the input dictionary
+    img = data['image']
+    # Convert to NumPy array
+    img = np.array(img)
+
+    # Handle RGBA or RGB images
+    if img.shape[-1] == 4:  # RGBA
+        img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
+    if img.shape[-1] == 3:  # RGB
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+
+    # Resize image to 28x28
+    img = cv2.resize(img, (28, 28))
+
+    # Normalize pixel values to [0, 1]
+    img = img / 255.0
+
+    # Reshape to match model input
+    img = img.reshape(1, 28, 28, 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)}
 
 # Top 3 classes
 label = gr.Label(num_top_classes=3)
@@ -54,4 +80,4 @@ interface = gr.Interface(
     description=head,
     article=ref
 )
-interface.launch()
+interface.launch(share=True)

requirements.txt

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