app.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image
from ultralytics import YOLO

# Load YOLO model
model = YOLO("best.pt")

def segment_image(input_image):
 # Resize the input image to 255x255
 img = np.array(input_image)
 img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

 # Perform object detection and segmentation
 # Perform object detection and segmentation
 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()
 print(overlay_image.shape, mask.shape )

 # 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

iface = gr.Interface(
 fn=segment_image,
 inputs=gr.components.Image(type="pil", label="Upload an image"),
 outputs=gr.components.Image(type="numpy", label="Segmented Image"),
 title="Aorta segmentation and Detection using YOLOv8 😃",
 description= 'This software generates the segementation mask for Aorta for the Point of Care Ultrasound (POCUS) images')

iface.launch()