import streamlit as st
from PIL import Image
import numpy as np
from tensorflow.keras.models import load_model

# Define a function to make predictions on a single image
def predict_single_image(image, model):
    try:
        # Resize and preprocess the image
        img = image.resize((128, 128))
        img = np.array(img) / 255.0  # Normalization
        img = np.expand_dims(img, axis=0)  # Add batch dimension

        # Make prediction using the provided model
        prediction = model.predict(img)

        # Thresholding prediction
        threshold = 0.5
        prediction_class = (prediction > threshold).astype(int)

        # Interpret prediction
        if prediction_class == 1:
            return "With Mask"
        else:
            return "Without Mask"
    except Exception as e:
        st.error(f"Error occurred during prediction: {str(e)}")

# Load the model from .h5 file
@st.cache(allow_output_mutation=True)
def load_model_from_h5():
    try:
        return load_model('model.h5')
    except Exception as e:
        st.error(f"Error occurred while loading the model: {str(e)}")

# Streamlit app
def main():
    # App title and description
    st.title("Face Mask Detection App")
    st.write("Upload an image and click 'Predict' to see the prediction.")

    # Custom CSS for styling
    custom_css = """
    <style>
    body {
        background-color: #f0f2f6;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        padding: 10px 15px;
        text-align: center;
        text-decoration: none;
        display: inline-block;
        font-size: 16px;
        margin: 4px 2px;
        cursor: pointer;
        border-radius: 8px;
        border: none;
    }
    .stButton>button:hover {
        background-color: #45a049;
    }
    .stTextInput>div>div>input {
        border-radius: 8px;
    }
    .stTextInput>div>div>input:focus {
        border-color: #4CAF50;
        box-shadow: none;
    }
    </style>
    """
    st.markdown(custom_css, unsafe_allow_html=True)

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

    if uploaded_file is not None:
        # Load the image
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image.', use_column_width=True)

        # Button to make prediction
        if st.button('Predict', key='predict_button'):
            # Load the model from .h5 file
            model_h5 = load_model_from_h5()
            if model_h5:
                # Make predictions using the provided model
                prediction = predict_single_image(image, model_h5)
                if prediction:
                    st.success(f"Prediction: {prediction}")

if __name__ == '__main__':
    main()