Update src/streamlit_app.py

src/streamlit_app.py

@@ -4,6 +4,7 @@ from PIL import Image, ImageOps
 import tensorflow as tf
 import matplotlib.pyplot as plt
 from pathlib import Path
+import cv2  # NEW
 
 st.set_page_config(page_title="MNIST Digit Recognizer", page_icon="✍️", layout="centered")
 
@@ -18,28 +19,79 @@ def load_model():
         )
     return tf.keras.models.load_model(MODEL_PATH)
 
-def preprocess_image(pil_img: Image.Image) -> np.ndarray:
+def center_digit(binary_28: np.ndarray) -> np.ndarray:
+    """
+    Centers the digit in a 28x28 binary image using its bounding box.
+    Expects values in [0,255] with digit in white (255) on black (0).
+    """
+    ys, xs = np.where(binary_28 > 0)
+    if len(xs) == 0 or len(ys) == 0:
+        return binary_28  # empty image fallback
+
+    x_min, x_max = xs.min(), xs.max()
+    y_min, y_max = ys.min(), ys.max()
+
+    digit = binary_28[y_min:y_max+1, x_min:x_max+1]
+
+    # Create blank canvas and paste centered
+    canvas = np.zeros((28, 28), dtype=np.uint8)
+    h, w = digit.shape
+    top = (28 - h) // 2
+    left = (28 - w) // 2
+
+    # Clip if something goes wrong (safety)
+    digit = digit[:min(h, 28), :min(w, 28)]
+    canvas[top:top+digit.shape[0], left:left+digit.shape[1]] = digit
+    return canvas
+
+def preprocess_image(pil_img: Image.Image,
+                     do_threshold: bool = True,
+                     threshold_value: int = 140,
+                     do_center: bool = True) -> np.ndarray:
     """
     Convert uploaded image to MNIST-like 28x28 grayscale tensor (1, 28, 28, 1).
+    Uses thresholding + optional centering for better robustness.
     """
     # 1) Convert to grayscale
     img = pil_img.convert("L")
 
-    # 2) Make it square (pad) to avoid distortion
-    img = ImageOps.pad(img, (max(img.size), max(img.size)), method=Image.Resampling.LANCZOS, color=255)
+    # 2) Make square (pad) to avoid distortion
+    img = ImageOps.pad(img, (max(img.size), max(img.size)),
+                       method=Image.Resampling.LANCZOS, color=255)
 
     # 3) Resize to 28x28
     img = img.resize((28, 28), Image.Resampling.LANCZOS)
 
-    # 4) Invert if background is dark (MNIST: black background, white digit)
-    arr = np.array(img).astype(np.float32)
+    # Convert to numpy
+    arr = np.array(img).astype(np.uint8)
+
+    # 4) Auto-invert if background is dark
+    # We want: black background (0), white digit (255) like MNIST (after threshold)
     if arr.mean() < 127:
-        arr = 255.0 - arr  # invert
+        arr = 255 - arr
+
+    # 5) Mild denoise (optional but helpful for uploaded images)
+    arr = cv2.GaussianBlur(arr, (3, 3), 0)
+
+    # 6) Thresholding (turn to clean black/white)
+    if do_threshold:
+        _, arr_bin = cv2.threshold(arr, threshold_value, 255, cv2.THRESH_BINARY_INV)
+        # After THRESH_BINARY_INV: digit becomes white (255) if it was dark on light background
+        # Ensure digit is white on black:
+        # arr_bin currently has digit as white and background as black -> perfect
+    else:
+        # If no threshold: convert to "digit white on black" approx by inverting intensity
+        arr_bin = 255 - arr
 
-    # 5) Normalize 0..1 and reshape
-    arr = arr / 255.0
-    arr = arr.reshape(1, 28, 28, 1)
-    return arr
+    # 7) Center the digit (optional)
+    if do_center:
+        arr_bin = center_digit(arr_bin)
+
+    # 8) Normalize to 0..1 (MNIST style: digit bright)
+    arr_norm = arr_bin.astype(np.float32) / 255.0
+
+    # Return shape (1, 28, 28, 1)
+    return arr_norm.reshape(1, 28, 28, 1)
 
 def plot_probabilities(probs: np.ndarray):
     plt.figure(figsize=(7, 3))
@@ -54,11 +106,17 @@ def plot_probabilities(probs: np.ndarray):
 st.title("✍️ MNIST Digit Recognizer")
 st.write("Upload an image of a handwritten digit (0–9). The model will predict the digit.")
 
-with st.expander("ℹ️ Requirements for best results"):
+with st.expander("⚙️ Preprocessing options"):
+    do_threshold = st.checkbox("Use thresholding (recommended)", value=True)
+    threshold_value = st.slider("Threshold value", 0, 255, 140, 5)
+    do_center = st.checkbox("Center the digit (recommended)", value=True)
+
+with st.expander("ℹ️ Tips for best results"):
     st.markdown(
         "- Use a **single digit** centered in the image.\n"
         "- Prefer **white background** and **dark digit**.\n"
-        "- PNG/JPG works fine."
+        "- PNG/JPG works fine.\n"
+        "- If prediction is wrong, try adjusting the **threshold** slider."
     )
 
 uploaded = st.file_uploader("Upload an image (png/jpg/jpeg)", type=["png", "jpg", "jpeg"])
@@ -71,22 +129,25 @@ except Exception as e:
 
 if uploaded is not None:
     pil_img = Image.open(uploaded)
+
     st.subheader("Your image")
     st.image(pil_img, use_container_width=True)
 
-    x = preprocess_image(pil_img)
+    x = preprocess_image(pil_img, do_threshold=do_threshold,
+                         threshold_value=threshold_value, do_center=do_center)
+
     probs = model.predict(x, verbose=0)[0]
     pred = int(np.argmax(probs))
     conf = float(np.max(probs))
 
     st.subheader("Prediction")
-    st.metric("Predicted Digit", pred, delta=None)
+    st.metric("Predicted Digit", pred)
     st.write(f"Confidence: **{conf:.3f}**")
 
     st.subheader("Model probabilities")
     plot_probabilities(probs)
 
     st.subheader("Preprocessed 28×28 image (what the model sees)")
-    st.image((x[0].reshape(28, 28) * 255).astype(np.uint8), clamp=True, use_container_width=False)
+    st.image((x[0].reshape(28, 28) * 255).astype(np.uint8), clamp=True)
 else:
     st.info("Upload an image to get a prediction.")