import streamlit as st
import tensorflow as tf
import numpy as np

from PIL import Image, ImageOps


st.set_page_config(
    page_title="MNIST Digit Recognition",
    page_icon="🧮"
)


st.title("🧮 MNIST Digit Recognition")

st.write(
    "Upload a handwritten digit image and the model will predict the digit (0-9)."
)


# Load Model

@st.cache_resource
def load_model():

    model = tf.keras.models.load_model(
        "mnist_digit_recognizer.keras"
    )

    return model


model = load_model()



uploaded_file = st.file_uploader(
    "Upload Digit Image",
    type=["png","jpg","jpeg"]
)



if uploaded_file is not None:

    image = Image.open(
        uploaded_file
    ).convert("L")


    # Resize to MNIST size

    image = image.resize(
        (28,28)
    )


    # Invert colors if needed

    image = ImageOps.invert(
        image
    )


    st.image(
        image,
        caption="Uploaded Image",
        width=150
    )


    # Preprocessing

    img_array = np.array(
        image
    )


    img_array = img_array / 255.0


    img_array = img_array.reshape(
        1,
        28,
        28,
        1
    )


    # Prediction

    prediction = model.predict(
        img_array
    )


    digit = np.argmax(
        prediction
    )


    confidence = np.max(
        prediction
    ) * 100


    st.success(
        f"Predicted Digit: {digit}"
    )


    st.info(
        f"Confidence: {confidence:.2f}%"
    )


    st.write(
        "Prediction probabilities:"
    )


    for i, probability in enumerate(prediction[0]):

        st.write(
            f"{i}: {probability*100:.2f}%"
        )