Update app.py

app.py

@@ -3,38 +3,61 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 
-# Load the trained model
-model = tf.keras.models.load_model("number_recognition_model_colab.keras")
+title = "Welcome on your first sketch recognition app!"
+
+head = (
+    "<center>"
+    "The robot was trained to classify numbers (from 0 to 9). To test it, write your number in the space provided."
+    "</center>"
+)
+
+ref = "Find the whole code [here](https://github.com/ovh/ai-training-examples/tree/main/apps/gradio/sketch-recognition)."
 
-# Image size and labels
 img_size = 28
 labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
-# Prediction function
+model = tf.keras.models.load_model("number_recognition_model_colab.keras")
+
 def predict(img):
     try:
-        # Ensure the image is a NumPy array and grayscale
-        img = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2GRAY)
+        # Convert the input image to a NumPy array if needed
+        if not isinstance(img, np.ndarray):
+            img = np.array(img)
+
+        # Print shape and type of the input image
+        print(f"Initial image type: {type(img)}, shape: {img.shape}")
+
+        # Ensure the image is in grayscale and has a single channel
+        if img.ndim == 3 and img.shape[-1] == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        elif img.ndim == 2:
+            img = np.expand_dims(img, axis=-1)
+
+        # Print the shape of the grayscale image
+        print(f"Grayscale image shape: {img.shape}")
+
+        # Resize the image
         img = cv2.resize(img, (img_size, img_size))
-        img = img.astype("float32") / 255.0
+
+        # Normalize the image
+        img = img.astype('float32') / 255.0
         img = img.reshape(1, img_size, img_size, 1)
 
-        # Make predictions
+        # Print the shape after resizing and normalizing
+        print(f"Processed image shape: {img.shape}")
+
         preds = model.predict(img)[0]
+
+        # Print the predictions
+        print("Predictions:", preds)
+
         return {label: float(pred) for label, pred in zip(labels, preds)}
     except Exception as e:
+        # Print the exception to the console
+        print(f"Error during prediction: {e}")
         return {"Error": str(e)}
 
-# Gradio interface
-if __name__ == "__main__":
-    demo = gr.Interface(
-        fn=predict,  # Function to call
-        inputs=gr.Sketchpad(label="Draw a number"),  # Sketchpad input
-        outputs=gr.Label(num_top_classes=3),  # Label output to show top 3 predictions
-        title="Number Recognition App",
-        description=(
-            "The model was trained to classify numbers (from 0 to 9). "
-            "Draw a number in the sketchpad below and see the prediction!"
-        )
-    )
-    demo.launch()
+label = gr.Label(num_top_classes=3)
+
+interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
+interface.launch(debug=True)