osherr commited on
Commit
73b3af8
·
verified ·
1 Parent(s): 495ad2b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +62 -0
app.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import tensorflow as tf
3
+ import numpy as np
4
+ from PIL import Image
5
+ import os
6
+
7
+ # === Fix font/matplotlib warnings for Hugging Face ===
8
+ os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
9
+ os.environ["XDG_CACHE_HOME"] = "/tmp"
10
+
11
+ # === Custom loss and metrics ===
12
+ def weighted_dice_loss(y_true, y_pred):
13
+ smooth = 1e-6
14
+ y_true_f = tf.reshape(y_true, [-1])
15
+ y_pred_f = tf.reshape(y_pred, [-1])
16
+ intersection = tf.reduce_sum(y_true_f * y_pred_f)
17
+ return 1 - ((2. * intersection + smooth) /
18
+ (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth))
19
+
20
+ def iou_metric(y_true, y_pred):
21
+ y_true = tf.cast(y_true > 0.5, tf.float32)
22
+ y_pred = tf.cast(y_pred > 0.5, tf.float32)
23
+ intersection = tf.reduce_sum(y_true * y_pred)
24
+ union = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) - intersection
25
+ return intersection / (union + 1e-6)
26
+
27
+ def bce_loss(y_true, y_pred):
28
+ return tf.keras.losses.binary_crossentropy(y_true, y_pred)
29
+
30
+ # === Load model ===
31
+ model_path = "final_model_after_third_iteration_WDL0.07_0.5155/"
32
+ @st.cache_resource
33
+ def load_model():
34
+ return tf.keras.models.load_model(
35
+ model_path,
36
+ custom_objects={
37
+ "weighted_dice_loss": weighted_dice_loss,
38
+ "iou_metric": iou_metric,
39
+ "bce_loss": bce_loss
40
+ }
41
+ )
42
+
43
+ model = load_model()
44
+
45
+ # === Streamlit UI ===
46
+ st.title("🕳️ Sinkhole Segmentation with EffV2-UNet")
47
+
48
+ uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
49
+ if uploaded_image:
50
+ image = Image.open(uploaded_image).convert("RGB")
51
+ st.image(image, caption="Original Image", use_column_width=True)
52
+
53
+ # Preprocess and predict
54
+ resized = image.resize((512, 512))
55
+ x = np.expand_dims(np.array(resized) / 255.0, axis=0)
56
+ y = model.predict(x)[0, :, :, 0]
57
+
58
+ y_norm = (y - y.min()) / (y.max() - y.min() + 1e-6)
59
+ mask = (y_norm * 255).astype(np.uint8)
60
+ result = Image.fromarray(mask)
61
+
62
+ st.image(result, caption="Predicted Segmentation", use_column_width=True)