Update app.py

app.py

@@ -4,11 +4,11 @@ import numpy as np
 from PIL import Image
 import os
 
-# === Fix font/matplotlib warnings for Hugging Face ===
+# Fix font/matplotlib warnings for Hugging Face
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
 os.environ["XDG_CACHE_HOME"] = "/tmp"
 
-# === Custom loss and metrics ===
+# Custom loss and metrics
 def weighted_dice_loss(y_true, y_pred):
     smooth = 1e-6
     y_true_f = tf.reshape(y_true, [-1])
@@ -27,12 +27,11 @@ def iou_metric(y_true, y_pred):
 def bce_loss(y_true, y_pred):
     return tf.keras.losses.binary_crossentropy(y_true, y_pred)
 
-# === Load model ===
-model_path = "final_model_after_third_iteration_WDL0.07_0.5155/"
+# Load model
 @st.cache_resource
 def load_model():
     return tf.keras.models.load_model(
-        model_path,
+        "final_model_after_third_iteration_WDL0.07_0.5155/",
         custom_objects={
             "weighted_dice_loss": weighted_dice_loss,
             "iou_metric": iou_metric,
@@ -42,57 +41,52 @@ def load_model():
 
 model = load_model()
 
-# === Title ===
+# App title
 st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
 
-# === Confidence threshold and predict trigger ===
-st.sidebar.header("Segmentation Settings")
-threshold = st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
+# Sidebar: Upload or select image
+st.sidebar.header("📁 Image Input")
+uploaded_file = st.sidebar.file_uploader("Upload your image", type=["jpg", "jpeg", "png", "tif", "tiff"])
 
-# === Image input section ===
-uploaded_image = st.file_uploader("📤 Upload an image", type=["png", "jpg", "jpeg", "tif", "tiff"])
-
-# === Example selector with preview ===
 example_dir = "examples"
 example_files = sorted([
     f for f in os.listdir(example_dir)
     if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
 ])
 
-selected_example_path = None
+if "selected_example" not in st.session_state:
+    st.session_state.selected_example = None
 
 if example_files:
-    st.subheader("🖼️ Try with an Example Image")
-    cols = st.columns(min(len(example_files), 4))
-    for i, file in enumerate(example_files):
-        img_path = os.path.join(example_dir, file)
-        img_preview = Image.open(img_path).convert("RGB").resize((128, 128))
-        with cols[i % len(cols)]:
-            st.image(img_preview, caption=file, use_column_width=True)
-            if st.button(f"Use {file}", key=file):
-                selected_example_path = img_path
-
-# === Set image to process ===
-if selected_example_path:
-    uploaded_image = selected_example_path
-
-# === Run prediction if button clicked ===
-if uploaded_image:
-    if isinstance(uploaded_image, str):
-        image = Image.open(uploaded_image).convert("RGB")
-    else:
-        image = Image.open(uploaded_image).convert("RGB")
-
-    st.image(image, caption="Input Image", use_column_width=True)
-
-    if st.button("Run Segmentation"):
-        resized = image.resize((512, 512))
+    st.sidebar.markdown("Or select from examples:")
+    for file in example_files:
+        path = os.path.join(example_dir, file)
+        image = Image.open(path)
+        st.sidebar.image(image, caption=file, width=120)
+        if st.sidebar.button(f"Use {file}"):
+            st.session_state.selected_example = path
+            uploaded_file = None  # clear uploaded file
+
+# Load image for display and processing
+input_image = None
+if uploaded_file is not None:
+    input_image = Image.open(uploaded_file).convert("RGB")
+elif st.session_state.selected_example is not None:
+    input_image = Image.open(st.session_state.selected_example).convert("RGB")
+
+if input_image:
+    st.image(input_image, caption="Input Image", use_column_width=True)
+
+    threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
+
+    if st.button("🔍 Run Segmentation"):
+        resized = input_image.resize((512, 512))
         x = np.expand_dims(np.array(resized), axis=0)
         y = model.predict(x)[0, :, :, 0]
 
         st.text(f"Prediction min/max: {y.min():.5f} / {y.max():.5f}")
 
-        mask_bin = (y > threshold).astype(np.uint8) * 255
-        mask_image = Image.fromarray(mask_bin)
-
-        st.image(mask_image, caption=f"Segmentation Mask (Threshold = {threshold:.2f})", use_column_width=True)
+        # Threshold and visualize
+        mask = (y > threshold).astype(np.uint8) * 255
+        mask_image = Image.fromarray(mask)
+        st.image(mask_image, caption=f"Segmentation (Threshold = {threshold:.2f})", use_column_width=True)