import streamlit as st
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
import base64

st.set_page_config(
    page_title="Pneumonia Detection App",
    page_icon=":lungs:",
    layout="wide"
)

custom_style = """
<style>
    div[data-testid="stToolbar"],
    div[data-testid="stDecoration"],
    div[data-testid="stStatusWidget"],
    #MainMenu,
    header,
    footer {
        visibility: hidden;
        height: 0%;
    }
</style>

"""
st.markdown(custom_style, unsafe_allow_html=True)

model = tf.keras.models.load_model('fixed_model.h5')
img_size = (224, 224)

def preprocess_image(img):
    img = image.load_img(img, target_size=img_size)
    img_array = image.img_to_array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

def predict_image(img):
    img_array = preprocess_image(img)
    prediction = np.squeeze(model.predict(img_array), axis=0)
    return prediction

def display_image_with_download(image_path, caption, download_text):
    image = Image.open(image_path)
    st.image(image, caption=caption, use_column_width=True)
    
    with open(image_path, 'rb') as f:
        data = f.read()
        base64_data = base64.b64encode(data).decode('utf-8')
        href = f'<a href="data:application/octet-stream;base64,{base64_data}" download="{download_text}.jpg">Download {download_text}</a>'
        st.markdown(href, unsafe_allow_html=True)

def main():
    st.title("Pneumonia Detection")

    uploaded_file = st.file_uploader("Upload a chest X-ray image...", type=["jpg", "png", "jpeg"])

    st.markdown("""
        Example Instructions:
        - Upload a chest X-ray image in JPG format.
        - Or, download sample images below and check the predictions.
    """)

    st.write("**Download/View Sample Images:**")
    
    normal_download = st.button("View Normal Image")
    pneumonic_download = st.button("View Pneumonic Image")

    if normal_download:
        normal_image_path = "test-normal_001.jpg"
        display_image_with_download(normal_image_path, "Normal Image", "Normal Image")

    if pneumonic_download:
        pneumonic_image_path = "test-pneumonia_028.jpg"
        display_image_with_download(pneumonic_image_path, "Pneumonic Image", "Pneumonic Image")

    if uploaded_file is not None:
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
        st.markdown('</div>', unsafe_allow_html=True)

        prediction = predict_image(uploaded_file)

        st.write("**Prediction:**")
        class_label = "Pneumonia" if prediction > 0.5 else "Normal"
        st.write(f"The image is classified as **{class_label}**.")

if __name__ == "__main__":
    main()