import streamlit as st
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
import cv2
from PIL import Image, ImageOps

# Load the trained model
model = load_model('digit_recognizer_model.h5')

# Streamlit app title
st.title("Handwritten Digit Recognizer")

# Instructions
st.write("Draw a digit below and click 'Predict' to see the model's prediction.")

# Create a canvas component
from streamlit_drawable_canvas import st_canvas

# Set up the canvas
canvas_result = st_canvas(
    fill_color="black",  # Drawing background color
    stroke_width=10,
    stroke_color="white",
    background_color="black",
    height=280,
    width=280,
    drawing_mode="freedraw",
    key="canvas",
)

# Predict button
if st.button('Predict'):
    if canvas_result.image_data is None:
        st.write("Please draw a digit first!")
    else:
        # Convert the canvas image to grayscale
        img = cv2.cvtColor(canvas_result.image_data.astype('uint8'), cv2.COLOR_BGR2GRAY)
        
        # Resize to 28x28 pixels, the input size for the model
        img_resized = cv2.resize(img, (28, 28))
        
        # Invert the image (white background, black digit)
        img_resized = cv2.bitwise_not(img_resized)
        
        # Normalize the image
        img_resized = img_resized / 255.0
        
        # Reshape for the model: (1, 28, 28, 1)
        img_resized = img_resized.reshape(1, 28, 28, 1)
        
        # Predict the digit
        prediction = model.predict(img_resized)
        predicted_digit = np.argmax(prediction)
        
        # Display the prediction
        st.write(f"Predicted Digit: {predicted_digit}")

# Clear button
if st.button('Clear'):
    st.experimental_rerun()