Update app.py

app.py

@@ -1,70 +1,47 @@
 import streamlit as st
-import cv2
-import numpy as np
-from keras.models import load_model
 from streamlit_drawable_canvas import st_canvas
+import numpy as np
+from tensorflow.keras.models import load_model
+from PIL import Image
+import cv2
 
-# Sidebar options
-drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw",))
-stroke_width = st.sidebar.slider("Stroke width:", 10, 25, 20)
-realtime_update = st.sidebar.checkbox("Update in realtime", True)
+st.set_page_config(page_title="MNIST Digit Recognizer", layout="centered")
+st.title("🖌️ Draw a digit (0-9)")
 
-# Load your trained MNIST model
+# Load pre-trained model (you can upload your own model to the space)
 @st.cache_resource
 def load_mnist_model():
-    return load_model("mnist.keras")
+    return load_model("digit_recog.keras")  # You must upload this file to your Space
 
 model = load_mnist_model()
 
-# Create the canvas
+# Create canvas component
 canvas_result = st_canvas(
-    fill_color="rgba(0,0,0,0)",  # Transparent fill
-    stroke_width=stroke_width,
-    stroke_color="#000000",      # Black drawing
-    background_color="#FFFFFF",  # White background
-    update_streamlit=realtime_update,
-    height=280,
-    width=280,
-    drawing_mode=drawing_mode,
-    key="canvas",
+    fill_color="#000000",  # Black background
+    stroke_width=10,
+    stroke_color="#FFFFFF",  # White digit
+    background_color="#000000",
+    width=200,
+    height=200,
+    drawing_mode="freedraw",
+    key="canvas"
 )
 
-def preprocess(image):
-    # Convert RGBA to grayscale
-    gray = cv2.cvtColor(image.astype("uint8"), cv2.COLOR_RGBA2GRAY)
-
-    # Invert so the digit is white on black (like MNIST)
-    gray = 255 - gray
-
-    # Apply binary threshold
-    _, thresh = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
-
-    # Find bounding box and crop
-    if np.sum(thresh) == 0:
-        return None
-    coords = cv2.findNonZero(thresh)
-    x, y, w, h = cv2.boundingRect(coords)
-    cropped = thresh[y:y+h, x:x+w]
-
-    # Resize to 20x20
-    resized = cv2.resize(cropped, (20, 20), interpolation=cv2.INTER_AREA)
-
-    # Pad to 28x28
-    padded = np.pad(resized, ((4,4),(4,4)), mode='constant', constant_values=0)
-
-    # Normalize and reshape for model
-    norm = padded / 255.0
-    return norm.reshape(1, 28, 28, 1)
-
-# Handle prediction
 if canvas_result.image_data is not None:
-    st.image(canvas_result.image_data, caption="Your Drawing")
-
-    processed = preprocess(canvas_result.image_data)
-
-    if processed is not None:
-        st.image(processed.reshape(28, 28), caption="Processed Image (28x28)")
-        prediction = model.predict(processed)
-        st.subheader(f"Prediction: {np.argmax(prediction)}")
-    else:
-        st.warning("Please draw a digit first.")
+    # Preprocess the image for prediction
+    img = canvas_result.image_data[:, :, 0]  # Get only one channel
+    img = cv2.resize(img, (28, 28))  # Resize to 28x28
+    img = img.astype("float32") / 255.0
+    img = np.expand_dims(img, axis=0)
+    img = np.expand_dims(img, axis=-1)
+
+    st.subheader("🧠 Model Prediction")
+    pred = model.predict(img)[0]
+    predicted_class = np.argmax(pred)
+
+    st.write(f"**Predicted Digit:** `{predicted_class}`")
+    st.bar_chart(pred)
+else:
+    st.info("Draw a digit above to see the prediction.")
+
+st.caption("Made with Streamlit ✨")