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

model = YOLO("runs/segment/train/weights/best.pt")

# Prediction function for image
def segment_image(image):
    results = model(image)[0]
    return np.array(results.plot())

# Prediction function for video
def segment_video(video_path):
    results = model(video_path)
    return results[0].save(save_dir="segmented_video/") or "segmented_video/predict.mp4"


with gr.Blocks(theme=gr.themes.Soft(),title='YOLOv11 Image Segmentation') as demo:
    gr.Markdown("# YOLOv11 Image Segmentation")
    with gr.Row():
        img_input = gr.Image(type="pil", label="Upload an Image")
        img_output = gr.Image(label="Segmented Image")
    img_button = gr.Button("Segment Image")
    img_button.click(segment_image, inputs=img_input, outputs=img_output)

    gr.Examples(
        examples=[
            "datasets/coco8-seg/images/train/000000000025.jpg",
            "datasets/coco8-seg/images/train/000000000009.jpg"
        ],
        inputs=img_input,
        outputs=img_output,
        fn=segment_image,
        label="Example Images"
    )

demo.launch(inbrowser=True)