Update app.py

app.py

@@ -3,64 +3,72 @@ import gradio as gr
 import cv2
 import numpy as np
 
-# 1. Load your saved model
+# 1. Load the model
 model = tf.keras.models.load_model('digit_recognizer.keras')
 
 def classify_digit(image):
-    # Error handling: if no image is provided
-    if image is None:
+    if image is None: 
         return None
-
-    # --- PREPROCESSING ---
-    # Convert to numpy array if it isn't already
+    
+    # Robust check: Gradio 4.x Sketchpad might return a dictionary
+    if isinstance(image, dict):
+        image = image['composite']
+        
     image = np.array(image)
-
-    # 1. Handle Color Channels
-    # If image has 4 channels (RGBA) from sketchpad, convert to Gray
+    
+    # --- PREPROCESSING ---
+    # 1. Handle different input formats (RGBA from sketchpad, RGB from upload)
     if image.shape[-1] == 4:
-        image = cv2.cvtColor(image, cv2.COLOR_RGBA2GRAY)
-    # If image has 3 channels (RGB) from upload, convert to Gray
-    elif image.shape[-1] == 3:
+        # RGBA: Composite onto white background then convert to Gray
+        background = np.ones((image.shape[0], image.shape[1], 3), dtype=np.uint8) * 255
+        alpha = image[:, :, 3] / 255.0
+        for c in range(3):
+            background[:, :, c] = alpha * image[:, :, c] + (1 - alpha) * background[:, :, c]
+        image = cv2.cvtColor(background, cv2.COLOR_RGB2GRAY)
+    elif len(image.shape) == 3 and image.shape[-1] == 3:
+        # RGB: Convert to Gray
         image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     
-    # 2. Resize to 28x28
-    # We use INTER_AREA for shrinking which preserves details better than default
+    # 2. Resize to 28x28 (Model Requirement)
     image = cv2.resize(image, (28, 28), interpolation=cv2.INTER_AREA)
-
-    # 3. Invert Colors (Critical Step)
-    # MNIST models expect White Text on Black Background.
-    # If the image is mostly bright (like white paper), we must invert it.
-    avg_brightness = np.mean(image)
-    if avg_brightness > 127:  # If the image is mostly white/light
-        image = 255 - image   # Invert to black background
-
-    # 4. Reshape for Model
-    # (1 sample, 28 height, 28 width, 1 channel)
-    image = image.reshape(1, 28, 28, 1)
     
-    # 5. Normalize (0 to 1)
-    image = image / 255.0
-
+    # 3. Invert Colors (Critical)
+    # MNIST expects white digit on black background. 
+    # If image is mostly bright (white paper/canvas), invert it.
+    if np.mean(image) > 127:
+        image = 255 - image
+        
+    # 4. Normalize & Reshape
+    image = image.reshape(1, 28, 28, 1) / 255.0
+    
     # --- PREDICTION ---
     prediction = model.predict(image).flatten()
     return {str(i): float(prediction[i]) for i in range(10)}
 
-# --- GRADIO INTERFACE ---
-# sources=["upload", "canvas"] enables both file upload and drawing
-interface = gr.Interface(
-    fn=classify_digit,
-    inputs=gr.Image(
-        type="numpy", 
-        label="Draw or Upload Digit", 
-        image_mode="L",  # "L" attempts to convert to grayscale immediately
-        sources=["upload", "canvas"],
-        height=400,
-        width=400
-    ),
-    outputs=gr.Label(num_top_classes=3),
-    title="Handwritten Digit Recognizer",
-    description="Draw a digit on the canvas OR upload a photo of a digit. The model will guess what it is."
-)
+# --- UI SETUP ---
+# We use gr.Blocks to create a custom layout with Tabs
+with gr.Blocks() as demo:
+    gr.Markdown("## Handwritten Digit Recognizer")
+    gr.Markdown("Draw a digit (0-9) or upload a photo to test the model.")
+    
+    with gr.Tabs():
+        # Tab 1: Drawing Interface
+        with gr.Tab("Draw Digit"):
+            sketchpad = gr.Sketchpad(label="Draw Here", type="numpy", brush=gr.Brush(color="#000000", thickness=20))
+            btn_draw = gr.Button("Predict Drawing", variant="primary")
+            
+        # Tab 2: Upload Interface
+        with gr.Tab("Upload Photo"):
+            # sources=["upload", "clipboard"] fixes your specific error
+            upload = gr.Image(label="Upload Image", sources=["upload", "clipboard"], type="numpy")
+            btn_upload = gr.Button("Predict Upload", variant="primary")
+    
+    # Output is shared
+    label = gr.Label(num_top_classes=3, label="Prediction")
+    
+    # Connect both buttons to the same function
+    btn_draw.click(fn=classify_digit, inputs=sketchpad, outputs=label)
+    btn_upload.click(fn=classify_digit, inputs=upload, outputs=label)
 
 if __name__ == "__main__":
-    interface.launch()
+    demo.launch()