|
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('unet_model_epoch_29_val_loss_0.5760.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) |