import streamlit as st
import numpy as np
import cv2
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model

@st.cache_resource
def load_unet_model():
    return load_model('best_unet_model.keras')

model = load_unet_model()

def preprocess_image(image):
    image = image.resize((256, 256))
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image

def predict_mask(image):
    processed_image = preprocess_image(image)
    predicted_mask = model.predict(processed_image)
    predicted_mask = (predicted_mask > 0.5).astype(np.uint8)
    return predicted_mask[0, :, :, 0]

st.title('Medical Image Segmentation with U-Net (Mohamed Arbi Nsibi)')
st.subheader("Note: The model's segmentation accuracy is not that accurate because of the small training dataset. Larger and more diverse data could improve performance ")

uploaded_file = st.file_uploader("Choose 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)
    
    if st.button('Segment Image'):
        mask = predict_mask(image)
        
        st.image(mask * 255, caption='Segmentation Mask', use_column_width=True)
        
        overlay = np.zeros((256, 256, 3), dtype=np.uint8)
        overlay[:,:,1] = mask * 255  
        original_resized = np.array(image.resize((256, 256)))
        overlayed_image = cv2.addWeighted(original_resized, 0.7, overlay, 0.3, 0)
        
        st.image(overlayed_image, caption='Segmentation Overlay', use_column_width=True)