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