app.py

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

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

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

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=[
            "images/00054_145.jpg",
            "images/00054_164.jpg",
            'images/00058_179.jpg',
            'images/00061_146.jpg'
        ],
        
        inputs=img_input,
        outputs=img_output,
        fn=segment_image,
        label="Example Images"
    )

demo.launch(inbrowser=True)