Update app.py

app.py

@@ -4,17 +4,17 @@ 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])
     y_pred_f = tf.reshape(y_pred, [-1])
     intersection = tf.reduce_sum(y_true_f * y_pred_f)
-    return 1 - ((2. * intersection + smooth) /
+    return 1 - ((2. * intersection + smooth) / 
                 (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth))
 
 def iou_metric(y_true, y_pred):
@@ -27,11 +27,12 @@ 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
+# === Load model ===
+model_path = "final_model_after_third_iteration_WDL0.07_0.5155/"
 @st.cache_resource
 def load_model():
     return tf.keras.models.load_model(
-        "final_model_after_third_iteration_WDL0.07_0.5155/",
+        model_path,
         custom_objects={
             "weighted_dice_loss": weighted_dice_loss,
             "iou_metric": iou_metric,
@@ -41,52 +42,58 @@ def load_model():
 
 model = load_model()
 
-# App title
+# === Title ===
 st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
 
-# 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"])
+# === Session state for selected example ===
+if "selected_example" not in st.session_state:
+    st.session_state.selected_example = None
 
+# === File uploader ===
+uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "tif", "tiff"])
+
+# === Example selector ===
 example_dir = "examples"
 example_files = sorted([
     f for f in os.listdir(example_dir)
     if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
 ])
 
-if "selected_example" not in st.session_state:
-    st.session_state.selected_example = None
-
 if example_files:
-    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")
+    st.subheader("🖼️ Try with an Example Image")
+    cols = st.columns(min(len(example_files), 4))
+
+    for i, file in enumerate(example_files):
+        with cols[i % len(cols)]:
+            img_path = os.path.join(example_dir, file)
+            example_img = Image.open(img_path)
+            st.image(example_img, caption=file, use_container_width=True)
+            if st.button(f"Run Segmentation", key=file):
+                st.session_state.selected_example = img_path
+
+# === Determine active image ===
+active_image = None
+if uploaded_image is not None:
+    active_image = uploaded_image
 elif st.session_state.selected_example is not None:
-    input_image = Image.open(st.session_state.selected_example).convert("RGB")
+    active_image = st.session_state.selected_example
+
+# === Confidence threshold slider ===
+threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
 
-if input_image:
-    st.image(input_image, caption="Input Image", use_column_width=True)
+# === Prediction ===
+if active_image:
+    image = Image.open(active_image).convert("RGB")
+    st.image(image, caption="Input Image", use_container_width=True)
 
-    threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
+    resized = image.resize((512, 512))
+    x = np.expand_dims(np.array(resized), axis=0)
+    y = model.predict(x)[0, :, :, 0]
 
-    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}")
 
-        st.text(f"Prediction min/max: {y.min():.5f} / {y.max():.5f}")
+    # Apply threshold
+    mask_bin = (y > threshold).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask_bin)
 
-        # 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)
+    st.image(mask_image, caption=f"Segmentation (Threshold = {threshold:.2f})", use_container_width=True)