Spaces:
Running
Running
| 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 | |