Update app.py

app.py

@@ -1,59 +1,55 @@
 import gradio as gr
-import cv2
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.models import load_model
+import cv2
 
-# Constants
-IMG_HEIGHT, IMG_WIDTH = 256, 256
+# Load your trained Keras model
+model = tf.keras.models.load_model("model.h5")
 
-# Custom metrics
-def iou_metric(y_true, y_pred):
-    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)
+# Image preprocessing function (same as used during training)
+def preprocess_image(img):
+    img_resized = cv2.resize(img, (256, 256))
+    img_normalized = img_resized / 255.0  # Normalize to 0-1
+    return img_normalized
 
-def dice_metric(y_true, y_pred):
-    y_pred = tf.cast(y_pred > 0.5, tf.float32)
-    intersection = tf.reduce_sum(y_true * y_pred)
-    return (2. * intersection) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + 1e-7)
+# Prediction and overlay function
+def predict(input_img):
+    # Ensure image is RGB and numpy array
+    img = np.array(input_img.convert("RGB"))
 
-# Load model
-model = load_model("unet_mask_segmentation.h5", custom_objects={'iou_metric': iou_metric, 'dice_metric': dice_metric})
+    # Preprocess
+    preprocessed_img = preprocess_image(img)
+    input_tensor = np.expand_dims(preprocessed_img, axis=0)  # Add batch dimension
 
-# Predict mask function
-def predict(image):
-    # Convert PIL to NumPy array
-    img_np = np.array(image)
-    img_resized = cv2.resize(img_np, (IMG_WIDTH, IMG_HEIGHT)) / 255.0
-    img_input = np.expand_dims(img_resized, axis=0)
+    # Model prediction
+    prediction = model.predict(input_tensor)[0]  # Remove batch dim
 
-    # Predict mask
-    mask_pred = model.predict(img_input)[0]
-    mask_binary = (mask_pred > 0.5).astype(np.uint8) * 255
+    # Post-processing mask
+    mask = (prediction > 0.5).astype(np.uint8)  # Binary mask
+    mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)
 
     # Create overlay
-    overlay = img_resized.copy()
-    overlay[mask_binary == 255] = [1.0, 0, 0]  # Red overlay in normalized scale
+    overlay = img.astype(np.float32) / 255.0  # Normalize input image
+    alpha = 0.5  # Transparency of overlay
+
+    # Create red mask in RGB format
+    red_mask = np.zeros_like(overlay)
+    red_mask[:, :, 0] = mask_resized  # Red channel
 
-    # Convert overlay and mask to displayable format
-    overlay_display = (overlay * 255).astype(np.uint8)
-    mask_display = mask_binary.astype(np.uint8)
+    # Alpha blend original image with red mask
+    blended = (1 - alpha) * overlay + alpha * red_mask
+    blended = np.clip(blended * 255, 0, 255).astype(np.uint8)
 
-    return overlay_display, mask_display
+    return blended
 
-# Gradio Interface
+# Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Image(type="numpy", label="Overlay"),
-        gr.Image(type="numpy", label="Predicted Mask")
-    ],
-    title="Face Mask Segmentation",
-    description="Upload a face image to get the predicted segmentation mask using U-Net."
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Image(type="numpy", label="Segmented Image"),
+    title="Image Segmentation App",
+    description="Upload an image and get the segmentation mask overlay using your trained model."
 )
 
-if __name__ == "__main__":
-    interface.launch()
+# Launch Gradio app (enable public link for Hugging Face Spaces)
+interface.launch(share=True)