Update app.py

app.py

@@ -1,47 +1,66 @@
 import tensorflow as tf
 import gradio as gr
-import numpy as np
 import cv2
+import numpy as np
 
-# 1. Load the trained model
+# 1. Load your saved model
 model = tf.keras.models.load_model('digit_recognizer.keras')
 
-# 2. Define the classification function
 def classify_digit(image):
+    # Error handling: if no image is provided
     if image is None:
         return None
+
+    # --- PREPROCESSING ---
+    # Convert to numpy array if it isn't already
+    image = np.array(image)
+
+    # 1. Handle Color Channels
+    # If image has 4 channels (RGBA) from sketchpad, convert to Gray
+    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:
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     
-    # Preprocessing to match MNIST data format
-    # Convert to grayscale if it isn't already
-    if len(image.shape) == 3:
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-        
-    # Resize the image to 28x28 pixels
-    image = cv2.resize(image, (28, 28))
-    
-    # Reshape to (1, 28, 28, 1) to match model input shape
-    # The '1' indicates a batch size of 1
+    # 2. Resize to 28x28
+    # We use INTER_AREA for shrinking which preserves details better than default
+    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)
     
-    # Normalize pixel values (0 to 1) just like in the training notebook
+    # 5. Normalize (0 to 1)
     image = image / 255.0
-    
-    # Predict
+
+    # --- PREDICTION ---
     prediction = model.predict(image).flatten()
-    
-    # Return dictionary for Gradio Label output
     return {str(i): float(prediction[i]) for i in range(10)}
 
-# 3. Build the Gradio Interface
-# We use Sketchpad so users can draw the digit
+# --- GRADIO INTERFACE ---
+# sources=["upload", "canvas"] enables both file upload and drawing
 interface = gr.Interface(
     fn=classify_digit,
-    inputs=gr.Sketchpad(label="Draw a 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 (0-9) on the canvas to see if the Neural Network recognizes it."
+    description="Draw a digit on the canvas OR upload a photo of a digit. The model will guess what it is."
 )
 
-# 4. Launch
 if __name__ == "__main__":
     interface.launch()