Update app.py

app.py

@@ -65,84 +65,96 @@
 import streamlit as st
 import cv2
 import numpy as np
-from streamlit_drawable_canvas import st_canvas
 from keras.models import load_model
+from streamlit_drawable_canvas import st_canvas
 
-# === Load models with caching ===
+# === Load models ===
 @st.cache_resource
 def load_single_digit_model():
-    return load_model("mnist_model.keras")
+    return load_model("single_digit_model.keras")
 
 @st.cache_resource
 def load_multi_digit_model():
-    return load_model("best_model.keras")
+    return load_model("best_model.keras")  # multi-digit model
+
+single_digit_model = load_single_digit_model()
+multi_digit_model = load_multi_digit_model()
 
-model_single = load_single_digit_model()
-model_multi = load_multi_digit_model()
+# === Helper function to clean prediction ===
+def clean_prediction(predicted_digits):
+    """
+    Removes junk or padded digits like trailing 0s or 1s and keeps only valid 0–9 digits.
+    You can further tune this logic based on training patterns.
+    """
+    digits = [str(d) for d in predicted_digits if 0 <= d <= 9]
+    return ''.join(digits)
 
-# === Sidebar settings ===
+# === 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")
+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)
 
-# === Streamlit layout ===
-st.title("🧠 Digit Recognition App")
-st.subheader("Draw a digit or number below 👇")
-
-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",
-)
-
-# === Prediction logic ===
+# === Title ===
+st.title("🖌️ Multi-Digit and Single-Digit Drawing: Predict Instantly")
+
+# === Create a two-column layout ===
+col1, col2 = st.columns([1, 1])
+
+# === Canvas for drawing ===
+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",
+    )
+
+# === Display original drawing ===
+with col2:
+    if canvas_result.image_data is not None:
+        st.subheader("Original Drawing")
+        st.image(canvas_result.image_data, use_column_width=True)
+
+# === Image preprocessing and prediction ===
 if canvas_result.image_data is not None:
     st.markdown("---")
-    st.subheader("Prediction Output")
-
-    img = canvas_result.image_data.astype("uint8")
-    img_gray = cv2.cvtColor(img, cv2.COLOR_RGBA2GRAY)
-    img_gray = 255 - img_gray  # Invert
-    _, thresh = cv2.threshold(img_gray, 30, 255, cv2.THRESH_BINARY)
-
-    # Detect bounding box of the drawing
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if contours:
-        x, y, w, h = cv2.boundingRect(np.vstack(contours))
+    st.subheader("Preprocessed Image & Prediction")
+    
+    # Preprocess image
+    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    img = 255 - img  # Invert colors
+    img_resized = cv2.resize(img, (28, 80))  # Resize to match multi-digit model input shape
+    img_normalized = img_resized / 255.0
+    final_img = img_normalized.reshape(1, 28, 80, 1)
+
+    # === Choose which model to use based on the image size ===
+    # If image is more likely to be a single digit (e.g., smaller width), use the single digit model
+    if img_resized.shape[1] < 50:  # This is an arbitrary threshold for width
+        model_to_use = single_digit_model
     else:
-        w = 0  # fallback
-
-    # Use width of the drawn content to decide
-    if w < 50:
-        st.info("✏️ Detected Single-Digit Input")
-        resized = cv2.resize(thresh, (28, 28))
-        input_img = resized.astype("float32") / 255.0
-        input_img = input_img.reshape(1, 28, 28, 1)
-
-        prediction = model_single.predict(input_img)
-        predicted_digit = np.argmax(prediction)
-
-        st.image(resized, caption="Preprocessed 28x28 Image", width=200)
-        st.success(f"🧠 Predicted Digit: **{predicted_digit}**")
+        model_to_use = multi_digit_model
 
+    # Predict using the selected model
+    preds = model_to_use.predict(final_img)
+    
+    # For multi-digit model, decode and clean prediction
+    if model_to_use == multi_digit_model:
+        predicted_digits = [np.argmax(p[0]) for p in preds]
+        predicted_str = clean_prediction(predicted_digits)
     else:
-        st.info("🔢 Detected Multi-Digit Input")
-        resized = cv2.resize(thresh, (100, 28))
-        input_img = resized.astype("float32") / 255.0
-        input_img = input_img.reshape(1, 28, 100, 1)
+        # For single digit model, directly decode
+        predicted_str = str(np.argmax(preds))
 
-        preds = model_multi.predict(input_img)
-        predicted_digits = [np.argmax(p[0]) for p in preds]
-        predicted_str = ''.join(str(d) for d in predicted_digits)
+    # Show prediction result
+    st.markdown(f"### 🧠 Predicted Number: **{predicted_str}**")
 
-        st.image(resized, caption="Preprocessed 100x28 Image", width=250)
-        st.success(f"🧠 Predicted Number: **{predicted_str}**")