Spaces:
Running
Running
| import streamlit as st | |
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import os | |
| # Fix font/matplotlib warnings for Hugging Face | |
| os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib" | |
| os.environ["XDG_CACHE_HOME"] = "/tmp" | |
| # 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) / | |
| (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) | |
| y_pred = tf.cast(y_pred > 0.5, tf.float32) | |
| intersection = tf.reduce_sum(y_true * y_pred) | |
| union = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) - intersection | |
| return intersection / (union + 1e-6) | |
| def bce_loss(y_true, y_pred): | |
| return tf.keras.losses.binary_crossentropy(y_true, y_pred) | |
| # Load model | |
| def load_model(): | |
| return tf.keras.models.load_model( | |
| "final_model_after_third_iteration_WDL0.07_0.5155/", | |
| custom_objects={ | |
| "weighted_dice_loss": weighted_dice_loss, | |
| "iou_metric": iou_metric, | |
| "bce_loss": bce_loss | |
| } | |
| ) | |
| model = load_model() | |
| # App 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"]) | |
| 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") | |
| 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}") | |
| # 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) | |