Spaces:
Running
Running
File size: 1,734 Bytes
90b4364 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
from typing import List, Tuple
from PIL import Image, ImageFont, ImageDraw
from hbutils.color import rnd_colors, Color
def plot_detection(pil_img: Image.Image, detection: List[Tuple[Tuple[float, float, float, float], int, float]],
captions: List[str], text_padding: int = 5, font: ImageFont.ImageFont = None,
no_label: bool = False):
new_img = pil_img.copy()
draw = ImageDraw.Draw(new_img, mode='RGBA')
_colors = list(map(str, rnd_colors(len(captions))))
_caption_map = {i: name for i, name in enumerate(captions)}
for (xmin, ymin, xmax, ymax), label, score in detection:
box_color = _colors[label]
draw.rectangle((xmin, ymin, xmax, ymax), outline=box_color, width=2)
if not no_label:
label_text = _caption_map.get(label, str(label))
label_text = f'{label_text}: {score * 100:.2f}%'
font = font or ImageFont.load_default()
_t_x0, _t_y0, _t_x1, _t_y1 = draw.textbbox((xmin, ymin), label_text, font=font)
_t_width, _t_height = _t_x1 - _t_x0, _t_y1 - _t_y0
if ymin - _t_height - text_padding < 0:
_t_text_rect = (xmin, ymin, xmin + _t_width + text_padding * 2, ymin + _t_height + text_padding * 2)
_t_text_co = (xmin + text_padding, ymin + text_padding)
else:
_t_text_rect = (xmin, ymin - _t_height - text_padding * 2, xmin + _t_width + text_padding * 2, ymin)
_t_text_co = (xmin + text_padding, ymin - _t_height - text_padding)
draw.rectangle(_t_text_rect, fill=str(Color(box_color, alpha=0.5)))
draw.text(_t_text_co, label_text, fill="black", font=font)
return new_img
|