engie_pipeline/utils.py

from PIL import ImageDraw


def draw_boxes(image, boxes, probs, color):
    draw = ImageDraw.Draw(image)

    for box, proba in zip(boxes, probs):
        draw.rectangle(box.tolist(), outline=color, width=6)

        draw.text(
            (box[0], box[1] - 15),
            text=str(round(proba.item(), 2)),
            fill=color,
            stroke_width=1,
        )

    return image