Update app.py

app.py

@@ -80,15 +80,6 @@ def load_multi_digit_model():
 single_digit_model = load_single_digit_model()
 multi_digit_model = 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 controls ===
 st.sidebar.title("Canvas Settings")
 drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"))
@@ -134,15 +125,16 @@ if canvas_result.image_data is not None:
     img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
     img = 255 - img  # Invert colors
     
-    # Resize image to match model input dimensions
-    img_resized = cv2.resize(img, (80, 28))  # Resize to match multi-digit model input shape
+    # Resize image to match the expected model input shape
+    img_resized = cv2.resize(img, (100, 28))  # Resize to match multi-digit model input shape
     img_normalized = img_resized / 255.0
-    final_img = img_normalized.reshape(1, 28, 80, 1)
+    final_img = img_normalized.reshape(1, 28, 100, 1)  # Adjust to model input shape
 
     # === 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
+        final_img = final_img.reshape(1, 28, 28, 1)  # For single digit, reshape accordingly
     else:
         model_to_use = multi_digit_model
 
@@ -152,7 +144,7 @@ if canvas_result.image_data is not None:
     # 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)
+        predicted_str = ''.join([str(d) for d in predicted_digits])
     else:
         # For single digit model, directly decode
         predicted_str = str(np.argmax(preds))