Update app.py

app.py

@@ -4,7 +4,7 @@ import numpy as np
 from PIL import Image
 import os
 
-# === Fix font/matplotlib warnings ===
+# === Fix font/matplotlib warnings for Hugging Face ===
 os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
 os.environ["XDG_CACHE_HOME"] = "/tmp"
 
@@ -14,7 +14,8 @@ 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)
@@ -26,7 +27,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():
@@ -41,35 +42,51 @@ 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 ===
+# === Title ===
 st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
 
+# === 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", ".png"))])
+example_files = sorted([
+    f for f in os.listdir(example_dir)
+    if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
+])
 
-# Display examples in columns
-cols = st.columns(len(example_files))
+if example_files:
+    st.subheader("🖼️ Try with an Example Image")
+    cols = st.columns(min(len(example_files), 4))  # up to 4 per row
 
-for i, filename in enumerate(example_files):
-    with cols[i]:
-        img_path = os.path.join(example_dir, filename)
+    for i, file in enumerate(example_files):
+        img_path = os.path.join(example_dir, file)
         example_img = Image.open(img_path)
-        st.image(example_img, caption=filename, use_column_width=True)
+        with cols[i % len(cols)]:
+            if st.button(file, key=file):
+                uploaded_image = img_path  # simulate upload
+                image = example_img.convert("RGB")
+                st.image(image, caption=f"Example: {file}", use_column_width=True)
+
+# === Prediction ===
+if uploaded_image:
+    if isinstance(uploaded_image, str):
+        image = Image.open(uploaded_image).convert("RGB")
+    else:
+        image = Image.open(uploaded_image).convert("RGB")
 
-        if st.button(f"Run on {filename}"):
-            st.subheader("Original Image")
-            st.image(example_img, use_column_width=True)
+    st.image(image, caption="Input Image", use_column_width=True)
 
-            st.subheader("Predicted Mask")
-            result = run_prediction(example_img)
-            st.image(result, use_column_width=True)
+    # Preprocess and predict
+    resized = 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}")
+
+    # Normalize output
+    y_norm = (y - y.min()) / (y.max() - y.min() + 1e-6)
+    mask = (y_norm * 255).astype(np.uint8)
+    result = Image.fromarray(mask)
 
+    st.image(result, caption="Predicted Segmentation", use_column_width=True)