Update app.py

app.py

@@ -2,118 +2,119 @@ import streamlit as st
 import cv2
 from streamlit_drawable_canvas import st_canvas
 from keras.models import load_model
-from keras.datasets import mnist
 import numpy as np
-import random
 
-# Apply custom CSS
+# Page setup
+st.set_page_config(page_title="Digit Recognizer", layout="centered")
+
+# Load the trained MNIST model
+@st.cache_resource
+def load_mnist_model():
+    return load_model("mnist_model.keras")  # Ensure this model is accurate (CNN preferred)
+
+model = load_mnist_model()
+
+# Styling
 st.markdown("""
     <style>
-    body {
-        background-color: #f5f7fa;
+    .main-title {
+        text-align: center;
+        font-size: 36px;
+        color: #2c3e50;
+        margin-bottom: 10px;
     }
-    .main {
-        background: linear-gradient(to right, #f8f9fa, #e0f7fa);
-        border-radius: 10px;
-        padding: 20px;
+    .subtitle {
+        text-align: center;
+        font-size: 18px;
+        color: #555;
     }
-    .title {
+    .result-box {
+        background-color: #e8f5e9;
+        padding: 10px;
+        border-radius: 8px;
+        margin-top: 15px;
         text-align: center;
-        font-size: 36px;
-        color: #0077b6;
-        font-family: 'Segoe UI', sans-serif;
-        margin-bottom: 20px;
     }
-    .st-cf {
-        background-color: white;
-        border-radius: 15px;
-        padding: 15px;
-        box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);
+    .digit {
+        font-size: 28px;
+        color: #2e7d32;
+        font-weight: bold;
     }
     </style>
 """, unsafe_allow_html=True)
 
-# Load MNIST model
-@st.cache_resource
-def load_mnist_model():
-    return load_model("mnist_model.keras")
-
-model = load_mnist_model()
+st.markdown('<div class="main-title">✏️ Draw a Digit</div>', unsafe_allow_html=True)
+st.markdown('<div class="subtitle">Draw a digit (0-9) and get an accurate prediction</div>', unsafe_allow_html=True)
+
+# Sidebar settings
+st.sidebar.header("Canvas Settings")
+stroke_width = st.sidebar.slider("Stroke Width", 5, 25, 15)
+stroke_color = st.sidebar.color_picker("Stroke Color", "#000000")
+bg_color = st.sidebar.color_picker("Background Color", "#FFFFFF")
+realtime = st.sidebar.checkbox("Update in Realtime", True)
+
+# Canvas
+canvas_result = st_canvas(
+    fill_color="rgba(255, 165, 0, 0.3)",
+    stroke_width=stroke_width,
+    stroke_color=stroke_color,
+    background_color=bg_color,
+    update_streamlit=realtime,
+    height=280,
+    width=280,
+    drawing_mode="freedraw",
+    key="canvas",
+)
+
+# Preprocessing function
+def preprocess_drawn_image(img_data):
+    img_gray = cv2.cvtColor(img_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    img_gray = 255 - img_gray  # Invert for white digit on black
+
+    # Threshold to remove background noise
+    _, img_thresh = cv2.threshold(img_gray, 50, 255, cv2.THRESH_BINARY)
+
+    # Find contours to crop the digit
+    contours, _ = cv2.findContours(img_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if len(contours) == 0:
+        return None
+
+    x, y, w, h = cv2.boundingRect(contours[0])
+    digit_crop = img_thresh[y:y+h, x:x+w]
+
+    # Fit into square and resize to 20x20
+    max_side = max(w, h)
+    square_digit = np.zeros((max_side, max_side), dtype=np.uint8)
+    x_offset = (max_side - w) // 2
+    y_offset = (max_side - h) // 2
+    square_digit[y_offset:y_offset+h, x_offset:x_offset+w] = digit_crop
+    digit_resized = cv2.resize(square_digit, (20, 20))
+
+    # Place in center of 28x28 image
+    final_img = np.zeros((28, 28), dtype=np.uint8)
+    final_img[4:24, 4:24] = digit_resized
+
+    # Normalize
+    final_img = final_img / 255.0
+    return final_img.reshape(1, 28, 28, 1)
+
+# Prediction
+if canvas_result.image_data is not None:
+    processed_img = preprocess_drawn_image(canvas_result.image_data)
+
+    if processed_img is not None:
+        st.image(processed_img.reshape(28, 28), caption="🧼 Preprocessed Image", clamp=True, channels="GRAY")
+
+        prediction = model.predict(processed_img)
+        pred_digit = int(np.argmax(prediction))
+        confidence = float(np.max(prediction)) * 100
+
+        st.markdown(f"""
+            <div class='result-box'>
+                🧠 Predicted Digit: <span class='digit'>{pred_digit}</span><br>
+                📊 Confidence: <strong>{confidence:.2f}%</strong>
+            </div>
+        """, unsafe_allow_html=True)
+    else:
+        st.warning("Couldn't detect a digit. Please try drawing again.")
 
-# Load MNIST test set for generator
-@st.cache_data
-def load_test_images():
-    (_, _), (x_test, y_test) = mnist.load_data()
-    return x_test, y_test
-
-x_test, y_test = load_test_images()
-
-# App title
-st.markdown('<div class="title">🧠 Mindist: Draw or Generate a Digit</div>', unsafe_allow_html=True)
-
-# Sidebar controls
-st.sidebar.title("🖌️ Canvas Settings")
-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", "#000000")
-bg_color = st.sidebar.color_picker("🌄 Background Color", "#FFFFFF")
-bg_image = st.sidebar.file_uploader("🖼️ Background Image", type=["png", "jpg"])
-realtime_update = st.sidebar.checkbox("🔄 Update in Realtime", True)
-
-# Tabs for drawing or generating
-tab1, tab2 = st.tabs(["🎨 Draw Your Own", "🤖 Generate Random Digit"])
-
-with tab1:
-    col1, col2 = st.columns([1, 1])
-
-    with col1:
-        st.subheader("🎨 Draw Here")
-        canvas_result = st_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",
-        )
-
-    with col2:
-        if canvas_result.image_data is not None:
-            st.subheader("🖼️ Your Drawing")
-            st.image(canvas_result.image_data, use_column_width=True)
-
-    if canvas_result.image_data is not None:
-        st.markdown("---")
-        st.subheader("📊 Prediction from 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)
-
-        col3, col4 = st.columns([1, 1])
-        with col3:
-            st.image(img_resized, caption="🧼 Preprocessed (28x28)", clamp=True, channels="GRAY")
-        with col4:
-            prediction = model.predict(final_img)
-            predicted_digit = int(np.argmax(prediction))
-            st.markdown(f"<h3 style='color:#00796b;'>✅ Predicted Digit: <strong>{predicted_digit}</strong></h3>", unsafe_allow_html=True)
-
-with tab2:
-    st.subheader("🎲 Random Digit Generator (MNIST)")
-    if st.button("🔁 Generate Random Digit"):
-        idx = random.randint(0, len(x_test) - 1)
-        random_img = x_test[idx]
-        true_label = y_test[idx]
-
-        st.image(random_img, width=150, caption=f"🧾 True Label: {true_label}", clamp=True, channels="GRAY")
-
-        input_img = random_img.reshape(1, 28, 28, 1) / 255.0
-        pred = model.predict(input_img)
-        pred_label = int(np.argmax(pred))
-
-        st.markdown(f"<h3 style='color:#1e88e5;'>🤖 Predicted Digit: <strong>{pred_label}</strong></h3>", unsafe_allow_html=True)