import gradio as gr
import numpy as np
import cv2
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

checkpoint = hf_hub_download(repo_id="bulkobubulko/cap", filename="yolo11-tapas.pt")
model = YOLO(checkpoint)
model.fuse()

def plot_bl(results):
    img = results[0].orig_img
    names = results[0].names
    scores = results[0].boxes.conf.numpy()
    classes = results[0].boxes.cls.numpy()
    boxes = results[0].boxes.xyxy.numpy().astype(np.int32)

    num_objects = len(boxes)

    for score, cls, box in zip(scores, classes, boxes):
        cl_label = names[cls]
        label = f"{cl_label}: {score:0.2f}"

        img = cv2.rectangle(img.copy(), (box[0], box[1]), (box[2], box[3]), color=(225, 0, 225), thickness=3)
        img = cv2.putText(img.copy(), label, (box[0], box[1]), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, color=(255, 0, 255), thickness=2)
        img = cv2.putText(img.copy(), f'num caps: {num_objects}', (100, 100), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1.0, color=(255, 0, 255), thickness=2)
    return img 

def predict(frame):
    frame = model(frame)
    frame = plot_bl(frame)
    return frame

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_frame = gr.Image(sources=["webcam"], type="numpy")
        with gr.Column():
            output_img = gr.Image(streaming=True)
        dep = input_frame.stream(predict, [input_frame], [output_img], time_limit=30, stream_every=0.1, concurrency_limit=30)

demo.launch()