Update app.py

app.py

@@ -3,26 +3,48 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-# Load MNIST model
+# =========================
+# Load trained model
+# =========================
+# Make sure you've trained and saved it as best_model.h5 in your notebook
 model = tf.keras.models.load_model("mnist_model.h5")
 
+# =========================
+# Prediction function
+# =========================
 def predict(image):
-    # Ensure PIL image, convert to grayscale, resize
+    """
+    Takes a PIL image, preprocesses it (grayscale + resize),
+    runs prediction using trained model, and returns predicted digit.
+    """
+    # Convert to grayscale + resize
     image = image.convert("L").resize((28, 28))
-    
-    # Normalize
-    image = np.array(image) / 255.0
-    image = image.reshape(1, 28, 28, 1)
+
+    # Convert to numpy and normalize
+    img_array = np.array(image) / 255.0
+    img_array = img_array.reshape(1, 28, 28, 1)  # batch shape
 
     # Predict
-    prediction = model.predict(image)
-    return str(np.argmax(prediction))
+    prediction = model.predict(img_array)
+    predicted_class = np.argmax(prediction, axis=1)[0]
+
+    # Also return top-3 predictions with probabilities
+    top3_indices = prediction[0].argsort()[-3:][::-1]
+    top3_probs = prediction[0][top3_indices]
+
+    result = {str(d): float(p) for d, p in zip(top3_indices, top3_probs)}
+    return result
 
-# Gradio interface (upload image)
+# =========================
+# Gradio interface
+# =========================
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", image_mode="L"),   # ← removed shape
-    outputs="label"
+    inputs=gr.Image(type="pil", image_mode="L"),
+    outputs=gr.Label(num_top_classes=3),   # show top 3 predictions
+    title="MNIST Digit Classifier",
+    description="Upload a handwritten digit (0–9) image. The model will predict the digit."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()