OptiDigit / Home.py
Sathwikchowdary's picture
Update Home.py
436d043 verified
import streamlit as st
from streamlit_drawable_canvas import st_canvas
from keras.models import load_model
import numpy as np
import cv2
# ๐ŸŽฏ App title and intro
st.set_page_config(page_title="MNIST Digit Recognizer", layout="centered")
st.title("๐Ÿค– VisionAI: MNIST Digit Predictor")
st.markdown("Draw a digit (0-9) below and click **Predict** to see the result!")
# ๐ŸŽจ Sidebar controls
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")
bg_image = st.sidebar.file_uploader("Background image (optional):", type=["png", "jpg"])
realtime_update = st.sidebar.checkbox("Update in realtime", True)
# ๐Ÿง  Load model from local path
@st.cache_resource
def load_mnist_model():
return load_model("digit_reco.keras")
model = load_mnist_model()
# ๐Ÿ–Œ๏ธ Canvas setup
canvas_result = st_canvas(
fill_color="rgba(255, 165, 0, 0.3)", # Transparent fill
stroke_width=stroke_width,
stroke_color=stroke_color,
background_color=bg_color,
update_streamlit=realtime_update,
height=280,
width=280,
drawing_mode=drawing_mode,
key="canvas",
)
# ๐Ÿ“ค Predict button
if st.button("๐Ÿ”ฎ Predict"):
if canvas_result.image_data is not None:
st.image(canvas_result.image_data, caption="๐Ÿ–ผ๏ธ Your Drawing", use_container_width=True) # Updated
# Preprocess image
img = cv2.cvtColor(canvas_result.image_data.astype("uint8"), cv2.COLOR_RGBA2GRAY)
img = 255 - img # Invert colors
img_resized = cv2.resize(img, (28, 28))
img_normalized = img_resized / 255.0
img_reshaped = img_normalized.reshape((1, 28, 28))
# Model prediction
prediction = model.predict(img_reshaped)
predicted_digit = np.argmax(prediction)
# Show result
st.markdown(f"## ๐Ÿง  Predicted Digit: **{predicted_digit}**")
else:
st.warning("Please draw something before predicting!")