osherr commited on
Commit
8277b50
ยท
verified ยท
1 Parent(s): 7fe897d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +45 -38
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
- "final_model_after_third_iteration_WDL0.07_0.5155/",
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
- # App title
45
  st.title("๐Ÿ•ณ๏ธ Sinkhole Segmentation with EffV2-UNet")
46
 
47
- # Sidebar: Upload or select image
48
- st.sidebar.header("๐Ÿ“ Image Input")
49
- uploaded_file = st.sidebar.file_uploader("Upload your image", type=["jpg", "jpeg", "png", "tif", "tiff"])
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.sidebar.markdown("Or select from examples:")
62
- for file in example_files:
63
- path = os.path.join(example_dir, file)
64
- image = Image.open(path)
65
- st.sidebar.image(image, caption=file, width=120)
66
- if st.sidebar.button(f"Use {file}"):
67
- st.session_state.selected_example = path
68
- uploaded_file = None # clear uploaded file
69
-
70
- # Load image for display and processing
71
- input_image = None
72
- if uploaded_file is not None:
73
- input_image = Image.open(uploaded_file).convert("RGB")
 
 
74
  elif st.session_state.selected_example is not None:
75
- input_image = Image.open(st.session_state.selected_example).convert("RGB")
 
 
 
76
 
77
- if input_image:
78
- st.image(input_image, caption="Input Image", use_column_width=True)
 
 
79
 
80
- threshold = st.slider("Confidence Threshold", 0.0, 1.0, 0.5, step=0.01)
 
 
81
 
82
- if st.button("๐Ÿ” Run Segmentation"):
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
- st.text(f"Prediction min/max: {y.min():.5f} / {y.max():.5f}")
 
 
88
 
89
- # Threshold and visualize
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)