import gradio as gr
from ultralyticsplus import YOLO, render_result

# load model
model = YOLO("keremberke/yolov8m-hard-hat-detection")

# set model parameters
model.overrides["conf"] = 0.25  # NMS confidence threshold
model.overrides["iou"] = 0.45  # NMS IoU threshold
model.overrides["agnostic_nms"] = False  # NMS class-agnostic
model.overrides["max_det"] = 1000  # maximum number of detections per image


def get_result(img):
    results = model.predict(img)
    return render_result(model=model, image=img, result=results[0])


title = "Hard Hat Detector"
description = "Upload an image to identify who is wearing a hard hat and who is not."
examples = ["https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg"]

iface = gr.Interface(
    title=title,
    description=description,
    examples=examples,
    fn=get_result,
    inputs=gr.components.Image(shape=(512, 512)),
    outputs=gr.components.Image(shape=(512, 512)),
)
iface.launch()