import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# Use caching to speed up model loading
@st.cache_resource
def load_model(model_path):
    return tf.keras.models.load_model(model_path, compile=False)

# Load the model once using cache
MODEL_PATH = r"D:\AMIT\amit\ODC\W4\D2\chest_xray\chest_x-ray_model.keras"
model = load_model(MODEL_PATH)

def preprocess_image(image, target_size=(256, 256)):
    try:
        image = image.convert("L")  # Convert to grayscale
        image = image.resize(target_size)
        image = np.array(image) / 255.0  # Normalize to [0, 1]
        image = np.expand_dims(image, axis=-1)  # Add grayscale channel dimension
        image = np.expand_dims(image, axis=0)  # Add batch dimension
        return image
    except Exception as e:
        raise ValueError(f"Error processing image: {str(e)}")

def plot_prediction(confidence):
    # Create a bar chart to show the prediction confidence
    labels = ['Healthy', 'Infected']
    values = [1 - confidence, confidence]
    
    fig, ax = plt.subplots()
    ax.bar(labels, values, color=['green', 'red'])
    ax.set_ylim(0, 1)
    ax.set_ylabel('Confidence')
    ax.set_title('Prediction Confidence')
    
    st.pyplot(fig)

st.title("Chest X-ray Diagnosis")
st.write("Upload a chest X-ray image for diagnosis.")

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

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)
    
    if st.button("Predict"):
        with st.spinner('Processing...'):
            processed_image = preprocess_image(image)
        
            predictions = model.predict(processed_image)
            confidence = float(predictions[0][0])
            predicted_class = "Infected" if confidence > 0.5 else "Healthy"
        
            st.write(f"### Result: {predicted_class}")
            st.write(f"### Confidence: {confidence:.2f}")
            
            # Add visualization
            plot_prediction(confidence)