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

# Load a pretrained YOLOv8n model
model = YOLO("yolov8n.pt")


def main_function(lbd, img):
    results = model(img)  # predict on an image
    r = results[0]
    im_bgr = r.plot()  # BGR-order numpy array
    im_rgb = Image.fromarray(im_bgr[..., ::-1])  # RGB-order PIL image
    new_img = im_rgb
    # res = results[0].save(filename="output.jpg")  # save the image
    # # load image
    # new_img = Image.open("output.jpg")
    return new_img


iface = gr.Interface(
    fn=main_function,
    inputs=["slider", gr.Image(type="pil")],
    outputs=gr.Image(type="pil"),
    examples=[
        [0, "bus.jpg"],
    ],
)
iface.launch()