Create app.py

app.py

@@ -0,0 +1,86 @@
+import numpy as np
+import tensorflow as tf
+import cv2
+from PIL import Image
+import gradio as gr
+import os
+
+# Define custom metrics
+smooth = 1e-15
+
+def dice_coef(y_true, y_pred):
+    y_true = tf.keras.layers.Flatten()(y_true)
+    y_pred = tf.keras.layers.Flatten()(y_pred)
+    intersection = tf.reduce_sum(y_true * y_pred)
+    return (2. * intersection + smooth) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + smooth)
+
+def dice_loss(y_true, y_pred):
+    return 1.0 - dice_coef(y_true, y_pred)
+
+# Load the model from the same directory as the script
+model_filename = "model.h5"  # Replace with your actual model filename
+model_path = os.path.join(os.path.dirname(__file__), model_filename) # Construct the full path
+
+
+def load_model(model_path):
+    try:
+        model = tf.keras.models.load_model(model_path, custom_objects={'dice_loss': dice_loss, 'dice_coef': dice_coef}, compile=False)
+        return model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
+
+# Perform inference
+model = load_model(model_path)
+
+def perform_inference(image):
+    if model is None:
+        print("Model not loaded properly.")
+        return None, None, None
+        
+     # Preprocess the image
+    original_shape = image.shape[:2]    
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image_resized = cv2.resize(image, (256, 256))
+    image_normalized = image_resized / 255.0
+    image_expanded = np.expand_dims(image_normalized, axis=0)
+
+    # Get the mask from the model prediction
+    mask = model.predict(image_expanded)[0]
+    mask_resized = cv2.resize(mask, (original_shape[1], original_shape[0]))
+    mask_binary = (mask_resized > 0.5).astype(np.uint8) * 255
+
+
+    # Apply the mask to the original image (for better visualization)
+    heatmap_img = cv2.applyColorMap(mask_binary, cv2.COLORMAP_JET)
+    segmented_image = cv2.addWeighted(image, 0.7, heatmap_img, 0.3, 0)
+
+    #Convert back to BGR
+    segmented_image = cv2.cvtColor(segmented_image, cv2.COLOR_RGB2BGR)
+
+    # Convert results to PIL Images
+    return (Image.fromarray(image),
+            Image.fromarray(mask_binary.astype(np.uint8)), #Mask is already multiplied by 255
+            Image.fromarray(segmented_image))
+
+
+# Gradio app
+def gradio_app():
+    with gr.Blocks() as demo:
+        with gr.Row():
+          with gr.Column():
+              input_image = gr.Image(label="Input Image", type="numpy")
+              submit_btn = gr.Button("Submit")
+
+          with gr.Column():
+              original_image_output = gr.Image(label="Original Image")
+              mask_output = gr.Image(label="Predicted Mask")
+              segmented_image_output= gr.Image(label="Segmented Image")
+
+
+        submit_btn.click(perform_inference, inputs=input_image, outputs=[original_image_output, mask_output, segmented_image_output])
+
+    demo.launch()
+
+if __name__ == "__main__":
+    gradio_app()