Update app.py

app.py

@@ -4,7 +4,7 @@ import numpy as np
 from PIL import Image
 import os
 
-# === Fix font/matplotlib warnings for Hugging Face ===
+# === Fix font/matplotlib warnings ===
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
 os.environ["XDG_CACHE_HOME"] = "/tmp"
 
@@ -14,8 +14,7 @@ def weighted_dice_loss(y_true, y_pred):
     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) /
-                (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + 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):
     y_true = tf.cast(y_true > 0.5, tf.float32)
@@ -27,7 +26,7 @@ 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():
@@ -42,25 +41,35 @@ def load_model():
 
 model = load_model()
 
+# === Inference Function ===
+def run_prediction(image):
+    image = image.convert("RGB").resize((512, 512))
+    x = np.expand_dims(np.array(image), axis=0)
+    y = model.predict(x)[0, :, :, 0]
+    y_norm = (y - y.min()) / (y.max() - y.min() + 1e-6)
+    mask = (y_norm * 255).astype(np.uint8)
+    return Image.fromarray(mask)
+
 # === Streamlit UI ===
 st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
 
-uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
-if uploaded_image:
-    image = Image.open(uploaded_image).convert("RGB")
-    st.image(image, caption="Original Image", use_column_width=True)
+example_dir = "examples"
+example_files = sorted([f for f in os.listdir(example_dir) if f.lower().endswith((".jpg", ".png"))])
 
-    # Preprocess and predict (NO normalization here)
-    resized = image.resize((512, 512))
-    x = np.expand_dims(np.array(resized), axis=0)
-    y = model.predict(x)[0, :, :, 0]
+# Display examples in columns
+cols = st.columns(len(example_files))
 
-    # Print min/max to logs
-    st.text(f"Prediction min/max: {y.min():.5f} / {y.max():.5f}")
+for i, filename in enumerate(example_files):
+    with cols[i]:
+        img_path = os.path.join(example_dir, filename)
+        example_img = Image.open(img_path)
+        st.image(example_img, caption=filename, use_column_width=True)
 
-    # Contrast stretch
-    y_norm = (y - y.min()) / (y.max() - y.min() + 1e-6)
-    mask = (y_norm * 255).astype(np.uint8)
-    result = Image.fromarray(mask)
+        if st.button(f"Run on {filename}"):
+            st.subheader("Original Image")
+            st.image(example_img, use_column_width=True)
+
+            st.subheader("Predicted Mask")
+            result = run_prediction(example_img)
+            st.image(result, use_column_width=True)
 
-    st.image(result, caption="Predicted Segmentation", use_column_width=True)