Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -4,17 +4,17 @@ 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):
|
@@ -27,11 +27,12 @@ def iou_metric(y_true, y_pred):
|
|
27 |
def bce_loss(y_true, y_pred):
|
28 |
return tf.keras.losses.binary_crossentropy(y_true, y_pred)
|
29 |
|
30 |
-
# Load model
|
|
|
31 |
@st.cache_resource
|
32 |
def load_model():
|
33 |
return tf.keras.models.load_model(
|
34 |
-
|
35 |
custom_objects={
|
36 |
"weighted_dice_loss": weighted_dice_loss,
|
37 |
"iou_metric": iou_metric,
|
@@ -41,52 +42,58 @@ def load_model():
|
|
41 |
|
42 |
model = load_model()
|
43 |
|
44 |
-
#
|
45 |
st.title("๐ณ๏ธ Sinkhole Segmentation with EffV2-UNet")
|
46 |
|
47 |
-
#
|
48 |
-
|
49 |
-
|
50 |
|
|
|
|
|
|
|
|
|
51 |
example_dir = "examples"
|
52 |
example_files = sorted([
|
53 |
f for f in os.listdir(example_dir)
|
54 |
if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
|
55 |
])
|
56 |
|
57 |
-
if "selected_example" not in st.session_state:
|
58 |
-
st.session_state.selected_example = None
|
59 |
-
|
60 |
if example_files:
|
61 |
-
st.
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
68 |
-
|
69 |
-
|
70 |
-
|
71 |
-
|
72 |
-
|
73 |
-
|
|
|
|
|
74 |
elif st.session_state.selected_example is not None:
|
75 |
-
|
|
|
|
|
|
|
76 |
|
77 |
-
|
78 |
-
|
|
|
|
|
79 |
|
80 |
-
|
|
|
|
|
81 |
|
82 |
-
|
83 |
-
resized = input_image.resize((512, 512))
|
84 |
-
x = np.expand_dims(np.array(resized), axis=0)
|
85 |
-
y = model.predict(x)[0, :, :, 0]
|
86 |
|
87 |
-
|
|
|
|
|
88 |
|
89 |
-
|
90 |
-
mask = (y > threshold).astype(np.uint8) * 255
|
91 |
-
mask_image = Image.fromarray(mask)
|
92 |
-
st.image(mask_image, caption=f"Segmentation (Threshold = {threshold:.2f})", use_column_width=True)
|
|
|
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):
|
|
|
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,
|
|
|
42 |
|
43 |
model = load_model()
|
44 |
|
45 |
+
# === Title ===
|
46 |
st.title("๐ณ๏ธ Sinkhole Segmentation with EffV2-UNet")
|
47 |
|
48 |
+
# === Session state for selected example ===
|
49 |
+
if "selected_example" not in st.session_state:
|
50 |
+
st.session_state.selected_example = None
|
51 |
|
52 |
+
# === File uploader ===
|
53 |
+
uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg", "tif", "tiff"])
|
54 |
+
|
55 |
+
# === Example selector ===
|
56 |
example_dir = "examples"
|
57 |
example_files = sorted([
|
58 |
f for f in os.listdir(example_dir)
|
59 |
if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
|
60 |
])
|
61 |
|
|
|
|
|
|
|
62 |
if example_files:
|
63 |
+
st.subheader("๐ผ๏ธ Try with an Example Image")
|
64 |
+
cols = st.columns(min(len(example_files), 4))
|
65 |
+
|
66 |
+
for i, file in enumerate(example_files):
|
67 |
+
with cols[i % len(cols)]:
|
68 |
+
img_path = os.path.join(example_dir, file)
|
69 |
+
example_img = Image.open(img_path)
|
70 |
+
st.image(example_img, caption=file, use_container_width=True)
|
71 |
+
if st.button(f"Run Segmentation", key=file):
|
72 |
+
st.session_state.selected_example = img_path
|
73 |
+
|
74 |
+
# === Determine active image ===
|
75 |
+
active_image = None
|
76 |
+
if uploaded_image is not None:
|
77 |
+
active_image = uploaded_image
|
78 |
elif st.session_state.selected_example is not None:
|
79 |
+
active_image = st.session_state.selected_example
|
80 |
+
|
81 |
+
# === Confidence threshold slider ===
|
82 |
+
threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
|
83 |
|
84 |
+
# === Prediction ===
|
85 |
+
if active_image:
|
86 |
+
image = Image.open(active_image).convert("RGB")
|
87 |
+
st.image(image, caption="Input Image", use_container_width=True)
|
88 |
|
89 |
+
resized = image.resize((512, 512))
|
90 |
+
x = np.expand_dims(np.array(resized), axis=0)
|
91 |
+
y = model.predict(x)[0, :, :, 0]
|
92 |
|
93 |
+
st.text(f"Prediction min/max: {y.min():.5f} / {y.max():.5f}")
|
|
|
|
|
|
|
94 |
|
95 |
+
# Apply threshold
|
96 |
+
mask_bin = (y > threshold).astype(np.uint8) * 255
|
97 |
+
mask_image = Image.fromarray(mask_bin)
|
98 |
|
99 |
+
st.image(mask_image, caption=f"Segmentation (Threshold = {threshold:.2f})", use_container_width=True)
|
|
|
|
|
|