import streamlit as st
from PIL import Image
import numpy as np
import tensorflow as tf

# Hide Streamlit UI elements
hide_streamlit_style = """
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
footer:after {
    content:'This app is in its early stage. We recommend you to seek professional advice from a dermatologist. Thank you.';
    visibility: visible;
    display: block;
    position: relative;
    padding: 5px;
    top: 2px;
}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)

# Center-align elements
st.markdown(
    """
    <style>
    .center { text-align: center; }
    </style>
    """,
    unsafe_allow_html=True
)

# Load the HDF5 model
model = tf.keras.models.load_model('best_model (1).h5')  # Updated to .h5

# Class labels
labels = {
    0: 'Chickenpox',
    1: 'Cowpox',
    2: 'HFMD',
    3: 'Healthy',
    4: 'Measles',
    5: 'MPOX'
}

def preprocess_image(image):
    image = image.resize((224, 224))
    image_array = np.array(image)
    image_array = np.expand_dims(image_array, axis=0)
    return image_array

def predict(image):
    processed_image = preprocess_image(image)
    prediction = model.predict(processed_image)
    label_index = np.argmax(prediction)
    predicted_label = labels[label_index]
    confidence = prediction[0][label_index] * 100
    return predicted_label, confidence

def main():
    st.markdown("<h1 class='center'>Skin Lesion Classifier</h1>", unsafe_allow_html=True)
    number = st.radio('Pick one', ['Upload from gallery', 'Capture by camera'])
    
    uploaded_file = None
    if number == 'Capture by camera':
        uploaded_file = st.camera_input("Take a picture")
    else:
        uploaded_file = st.file_uploader("Choose an image", type=["jpg", "png", "jpeg", "bmp"])
    
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        predicted_label, confidence = predict(image)
        
        st.markdown("<h3 class='center'>This might be:</h3>", unsafe_allow_html=True)
        st.markdown(f"<h1 class='center'>{predicted_label}</h1>", unsafe_allow_html=True)
        st.markdown(f"<p class='center'>Confidence: {confidence:.2f}%</p>", unsafe_allow_html=True)

if __name__ == '__main__':
    main()