Update app.py

app.py

@@ -7,55 +7,33 @@ import numpy as np
 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):
-    try:
-        # Convert the input image to a NumPy array if needed
-        if not isinstance(img, np.ndarray):
-            img = np.array(img)
 
-        # 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)
+def recognize_digit(image):
+    if image is not None:
+        image = image.reshape((1, 28, 28, 1)).astype('float32')/255
 
-        # Resize the image to the expected input size
-        img = cv2.resize(img, (img_size, img_size))
+        prediction = model.predict(image)
 
-        # Normalize the image
-        img = img.astype('float32') / 255.0
-        img = img.reshape(1, img_size, img_size, 1)
+        return {str(i): float(prediction[0][i]) for i in range(10)}
+    else:
+        return ''
 
-        # Get predictions from the model
-        preds = model.predict(img)[0]
 
-        # Return the predicted probabilities for each class
-        return {label: float(pred) for label, pred in zip(labels, preds)}
-
-    except Exception as e:
-        return {"Error": str(e)}
-
-# 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
-)
-
-# Launch the interface
-interface.launch(debug=True)
+# Build and launch the Gradio interface
+demo = gr.Interface(
+    fn = recognize_digit,
+    inputs = gr.Image(
+        shape=(28, 28), 
+        image_mode='L', 
+        invert_colors=True, 
+        source='canvas',
+        brush_radius=1, 
+        tool="color-sketch",
+    ),
+    outputs = gr.Label(num_top_classes=3),
+    live = True
+).launch(share=True)