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

# Define available YOLO models
available_models = {
    "X-ray": YOLO("xray.pt"),
    "CT scan": YOLO("CT.pt"),
    "Ultrasound": YOLO("ultrasound.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="X-ray"
        )
    ],
    outputs=gr.outputs.Image(type="numpy", label="Segmented Image"),
    title="YOLOv8 with SAM 😃",
    description='This software generates the segmentation mask for Aorta for Point of Care Ultrasound (POCUS) images'
)

iface.launch()