Update app.py

app.py

@@ -1,50 +1,47 @@
+import tensorflow as tf
 import gradio as gr
-import torch
-import torch.nn as nn
-from torchvision import models, transforms
-from PIL import Image
+import numpy as np
+import cv2
 
-# 1. SETUP MODEL
-# We use ResNet18 structure to match your training
-model = models.resnet18(weights=None) 
-model.fc = nn.Linear(model.fc.in_features, 10) # Adjust head to 10 classes
+# 1. Load the trained model
+model = tf.keras.models.load_model('digit_recognizer.keras')
 
-# Load your 98.79% accuracy weights
-try:
-    state_dict = torch.load("fulldigits.pt", map_location="cpu")
-    model.load_state_dict(state_dict)
-    model.eval()
-except Exception as e:
-    print(f"Error loading model: {e}")
-
-# 2. PREPROCESSING
-# Must use the ImageNet stats you trained with!
-transform = transforms.Compose([
-    transforms.Lambda(lambda x: x.convert("RGB")), # Force RGB
-    transforms.Resize((128, 128)),                 # Match training size
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
-# 3. PREDICT FUNCTION
-def predict(image):
-    if image is None: return None
-    img_tensor = transform(image).unsqueeze(0)
+# 2. Define the classification function
+def classify_digit(image):
+    if image is None:
+        return None
+    
+    # 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
+    image = image.reshape(1, 28, 28, 1)
+    
+    # Normalize pixel values (0 to 1) just like in the training notebook
+    image = image / 255.0
     
-    with torch.no_grad():
-        output = model(img_tensor)
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    # Predict
+    prediction = model.predict(image).flatten()
     
-    return {str(i): float(probabilities[i]) for i in range(10)}
+    # Return dictionary for Gradio Label output
+    return {str(i): float(prediction[i]) for i in range(10)}
 
-# 4. INTERFACE
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Draw or Upload Digit"),
+# 3. Build the Gradio Interface
+# We use Sketchpad so users can draw the digit
+interface = gr.Interface(
+    fn=classify_digit,
+    inputs=gr.Sketchpad(label="Draw a Digit"),
     outputs=gr.Label(num_top_classes=3),
     title="Handwritten Digit Recognizer",
-    description="A ResNet18 model fine-tuned to 98.79% accuracy."
+    description="Draw a digit (0-9) on the canvas to see if the Neural Network recognizes it."
 )
 
+# 4. Launch
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()