Update app.py

app.py

@@ -3,15 +3,14 @@ import tensorflow as tf
 import numpy as np
 from PIL import Image
 
-st.title("🌊 Flood Segmentation with U-Net (TF SavedModel)")
+st.title("🌊 Flood Segmentation with U-Net")
 
 @st.cache_resource
 def load_model():
-    model = tf.saved_model.load("unet_savedmodel")
-    infer = model.signatures["serving_default"]
-    return infer
+    model = tf.keras.models.load_model("unet_savedmodel", compile=False)
+    return model
 
-infer = load_model()
+model = load_model()
 
 uploaded = st.file_uploader("Upload flood image", type=["jpg", "jpeg", "png"])
 
@@ -22,9 +21,9 @@ if uploaded:
     x = np.array(img_resized, dtype=np.float32) / 255.0
     x = np.expand_dims(x, axis=0)
 
-    pred = infer(tf.constant(x))["conv2d_9"].numpy()[0, :, :, 0]
+    pred = model.predict(x)[0, :, :, 0]
 
     mask = (pred > 0.5).astype(np.uint8) * 255
 
     st.image(img, caption="Original Image")
-    st.image(mask, caption="Predicted Mask", clamp=True)
+    st.image(mask, caption="Segmentation Mask", clamp=True)