Update app.py

app.py

@@ -1,44 +1,71 @@
-import gradio as gr
+import cv2
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-import cv2
+import matplotlib.pyplot as plt
+import os
 
-# Define the IOU metric exactly as provided
+# Define image dimensions
+IMG_HEIGHT = 256
+IMG_WIDTH = 256
+
+# Load trained model
 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)
 
-# Load the trained U-Net model
-model = load_model("unet_mask_segmentation.h5", custom_objects={'iou_metric': iou_metric})
-
-# Preprocess the uploaded image for prediction
-def preprocess_image(image, target_size=(256, 256)):
-    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 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]  # Remove batch dimension
-    binary_mask = (pred_mask > 0.5).astype(np.uint8) * 255
-    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
-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()
+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)
+
+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.")
+
+isse bana de app.py