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

# Load the saved model
@st.cache_resource
def load_model():
    model = tf.keras.models.load_model('pneumonia_cnn_model.keras')  
    return model

model = load_model()

st.title("🫁 Pneumonia Detection from Chest X-ray Images")

st.markdown("Upload your own X-ray or try one of the sample images below.")

# === Sample Image Section ===
sample_images = {
    "Choose a sample image": None,
    "🧍 Normal Sample": "samples/normal.jpg",
    "🤒 Pneumonia Sample": "samples/pneumonia.jpg"
}

selected_sample = st.selectbox("📂 Select a sample image", list(sample_images.keys()))
uploaded_file = st.file_uploader("📁 Or upload a chest X-ray image...", type=["jpg", "jpeg", "png"])

# Determine which image to use
if selected_sample != "Choose a sample image":
    image_path = sample_images[selected_sample]
    image = Image.open(image_path).convert("RGB")
    st.image(image, caption=f'🖼️ {selected_sample}', use_column_width=True)
elif uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption='🖼️ Uploaded Image', use_column_width=True)
else:
    image = None

# === Predict Button ===
if image and st.button('🔍 Predict'):
    img = image.resize((150, 150))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    prediction = model.predict(img_array)

    if prediction[0][0] > 0.5:
        st.error("🩺 **Prediction: Pneumonia Detected**")
    else:
        st.success("✅ **Prediction: Normal**")