Update app.py

app.py

@@ -1,64 +1,150 @@
+# import streamlit as st
+# import cv2
+# from streamlit_drawable_canvas import st_canvas
+# from keras.models import load_model
+# import numpy as np
+
+# # 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 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)
+
+# # Load model with caching
+# @st.cache_resource
+# def load_mnist_model():
+#     return load_model("mnist_model.keras")
+
+# model = load_mnist_model()
+
+# st.title("🖌️ Mindist: Draw a Number, Predict Instantly")
+
+# # Create a two-column layout
+# 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("Original Drawing")
+#         st.image(canvas_result.image_data, use_column_width=True)
+
+# # Below the two columns: Show preprocessing and prediction
+# if canvas_result.image_data is not None:
+#     st.markdown("---")
+#     st.subheader("Preprocessed Image & Prediction")
+    
+#     img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+#     img = 255 - img  # Invert colors
+#     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="28x28 Preprocessed", clamp=True, channels="GRAY")
+#     with col4:
+#         prediction = model.predict(final_img)
+#         predicted_digit = np.argmax(prediction)
+#         st.markdown(f"### 🧠 Predicted Digit: **{predicted_digit}**")
 import streamlit as st
 import cv2
-from streamlit_drawable_canvas import st_canvas
-from keras.models import load_model
 import numpy as np
+from streamlit_drawable_canvas import st_canvas
+from tensorflow.keras.models import load_model
 
-# Sidebar controls
-st.sidebar.title("Canvas Settings")
-drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"))
+# === Load model (trained on single digits) ===
+@st.cache_resource
+def load_digit_model():
+    return load_model("mnist_model.keras")
+
+model = load_digit_model()
+
+# === Sidebar Controls ===
+st.sidebar.title("🛠️ Settings")
+mode = st.sidebar.radio("Choose Prediction Mode", ("Single Digit", "Multi Digit"))
 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"])
+stroke_color = st.sidebar.color_picker("Stroke color: ", "#000000")
+bg_color = st.sidebar.color_picker("Background color: ", "#FFFFFF")
 realtime_update = st.sidebar.checkbox("Update in realtime", True)
 
-# Load model with caching
-@st.cache_resource
-def load_mnist_model():
-    return load_model("mnist_model.keras")
+# === Title ===
+st.title("🧠 MNIST Digit Recognizer")
+st.caption("Draw digits and let the model predict them! Choose **Single** or **Multi** mode from the sidebar.")
 
-model = load_mnist_model()
-
-st.title("🖌️ Mindist: Draw a Number, Predict Instantly")
-
-# Create a two-column layout
-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("Original Drawing")
-        st.image(canvas_result.image_data, use_column_width=True)
-
-# Below the two columns: Show preprocessing and prediction
+# === Drawing 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_update,
+    height=280,
+    width=280,
+    drawing_mode="freedraw",
+    key="canvas",
+)
+
+# === Prediction Logic ===
 if canvas_result.image_data is not None:
-    st.markdown("---")
-    st.subheader("Preprocessed Image & Prediction")
-    
-    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
-    img = 255 - img  # Invert colors
-    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="28x28 Preprocessed", clamp=True, channels="GRAY")
-    with col4:
-        prediction = model.predict(final_img)
-        predicted_digit = np.argmax(prediction)
-        st.markdown(f"### 🧠 Predicted Digit: **{predicted_digit}**")
+    st.subheader("🖼️ Original Drawing")
+    st.image(canvas_result.image_data, use_column_width=False, width=280)
+
+    img = canvas_result.image_data.astype('uint8')
+    gray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
+    gray = 255 - gray  # Invert for white digits on black
+    _, thresh = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
+
+    if mode == "Single Digit":
+        st.subheader("🔢 Single Digit Mode")
+        resized = cv2.resize(thresh, (28, 28))
+        normalized = resized.astype("float32") / 255.0
+        input_tensor = normalized.reshape(1, 28, 28, 1)
+        pred = model.predict(input_tensor, verbose=0)
+        digit = np.argmax(pred)
+        st.image(resized, width=100, caption="🧼 Cleaned & Resized")
+        st.success(f"✅ Predicted Digit: **{digit}**")
+
+    elif mode == "Multi Digit":
+        st.subheader("🔢 Multi Digit Mode")
+
+        # Detect contours
+        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        boxes = [cv2.boundingRect(c) for c in contours if cv2.contourArea(c) > 50]
+        boxes = sorted(boxes, key=lambda b: b[0])  # sort by x
+
+        predictions = []
+        for x, y, w, h in boxes:
+            digit_img = thresh[y:y+h, x:x+w]
+            digit_img = cv2.resize(digit_img, (28, 28))
+            normalized = digit_img.astype("float32") / 255.0
+            input_tensor = normalized.reshape(1, 28, 28, 1)
+            pred = model.predict(input_tensor, verbose=0)
+            digit = np.argmax(pred)
+            predictions.append((digit_img, digit))
+
+        if predictions:
+            st.markdown("### ✂️ Segmented Digits & Predictions")
+            cols = st.columns(len(predictions))
+            for i, (img, digit) in enumerate(predictions):
+                with cols[i]:
+                    st.image(img, width=64, caption=f"➡️ {digit}")
+            full_number = ''.join(str(d) for (_, d) in predictions)
+            st.success(f"📌 Final Multi-Digit Prediction: **{full_number}**")
+        else:
+            st.warning("⚠️ Couldn't detect any digits. Try writing more clearly.")