Update app.py

app.py

@@ -109,40 +109,40 @@ if canvas_result.image_data is not None:
     st.markdown("---")
     st.subheader("Prediction Output")
 
-    # 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)
 
-    # Resize to both possible formats
-    resized_single = cv2.resize(thresh, (28, 28))
-    resized_multi = cv2.resize(thresh, (100, 28))  # width x height
+    # 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))
+    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)
 
-    # 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
+        st.image(resized, caption="Preprocessed 28x28 Image", width=200)
+        st.success(f"🧠 Predicted Digit: **{predicted_digit}**")
 
-    if is_multi:
+    else:
         st.info("🔢 Detected Multi-Digit Input")
-        input_img = resized_multi.astype("float32") / 255.0
+        resized = cv2.resize(thresh, (100, 28))
+        input_img = resized.astype("float32") / 255.0
         input_img = input_img.reshape(1, 28, 100, 1)
 
         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)
 
-        st.image(resized_multi, caption="Preprocessed 100x28 Image", width=250)
+        st.image(resized, 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}**")
-