Update app.py

app.py

@@ -45,30 +45,38 @@ model = load_model()
 # === Title ===
 st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
 
-# === File uploader ===
-uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "tif", "tiff"])
+# === 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)
 
-# === Example selector ===
+# === 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 example_files:
     st.subheader("🖼️ Try with an Example Image")
-    cols = st.columns(min(len(example_files), 4))  # up to 4 per row
-
+    cols = st.columns(min(len(example_files), 4))
     for i, file in enumerate(example_files):
         img_path = os.path.join(example_dir, file)
-        example_img = Image.open(img_path)
+        img_preview = Image.open(img_path).convert("RGB").resize((128, 128))
         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)
+            st.image(img_preview, caption=file, use_column_width=True)
+            if st.button(f"Use {file}", key=file):
+                selected_example_path = img_path
 
-# === Prediction ===
+# === 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")
@@ -77,18 +85,14 @@ if uploaded_image:
 
     st.image(image, caption="Input Image", use_column_width=True)
 
-    # Confidence threshold slider
-    threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
-
-    # Preprocess and predict
-    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 = 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 thresholding if needed
-    mask_bin = (y > threshold).astype(np.uint8) * 255
-    mask_image = Image.fromarray(mask_bin)
+        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)
+        st.image(mask_image, caption=f"Segmentation Mask (Threshold = {threshold:.2f})", use_column_width=True)