Update app.py

app.py

@@ -1,45 +1,45 @@
-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)
-        
+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)