Update app.py

app.py

@@ -66,28 +66,37 @@ import streamlit as st
 import cv2
 import numpy as np
 from streamlit_drawable_canvas import st_canvas
-from tensorflow.keras.models import load_model
+from keras.models import load_model
 
-# === Load model (trained on single digits) ===
+# 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)
+
+# Mode selection
+mode = st.sidebar.radio("Select Prediction Mode", ["Single Digit", "Multi Digit"])
+
+# === Load models ===
 @st.cache_resource
-def load_digit_model():
+def load_single_digit_model():
     return load_model("mnist_model.keras")
 
-model = load_digit_model()
+@st.cache_resource
+def load_multi_digit_model():
+    return load_model("best_model.keras")  # Your multi-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: ", "#000000")
-bg_color = st.sidebar.color_picker("Background color: ", "#FFFFFF")
-realtime_update = st.sidebar.checkbox("Update in realtime", True)
+model_single = load_single_digit_model()
+model_multi = load_multi_digit_model()
 
-# === Title ===
-st.title("🧠 MNIST Digit Recognizer")
-st.caption("Draw digits and let the model predict them! Choose **Single** or **Multi** mode from the sidebar.")
+# === Streamlit UI ===
+st.title("🧠 Digit Recognition App")
+st.subheader(f"✏️ Mode: {mode}")
 
-# === Drawing Canvas ===
+# Create drawing canvas
 canvas_result = st_canvas(
     fill_color="rgba(255, 165, 0, 0.3)",
     stroke_width=stroke_width,
@@ -96,55 +105,47 @@ canvas_result = st_canvas(
     update_streamlit=realtime_update,
     height=280,
     width=280,
-    drawing_mode="freedraw",
+    drawing_mode=drawing_mode,
     key="canvas",
 )
 
-# === Prediction Logic ===
+# Prediction Section
 if canvas_result.image_data is not None:
-    st.subheader("🖼️ Original Drawing")
-    st.image(canvas_result.image_data, use_column_width=False, width=280)
+    st.markdown("---")
+    st.subheader("🧪 Prediction Results")
 
-    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)
+    # Preprocess drawing
+    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    img = 255 - img  # Invert
+    _, thresh = cv2.threshold(img, 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}**")
+        norm = resized.astype("float32") / 255.0
+        input_img = norm.reshape(1, 28, 28, 1)
+
+        prediction = model_single.predict(input_img)
+        digit = np.argmax(prediction)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(resized, width=200, caption="28x28 Preprocessed")
+        with col2:
+            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.")
+        resized = cv2.resize(thresh, (80, 28))  # Resize to match your model (width=80, height=28)
+        norm = resized.astype("float32") / 255.0
+        input_seq = norm.reshape(1, 28, 80, 1)
+
+        preds = model_multi.predict(input_seq)
+
+        # Decode predictions for each digit
+        predicted_digits = [np.argmax(p[0]) for p in preds]
+        predicted_str = ''.join(str(d) for d in predicted_digits)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(resized, width=300, caption="80x28 Multi-digit Input")
+        with col2:
+            st.success(f"🧠 Predicted Sequence: **{predicted_str}**")