import gradio as gr
from ultralytics import YOLO
import cv2

model = YOLO('head_yolov8s.pt')

def head_detect(path):
    img = cv2.imread(path)
    output = model(source=img)
    res = output[0].cpu().numpy()
    
    # Extract bbox, cls id and conf
    bboxes = res.boxes.xyxy
    class_ids = res.boxes.cls
    conf_scores = res.boxes.conf

    for i in range(len(bboxes)):
        xmin, ymin, xmax, ymax = int(bboxes[i][0]), int(bboxes[i][1]), int(bboxes[i][2]), int(bboxes[i][3])
        conf = conf_scores[i]
        cls_id = int(class_ids[i])
        label = model.names[cls_id]  # Get the label name

        # Draw rectangle for bounding box
        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)

        # Prepare label text with confidence score
        label_text = f'{label} {conf:.2f}'

        # Put text (label) on the image
        cv2.putText(img, label_text, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 1, lineType=cv2.LINE_AA)

    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

input_image = [
    gr.components.Image(type='filepath', label='Input Image'),
]

output_image = [
    gr.components.Image(type='numpy', label='Prediction'),
]

interface = gr.Interface(
    fn=head_detect,
    inputs=input_image,
    outputs=output_image,
    title='Top Head Detection',
    cache_examples=False,
)

gr.TabbedInterface(
    [interface],
    tab_names=['Image Inference']
).queue().launch()