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

# Load Models
@st.cache_resource
def load_models():
    mri_model = tf.keras.models.load_model("updatedmri_tumor_classifier.keras")
    ultrasound_model = tf.keras.models.load_model("uls_tumor_classifier.keras")
    mammogram_model = tf.keras.models.load_model("updated_mammogram_classifier.keras")
    return mri_model, ultrasound_model, mammogram_model

mri_model, ultrasound_model, mammogram_model = load_models()

# Class mappings
mri_labels = ['Sick', 'Healthy']
ultrasound_labels = ['benign', 'malignant', 'normal']
mammogram_labels = ['Malignant', 'Benign']

# Image preprocessing function
def preprocess_image(image):
    image = image.resize((128, 128))
    image = np.array(image) / 255.0
    return np.expand_dims(image, axis=0)

# Streamlit UI
st.title("Breast Cancer Classification Web App")
st.write("Upload a breast scan and select the type to detect if it is healthy or has a benign tumor.")

# User input
image_type = st.selectbox("Select the image type", ["MRI", "Ultrasound", "Mammogram"])
uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])

if uploaded_image:
    image = Image.open(uploaded_image).convert("RGB")
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Preprocess image
    processed_image = preprocess_image(image)

    # Model prediction
    if image_type == "MRI":
        prediction = mri_model.predict(processed_image)
        label = mri_labels[np.argmax(prediction)]
        if label == "Healthy":
            st.success("This image is healthy")
        else:
            st.success("This image is not healthy")
    elif image_type == "Ultrasound":
        prediction = ultrasound_model.predict(processed_image)
        label = ultrasound_labels[np.argmax(prediction)]
        if label == "normal":
            st.success("This image is healthy")
        elif label == "benign":
            st.success("This image has a benign tumor")
        elif label == "malignant":
            st.success("This image has a malignant tumor")

    elif image_type == "Mammogram":
        prediction = mammogram_model.predict(processed_image)
        label = mammogram_labels[np.argmax(prediction)]
        if label == "Benign":
            st.success("This image has a benign tumor")
        else:
            st.success("This image has a malignant tumor")
            
    # Display result
   # if label == "Healthy" or "normal":
   #     st.success("This image is healthy.")
 #   elif label == "benign" or "Benign":
   #     st.success("This image has a benign tumor.")
 #   elif label == "malignant" or "Malignant":
     #   st.success("This image has a malignant tumor.")
    #elif label == "Sick":
    #    st.success("This image is not healthy")