Create main.py

main.py

@@ -0,0 +1,43 @@
+import gradio as gr
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+import cv2
+
+# Custom metric (if any, otherwise comment this out)
+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
+model = load_model("unet_mask_segmentation.h5", custom_objects={'iou_metric': iou_metric})
+
+# Preprocess image
+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)
+
+# Predict mask
+def segment_image(input_image):
+    preprocessed = preprocess_image(input_image)
+    pred_mask = model.predict(preprocessed)[0]
+    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
+
+# Gradio interface
+interface = gr.Interface(
+    fn=segment_image,
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=gr.Image(type="numpy", label="Segmented Mask"),
+    title="Image Segmentation with U-Net",
+    description="Upload an image to see the segmentation mask predicted by the U-Net model."
+)
+
+if __name__ == "__main__":
+    interface.launch()