Update app.py

app.py

@@ -68,37 +68,32 @@ import numpy as np
 from streamlit_drawable_canvas import st_canvas
 from keras.models import load_model
 
-# 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 ===
+# === Load models with caching ===
 @st.cache_resource
 def load_single_digit_model():
     return load_model("mnist_model.keras")
 
 @st.cache_resource
 def load_multi_digit_model():
-    return load_model("best_model.keras")  # Your multi-digit model
+    return load_model("best_model.keras")
 
 model_single = load_single_digit_model()
 model_multi = load_multi_digit_model()
 
-# === Streamlit UI ===
+# === Sidebar settings ===
+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")
+realtime_update = st.sidebar.checkbox("Update in realtime", True)
+
+# === Streamlit layout ===
 st.title("🧠 Digit Recognition App")
-st.subheader(f"✏️ Mode: {mode}")
+st.subheader("Draw a digit or number below 👇")
 
-# Create drawing canvas
 canvas_result = st_canvas(
-    fill_color="rgba(255, 165, 0, 0.3)",
+    fill_color="rgba(255,165,0,0.3)",
     stroke_width=stroke_width,
     stroke_color=stroke_color,
     background_color=bg_color,
@@ -109,43 +104,45 @@ canvas_result = st_canvas(
     key="canvas",
 )
 
-# Prediction Section
+# === Prediction logic ===
 if canvas_result.image_data is not None:
     st.markdown("---")
-    st.subheader("🧪 Prediction Results")
-
-    # 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)
+    st.subheader("Prediction Output")
 
-    if mode == "Single Digit":
-        resized = cv2.resize(thresh, (28, 28))
-        norm = resized.astype("float32") / 255.0
-        input_img = norm.reshape(1, 28, 28, 1)
+    # Preprocess the drawing
+    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)
 
-        prediction = model_single.predict(input_img)
-        digit = np.argmax(prediction)
+    # Resize to both possible formats
+    resized_single = cv2.resize(thresh, (28, 28))
+    resized_multi = cv2.resize(thresh, (100, 28))  # width x height
 
-        col1, col2 = st.columns(2)
-        with col1:
-            st.image(resized, width=200, caption="28x28 Preprocessed")
-        with col2:
-            st.success(f"🧠 Predicted Digit: **{digit}**")
+    # Decide based on content width if it's likely multi-digit
+    nonzero_cols = np.count_nonzero(np.sum(thresh, axis=0) > 10)
+    is_multi = nonzero_cols > 40  # simple heuristic
 
-    elif mode == "Multi Digit":
-        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)
+    if is_multi:
+        st.info("🔢 Detected Multi-Digit Input")
+        input_img = resized_multi.astype("float32") / 255.0
+        input_img = input_img.reshape(1, 28, 100, 1)
 
-        preds = model_multi.predict(input_seq)
-
-        # Decode predictions for each digit
+        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)
 
-        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}**")
+        st.image(resized_multi, caption="Preprocessed 100x28 Image", width=250)
+        st.success(f"🧠 Predicted Number: **{predicted_str}**")
+
+    else:
+        st.info("✏️ Detected Single-Digit Input")
+        input_img = resized_single.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_single, caption="Preprocessed 28x28 Image", width=200)
+        st.success(f"🧠 Predicted Digit: **{predicted_digit}**")
+