import streamlit as st
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import CustomObjectScope
from tensorflow.keras.initializers import glorot_uniform
from PIL import Image
import numpy as np
import os

# Define models and their validation accuracies
model_options = {
    "Model1": {
        "path": "model_name.h5",
        "accuracy": 50
    },
    "Model2": {
        "path": "pneu_cnn_model.h5",
        "accuracy": 76
    }
}

# Load the model with custom objects if necessary
def load_model_with_custom_objects(model_path):
    if not os.path.isfile(model_path):
        raise FileNotFoundError(f"Model file not found: {model_path}")

    custom_objects = {
        'GlorotUniform': glorot_uniform
        # Add other custom objects if needed
    }

    try:
        with CustomObjectScope(custom_objects):
            model = load_model(model_path)
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        raise

    return model

# Image preprocessing (adjust as needed for your model)
def preprocess_image(image):
    # Convert image to grayscale and resize
    image = image.convert('L')  # Convert to grayscale if necessary
    image = image.resize((64, 64))  # Resize to match the model input shape
    image_array = np.array(image)
    image_array = image_array.astype('float32') / 255.0  # Normalize
    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
    image_array = np.expand_dims(image_array, axis=-1)  # Add channel dimension if needed
    return image_array

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

    model_name = st.selectbox("Select a model", list(model_options.keys()))
    model_path = model_options[model_name]["path"]
    model_accuracy = model_options[model_name]["accuracy"]

    # Load the selected model
    try:
        model = load_model_with_custom_objects(model_path)
    except Exception as e:
        st.error(f"Failed to load model: {e}")
        return

    st.write(f"Model: {model_name}")
    st.write(f"Validation Accuracy: {model_accuracy}%")

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

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        # Perform prediction using the model
        img_array = preprocess_image(image)
        
        try:
            prediction = model.predict(img_array)
            predicted_class = "Pneumonia" if prediction[0][0] > 0.5 else "Normal"
            st.write(f"Prediction: {predicted_class}")
        except Exception as e:
            st.error(f"Error during prediction: {str(e)}")

if __name__ == "__main__":
    main()