Update app.py

app.py

@@ -3,25 +3,25 @@ import cv2
 from streamlit_drawable_canvas import st_canvas
 from keras.models import load_model
 import numpy as np
+from scipy.ndimage import interpolation
 
-# Sidebar configs
-drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw", "line", "rect", "circle", "transform"))
-stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 15)
-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"])
+# Sidebar: Canvas controls
+drawing_mode = st.sidebar.selectbox("Drawing tool:", ("freedraw",))
+stroke_width = st.sidebar.slider("Stroke width: ", 10, 25, 20)
+stroke_color = "#000000"  # Black
+bg_color = "#FFFFFF"      # White
 realtime_update = st.sidebar.checkbox("Update in realtime", True)
 
-# Load trained MNIST model
+# Load MNIST model
 @st.cache_resource
 def load_mnist_model():
     return load_model("mnist.keras")
 
 model = load_mnist_model()
 
-# Canvas for user drawing
+# Streamlit drawable canvas
 canvas_result = st_canvas(
-    fill_color="rgba(255, 165, 0, 0.3)",  # translucent fill
+    fill_color="rgba(0,0,0,0)",  # Transparent fill
     stroke_width=stroke_width,
     stroke_color=stroke_color,
     background_color=bg_color,
@@ -32,31 +32,52 @@ canvas_result = st_canvas(
     key="canvas",
 )
 
-# Prediction block
-if canvas_result.image_data is not None:
-    st.image(canvas_result.image_data, caption="Original Drawing")
+def preprocess(img):
+    # Convert to grayscale
+    gray = cv2.cvtColor(img.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+
+    # Invert (black bg, white digit)
+    gray = 255 - gray
 
-    # Convert RGBA to grayscale
-    img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
+    # Apply threshold
+    _, thresh = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
 
-    # Invert colors: black background with white digit
-    img = 255 - img
+    # Crop the digit (remove empty rows/cols)
+    if np.sum(thresh) == 0:
+        return None  # blank canvas
+    coords = cv2.findNonZero(thresh)
+    x, y, w, h = cv2.boundingRect(coords)
+    cropped = thresh[y:y+h, x:x+w]
 
-    # Resize to 28x28
-    img_resized = cv2.resize(img, (28, 28), interpolation=cv2.INTER_AREA)
+    # Resize keeping aspect ratio
+    h, w = cropped.shape
+    if h > w:
+        resized = cv2.resize(cropped, (int(20 * w / h), 20))
+    else:
+        resized = cv2.resize(cropped, (20, int(20 * h / w)))
 
-    # Apply binary threshold to remove noise
-    _, img_thresh = cv2.threshold(img_resized, 100, 255, cv2.THRESH_BINARY)
+    # Add padding to get 28x28
+    h, w = resized.shape
+    top = (28 - h) // 2
+    bottom = 28 - h - top
+    left = (28 - w) // 2
+    right = 28 - w - left
+    padded = cv2.copyMakeBorder(resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
 
     # Normalize
-    img_normalized = img_thresh / 255.0
+    norm = padded / 255.0
 
-    # Reshape to match model input
-    final_img = img_normalized.reshape(1, 28, 28, 1)
+    # Final reshape
+    return norm.reshape(1, 28, 28, 1)
+
+if canvas_result.image_data is not None:
+    st.image(canvas_result.image_data, caption="Original Drawing")
 
-    # Show processed image
-    st.image(img_thresh, caption="Preprocessed (28x28 Thresholded)")
+    processed = preprocess(canvas_result.image_data)
 
-    # Predict and display result
-    prediction = model.predict(final_img)
-    st.write("Prediction:", np.argmax(prediction))
+    if processed is not None:
+        st.image(processed.reshape(28, 28), caption="Processed Input (28x28)")
+        pred = model.predict(processed)
+        st.subheader(f"Prediction: {np.argmax(pred)}")
+    else:
+        st.warning("Please draw a digit!")