Create app.py

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import numpy as np
+import cv2
+from PIL import Image
+from ultralytics import YOLO
+
+# Define available YOLO models
+available_models = {
+    "Model 1 (best.pt)": YOLO("best.pt"),
+    "Model 2 (another_model.pt)": YOLO("another_model.pt"),
+    # Add more models as needed
+}
+
+# Create a function to perform image segmentation using the selected model
+def segment_image(input_image, selected_model):
+    # Resize the input image to 255x255
+    img = np.array(input_image)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    # Perform object detection and segmentation using the selected model
+    model = available_models[selected_model]
+    results = model(img)
+    mask = results[0].masks.data.numpy()
+    target_height = img.shape[0]
+    target_width = img.shape[1]
+
+    # Resize the mask using OpenCV
+    resized_mask = cv2.resize(mask[0], (target_width, target_height))
+    resized_mask = (resized_mask * 255).astype(np.uint8)
+
+    # Create a copy of the original image
+    overlay_image = img.copy()
+
+    # Apply the resized mask to the overlay image
+    overlay_image[resized_mask > 0] = [100, 0, 0]  # Overlay in green
+
+    # Convert the overlay image to PIL format
+    overlay_pil = Image.fromarray(overlay_image)
+
+    return overlay_pil
+
+# Create the Gradio interface with a dropdown for model selection
+iface = gr.Interface(
+    fn=segment_image,
+    inputs=[
+        gr.inputs.Image(type="pil", label="Upload an image"),
+        gr.inputs.Dropdown(
+            choices=list(available_models.keys()),
+            label="Select YOLO Model",
+            default="Model 1 (best.pt)"
+        )
+    ],
+    outputs=gr.outputs.Image(type="numpy", label="Segmented Image"),
+    title="Aorta segmentation and Detection using YOLOv8 😃",
+    description='This software generates the segmentation mask for Aorta for Point of Care Ultrasound (POCUS) images'
+)
+
+iface.launch()