import gradio as gr
import cv2
from ultralytics import YOLO

def run(source):
    global model
    res = model(source, conf=.5, iou=.5)
    res_plotted = res[0].plot()
    # converting BGR to RGB
    result = cv2.cvtColor(res_plotted, cv2.COLOR_BGR2RGB)
    return result
    
model = YOLO("yolov8n-nckh2023.pt") # Select YOLO model

gr.Interface(
    run,
    inputs=gr.Image(label="Upload image", type="filepath"),
    outputs=gr.Image(label="Your result"),
    title="Motorcyclist, helmet, and license plate detection",
).launch()