Update app.py

app.py

@@ -1,15 +1,13 @@
+import gradio as gr
 import cv2
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-import matplotlib.pyplot as plt
-import os
 
-# Define image dimensions
-IMG_HEIGHT = 256
-IMG_WIDTH = 256
+# Constants
+IMG_HEIGHT, IMG_WIDTH = 256, 256
 
-# Load trained model
+# 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)
@@ -21,50 +19,41 @@ def dice_metric(y_true, y_pred):
     intersection = tf.reduce_sum(y_true * y_pred)
     return (2. * intersection) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + 1e-7)
 
+# Load model
 model = load_model("unet_mask_segmentation.h5", custom_objects={'iou_metric': iou_metric, 'dice_metric': dice_metric})
 
-# Function to preprocess input image
-def preprocess_image(image_path):
-    img = cv2.imread(image_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
-    img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT)) / 255.0  # Resize and normalize
-    return np.expand_dims(img, axis=0)  # Add batch dimension
-
-# Function to make predictions
-def predict_mask(image_path):
-    img = preprocess_image(image_path)
-    mask_pred = model.predict(img)[0]  # Remove batch dimension
-    mask_pred = (mask_pred > 0.5).astype(np.uint8)  # Threshold mask
-    return mask_pred
-
-# Function to visualize results
-def visualize_results(image_path, mask_pred):
-    img = cv2.imread(image_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert to RGB
-    img_resized = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
-
-    fig, ax = plt.subplots(1, 2, figsize=(12, 4))
-    ax[0].imshow(img_resized)
-    ax[0].set_title("Original Image")
-    ax[0].axis("off")
-
-    ax[1].imshow(mask_pred, cmap="gray")
-    ax[1].set_title("Predicted Mask")
-    ax[1].axis("off")
-
-    # overlay = img_resized.copy()
-    # overlay[mask_pred == 1] = [255, 0, 0]  # Apply red overlay on mask
-    # ax[2].imshow(overlay)
-    # ax[2].set_title("Overlay")
-    # ax[2].axis("off")
-
-    plt.show()
-
-# Test on an image
-TEST_IMAGE = "/content/MSFD/1/face_crop/000020_1.jpg"  # Change this to your test image path
-if os.path.exists(TEST_IMAGE):
-    mask_prediction = predict_mask(TEST_IMAGE)
-    visualize_results(TEST_IMAGE, mask_prediction)
-else:
-    print(f"Test image '{TEST_IMAGE}' not found. Please check the path.")
-
+# 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)
+
+    # Predict mask
+    mask_pred = model.predict(img_input)[0]
+    mask_binary = (mask_pred > 0.5).astype(np.uint8) * 255
+
+    # Create overlay
+    overlay = img_resized.copy()
+    overlay[mask_binary == 255] = [1.0, 0, 0]  # Red overlay in normalized scale
+
+    # Convert overlay and mask to displayable format
+    overlay_display = (overlay * 255).astype(np.uint8)
+    mask_display = mask_binary.astype(np.uint8)
+
+    return overlay_display, mask_display
+
+# 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."
+)
+
+if __name__ == "__main__":
+    interface.launch()