|
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) |