models/tools/draw.py

import random
from PIL import ImageDraw


def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    Helper Functions for Plotting BBoxes
    :param x:
    :param img:
    :param color:
    :param label:
    :param line_thickness:
    :return:
    """
    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))


def add_bboxes(pil_img, result, confidence=0.6):
    """
    Plotting Bounding Box on img
    :param pil_img:
    :param result:
    :param confidence:
    :return:
    """
    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 add_bboxes2(pil_img, result, confidence=0.6):
    """
    Plotting Bounding Box on img
    :param pil_img:
    :param result:
    :param confidence:
    :return:
    """
    for box in result['boxes']:
        cl = box['cls']
        conf = box['conf']
        if conf < confidence:
            continue
        rect = box['xyxy']
        text = f'{cl}: {conf: 0.2f}'
        plot_one_box(x=rect, img=pil_img, label=text)

    return pil_img