Update app.py

app.py

@@ -1,47 +1,42 @@
 import streamlit as st
+import cv2
 from streamlit_drawable_canvas import st_canvas
+from keras.models import load_model
 import numpy as np
-from tensorflow.keras.models import load_model
-from PIL import Image
-import cv2
 
-st.set_page_config(page_title="MNIST Digit Recognizer", layout="centered")
-st.title("🖌️ Draw a digit (0-9)")
 
-# Load pre-trained model (you can upload your own model to the space)
+drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"))
+stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 10)
+stroke_color = st.sidebar.color_picker("Stroke color hex: ", "#000000")  # black
+bg_color = st.sidebar.color_picker("Background color hex: ", "#FFFFFF")  # white
+bg_image = st.sidebar.file_uploader("Background image:", type=["png", "jpg"])
+realtime_update = st.sidebar.checkbox("Update in realtime", True)
+
 @st.cache_resource
 def load_mnist_model():
-    return load_model("digit_recog.keras")  # You must upload this file to your Space
+    return load_model("digit_recog.keras")
 
 model = load_mnist_model()
 
-# Create canvas component
 canvas_result = st_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"
+    fill_color="rgba(255, 165, 0, 0.3)",
+    stroke_width=stroke_width,
+    stroke_color=stroke_color,
+    background_color=bg_color,
+    update_streamlit=realtime_update,
+    height=280,
+    width=280,
+    drawing_mode=drawing_mode,
+    key="canvas",
 )
 
 if canvas_result.image_data is not None:
-    # 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 ✨")
+    st.image(canvas_result.image_data, caption="Original Drawing")
+    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    img = 255 - img
+    img_resized = cv2.resize(img, (28, 28))
+    img_normalized = img_resized / 255.0
+    final_img = img_normalized.reshape(1, 28, 28, 1)
+    st.image(img_resized, caption="Preprocessed (28x28)")
+    prediction = model.predict(final_img)
+    st.write("Prediction:", np.argmax(prediction))