import streamlit as st
import numpy as np
import joblib
from PIL import Image, ImageOps

st.title('Handwritten Digit Recognizer')

# Load the model
try:
    model = joblib.load('src/digit_rf_model.joblib')
except Exception as e:
    st.error(f"Error loading model: {e}")

uploaded_file = st.file_uploader("Upload a digit image (28x28 grayscale)", type=["png", "jpg", "jpeg"])

def preprocess_image(img):
    # Convert to grayscale, resize to 28x28, flatten
    img = ImageOps.grayscale(img)
    img = img.resize((28, 28))
    arr = np.array(img).reshape(1, -1)
    return arr

if uploaded_file is not None:
    try:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        input_data = preprocess_image(image)
        if st.button('Predict Digit'):
            prediction = model.predict(input_data)
            st.success(f'Predicted Digit: {int(prediction[0])}')
    except Exception as e:
        st.error(f"Error processing image or making prediction: {e}")