File size: 6,596 Bytes
e18f153
 
 
 
 
 
 
0da14c5
e18f153
a84ccd1
0da14c5
 
 
 
 
 
 
 
 
 
 
 
 
a84ccd1
 
0da14c5
e18f153
a84ccd1
0da14c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e18f153
 
0da14c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c33e07b
0da14c5
 
 
 
 
 
 
 
 
 
 
 
 
c33e07b
a84ccd1
 
0da14c5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# This file is used to visualize bounding boxes on an image
from urllib.parse import urlparse
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import requests
from typing import List
from functools import cache
import matplotlib.colors as colors

DEFAULTS = {
    'bbox_outline_width': 2,
    # color name or hex code or tuple of RGBA or tuple of RGB or tuple (color_name, alpha)
    # between 0 (fully transparent) and 255 (fully opaque)
    'bbox_outline_color': ('blue', 123),
    # color name or hex code or tuple of RGBA or tuple of RGB or tuple (color_name, alpha)
    # between 0 (fully transparent) and 255 (fully opaque)
    'bbox_fill_color': ('red', 50),
    'label_text_color': "black",
    'label_fill_color': "red",
    'label_text_padding': 0,
    'label_rectangle_left_margin': 0,
    'label_rectangle_top_margin': 0,
    'label_text_size': 12,
}


@cache
def get_font(path_or_url: str = 'https://github.com/googlefonts/roboto/raw/main/src/hinted/Roboto-Regular.ttf', size: int = DEFAULTS['label_text_size']):
    if urlparse(path_or_url).scheme in ["http", "https"]:  # Online
        return ImageFont.truetype(requests.get(path_or_url, stream=True).raw, size=size)
    else:  # Local
        return ImageFont.truetype(path_or_url, size=size)

named_colors_mapping = colors.get_named_colors_mapping()
@cache
def get_color(color: str | tuple) -> tuple | str:
    if isinstance(color, tuple):
        if len(color) == 2:
            real_color, alpha = (color[0], int(color[1]))
            if colors.is_color_like(real_color):
                real_color_rgb = colors.hex2color(named_colors_mapping.get(real_color, real_color))
                if len(real_color_rgb) == 3:
                    real_color_alpha = (np.array(real_color_rgb, dtype=int) * 255).tolist() + [alpha]
                    return tuple(real_color_alpha)
    return color

def visualize_bboxes_on_image(
        image: Image.Image,
        bboxes: List[List[int]],
        labels: List[str] = None,
        bbox_outline_width=DEFAULTS["bbox_outline_width"],
        bbox_outline_color=DEFAULTS["bbox_outline_color"],
        bbox_fill_color: str | list[tuple | str] = DEFAULTS["bbox_fill_color"],
        label_text_color: str | list[tuple |
                                     str] = DEFAULTS["label_text_color"],
        label_fill_color=DEFAULTS["label_fill_color"],
        label_text_padding=DEFAULTS["label_text_padding"],
        label_rectangle_left_margin=DEFAULTS["label_rectangle_left_margin"],
        label_rectangle_top_margin=DEFAULTS['label_rectangle_top_margin'],
        label_text_size=DEFAULTS["label_text_size"],
        convert_to_x0y0x1y1=None) -> Image.Image:
    '''
    Visualize bounding boxes on an image
    Args:
      image: Image to visualize
      bboxes: List of bounding boxes
      labels: Titles of the bounding boxes
      bbox_outline_width: Width of the bounding box
      bbox_outline_color: Color of the bounding box
      bbox_fill_color: Fill color of the bounding box
      label_text_color: Color of the label text
      label_fill_color: Color of the label rectangle
      label_text_padding: Padding of the label text
      label_rectangle_left_margin: Left padding of the label rectangle
      label_rectangle_top_margin: Top padding of the label rectangle
      label_text_size: Font size of the label text
      convert_to_x0y0x1y1: Function to convert bounding box to x0y0x1y1 format
    Returns:
      Image: Image annotated with bounding boxes
    '''
    image = image.copy().convert("RGB")
    draw = ImageDraw.Draw(image)
    font = get_font(size=label_text_size)
    labels = (labels or []) + np.full(len(bboxes) -
                                      len(labels or []), None).tolist()
    bbox_fill_colors = bbox_fill_color if isinstance(bbox_fill_color, list) else [
        bbox_fill_color] * len(bboxes)
    bbox_outline_colors = bbox_outline_color if isinstance(
        bbox_outline_color, list) else [bbox_outline_color] * len(bboxes)
    
    for bbox, label, _bbox_fill_color, _bbox_outline_color in zip(bboxes, labels, bbox_fill_colors, bbox_outline_colors):
        x0, y0, x1, y1 = convert_to_x0y0x1y1(
            bbox) if convert_to_x0y0x1y1 is not None else bbox
        _bbox_fill_color = get_color(_bbox_fill_color)
        _bbox_outline_color = get_color(_bbox_outline_color)
        rectangle_image = Image.new('RGBA', image.size)
        rectangle_image_draw = ImageDraw.Draw(rectangle_image)
        rectangle_image_draw.rectangle(
            xy=[x0, y0, x1, y1],
            fill=_bbox_fill_color,
            outline=_bbox_outline_color,
            width=bbox_outline_width)
        image.paste(im=rectangle_image, mask=rectangle_image)

        if label is not None:
            draw_text_on_image(
                draw,
                [x0, y0],
                label,
                label_text_color,
                label_fill_color,
                label_text_padding,
                label_rectangle_left_margin,
                label_rectangle_top_margin,
                label_text_size,
                font)
    return image


def draw_text_on_image(
        image_or_draw: Image.Image | ImageDraw.ImageDraw,
        text_position_xy: List[int],
        label: str,
        label_text_color=DEFAULTS["label_text_color"],
        label_fill_color=DEFAULTS["label_fill_color"],
        label_text_padding=DEFAULTS["label_text_padding"],
        label_rectangle_left_margin=DEFAULTS["label_rectangle_left_margin"],
        label_rectangle_top_margin=DEFAULTS['label_rectangle_top_margin'],
        label_text_size=DEFAULTS["label_text_size"],
        font: ImageFont.FreeTypeFont = None) -> Image.Image:
    is_image = isinstance(image_or_draw, Image.Image)
    image = image_or_draw.copy().convert("RGB") if is_image else None
    font = font or get_font(size=label_text_size)
    x0, y0 = text_position_xy
    text_position = (x0 - label_rectangle_left_margin + label_text_padding,
                     y0 - label_rectangle_top_margin + label_text_padding)
    draw = ImageDraw.Draw(image) if is_image else image_or_draw
    _, _, text_bbox_right, text_bbox_bottom = draw.textbbox(
        text_position, label, font=font)
    xy = [
        text_position[0] - label_text_padding,
        text_position[1] - label_text_padding,
        text_bbox_right + label_text_padding + label_text_padding,
        text_bbox_bottom + label_text_padding + label_text_padding
    ]
    draw.rectangle(xy, fill=label_fill_color)
    draw.text(text_position, label, font=font, fill=label_text_color)
    return image