import random
from PIL import ImageDraw, Image
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

# Load model
weight_file = hf_hub_download("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
model = YOLO(weight_file)  # pretrained YOLOv8n model

# Helper Functions for Plotting BBoxes
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    width, height = img.size
    tl = line_thickness or round(0.002 * (width + height) / 2) + 1  # line/font thickness
    color = color or (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    img_draw = ImageDraw.Draw(img)
    img_draw.rectangle((c1[0], c1[1], c2[0], c2[1]), outline=color, width=tl)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        x1, y1, x2, y2 = img_draw.textbbox(c1, label, stroke_width=tf)
        img_draw.rectangle((x1, y1, x2, y2), fill=color)
        img_draw.text((x1, y1), label, fill=(255, 255, 255))


# Ploting Bounding Box on img
def add_bboxes(pil_img, result, confidence=0.6):
    for box in result.boxes:
        [cl] = box.cls.tolist()
        [conf] = box.conf.tolist()
        if conf < confidence:
            continue
        [rect] = box.xyxy.tolist()
        text = f'{result.names[cl]}: {conf: 0.2f}'
        plot_one_box(x=rect, img=pil_img, label=text)

    return pil_img


def detect(im, confidence):
    results = model(source=im)
    res_img = im
    for result in results:
        res_img = add_bboxes(res_img, result, confidence)
    return res_img

if __name__ == '__main__':
    im = Image.open("./tests/fire1.jpg")
    results = model(source=im)
    for result in results:
        im = add_bboxes(im, result)
    im.show()