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()