Spaces:
Sleeping
Sleeping
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 === | |
model_path = "final_model_after_third_iteration_WDL0.07_0.5155/" | |
def load_model(): | |
return tf.keras.models.load_model( | |
model_path, | |
custom_objects={ | |
"weighted_dice_loss": weighted_dice_loss, | |
"iou_metric": iou_metric, | |
"bce_loss": bce_loss | |
} | |
) | |
model = load_model() | |
# === Title === | |
st.title("🕳️ SinkSAM-Net - Self Supervised Sinkhole segmentation") | |
# === 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 example_files: | |
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: | |
active_image = st.session_state.selected_example | |
# === Confidence threshold slider === | |
threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01) | |
# === Prediction === | |
if active_image: | |
image = Image.open(active_image).convert("RGB") | |
st.image(image, caption="Input Image", use_container_width=True) | |
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}") | |
# Apply threshold | |
mask_bin = (y > threshold).astype(np.uint8) * 255 | |
mask_image = Image.fromarray(mask_bin) | |
st.image(mask_image, caption=f"Segmentation (Threshold = {threshold:.2f})", use_container_width=True) | |