Update app.py

app.py

@@ -4,33 +4,34 @@ import tensorflow as tf
 from tensorflow.keras.models import load_model
 import cv2
 
-# Custom metric (if any, otherwise comment this out)
+# Define the IOU metric exactly as provided
 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-7)
 
-# Load the model
+# Load the trained U-Net model
 model = load_model("unet_mask_segmentation.h5", custom_objects={'iou_metric': iou_metric})
 
-# Preprocess image
+# Preprocess the uploaded image for prediction
 def preprocess_image(image, target_size=(256, 256)):
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-    image = cv2.resize(image, target_size)
-    image = image / 255.0
-    return np.expand_dims(image, axis=0)
+    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert RGB -> BGR for OpenCV ops
+    image_resized = cv2.resize(image_bgr, target_size)
+    image_resized = image_resized / 255.0  # Normalize
+    return np.expand_dims(image_resized, axis=0)  # Add batch dimension
 
-# Predict mask
+# Predict and return the segmented mask
 def segment_image(input_image):
+    original_size = (input_image.shape[1], input_image.shape[0])  # (width, height)
     preprocessed = preprocess_image(input_image)
-    pred_mask = model.predict(preprocessed)[0]
+    pred_mask = model.predict(preprocessed)[0]  # Remove batch dimension
     binary_mask = (pred_mask > 0.5).astype(np.uint8) * 255
-    binary_mask = cv2.resize(binary_mask, (input_image.shape[1], input_image.shape[0]))
-    return binary_mask
+    binary_mask_resized = cv2.resize(binary_mask, original_size)  # Resize mask to original
+    binary_mask_rgb = cv2.cvtColor(binary_mask_resized, cv2.COLOR_GRAY2RGB)  # Convert to 3-channel RGB
+    return binary_mask_rgb
 
-# Gradio interface
+# Gradio Interface
 interface = gr.Interface(
     fn=segment_image,
     inputs=gr.Image(type="numpy", label="Upload Image"),