Update app.py

app.py

@@ -3,19 +3,15 @@ import gradio as gr
 import tensorflow as tf
 import numpy as np
 
-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)."
+# Title and description for the interface
+title = "Welcome to your first sketch recognition app!"
+head = "<center>The robot was trained to classify numbers (0 to 9). To test it, write your number in the space provided.</center>"
 
+# Image size and label mapping
 img_size = 28
 labels = ["zero", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine"]
 
+# Load the trained model
 model = tf.keras.models.load_model("number_recognition_model_colab.keras")
 
 def predict(img):
@@ -24,41 +20,42 @@ def predict(img):
         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
+        # Convert the image to grayscale if it's not already
         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
+        # Resize the image to the expected input size
         img = cv2.resize(img, (img_size, img_size))
 
         # Normalize the image
         img = img.astype('float32') / 255.0
         img = img.reshape(1, img_size, img_size, 1)
 
-        # Print the shape after resizing and normalizing
-        print(f"Processed image shape: {img.shape}")
-
+        # Get predictions from the model
         preds = model.predict(img)[0]
 
-        # Tahmin sonuçlarını formatla
-        sorted_preds = sorted(zip(labels, preds), key=lambda x: x[1], reverse=True)[:3]
-        formatted_preds = "\n".join([f"{label}: {prob:.2f}" for label, prob in sorted_preds])
+        # Return the predicted probabilities for each class
+        return {label: float(pred) for label, pred in zip(labels, preds)}
 
-        return formatted_preds
     except Exception as e:
-        # Print the exception to the console
-        print(f"Error during prediction: {e}")
         return {"Error": str(e)}
 
-label = gr.Label(num_top_classes=3)
+# Use a sketchpad as input for drawing
+input_component = gr.Sketchpad()
+
+# Output will show the top 3 predicted classes
+output_component = gr.Label(num_top_classes=3)
+
+# Create the Gradio interface
+interface = gr.Interface(
+    fn=predict, 
+    inputs=input_component, 
+    outputs=output_component, 
+    title=title, 
+    description=head
+)
 
-interface = gr.Interface(fn=predict, inputs="sketchpad", outputs=label, title=title, description=head, article=ref)
-interface.launch(debug=True)
+# Launch the interface
+interface.launch(debug=True)