-
Notifications
You must be signed in to change notification settings - Fork 1
/
bounding_box.py
executable file
·127 lines (95 loc) · 4.25 KB
/
bounding_box.py
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
from __future__ import division as _division
from __future__ import print_function as _print_function
import os as _os
import os.path as _path
import numpy as np
import cv2 as _cv2
from PIL import ImageFont
import numpy as _np
from hashlib import md5 as _md5
_LOC = _path.realpath(_path.join(_os.getcwd(),_path.dirname(__file__)))
#https://clrs.cc/
_COLOR_NAME_TO_RGB = dict(
navy=((0, 38, 63), (119, 193, 250)),
blue=((0, 120, 210), (173, 220, 252)),
aqua=((115, 221, 252), (0, 76, 100)),
teal=((15, 205, 202), (5, 5, 5)),
olive=((52, 153, 114), (25, 58, 45)),
green=((0, 204, 84), (15, 64, 31)),
lime=((1, 255, 127), (0, 102, 53)),
yellow=((255, 216, 70), (103, 87, 28)),
orange=((255, 125, 57), (104, 48, 19)),
red=((255, 47, 65), (131, 0, 17)),
maroon=((135, 13, 75), (239, 117, 173)),
fuchsia=((246, 0, 184), (103, 0, 78)),
purple=((179, 17, 193), (241, 167, 244)),
black=((24, 24, 24), (220, 220, 220)),
# gray=((168, 168, 168), (5, 5, 5)),
silver=((220, 220, 220), (5, 5, 5))
)
_COLOR_NAMES = list(_COLOR_NAME_TO_RGB)
_DEFAULT_COLOR_NAME = "green"
_FONT_PATH = _os.path.join(_LOC, "Ubuntu-B.ttf")
_FONT_HEIGHT = 25
_FONT = ImageFont.truetype(_FONT_PATH, _FONT_HEIGHT)
def get_font_with_size(font_size):
return ImageFont.truetype(_FONT_PATH, font_size)
def _rgb_to_bgr(color):
# return list(reversed(color))
return color
def _color_image(image, font_color, background_color):
return background_color + (font_color - background_color) * image / 255
def _get_label_image(text, font_color_tuple_bgr, background_color_tuple_bgr, font):
text_image = font.getmask(text)
shape = list(reversed(text_image.size))
bw_image = np.array(text_image).reshape(shape)
image = [
_color_image(bw_image, font_color, background_color)[None, ...]
for font_color, background_color
in zip(font_color_tuple_bgr, background_color_tuple_bgr)
]
return np.concatenate(image).transpose(1, 2, 0)
def add(image, left, top, right, bottom, label=None, color=None, font=_FONT):
if type(image) is not _np.ndarray:
raise TypeError("'image' parameter must be a numpy.ndarray")
try:
left, top, right, bottom = int(left), int(top), int(right), int(bottom)
except ValueError:
raise TypeError("'left', 'top', 'right' & 'bottom' must be a number")
if label and type(label) is not str:
raise TypeError("'label' must be a str")
if label and not color:
hex_digest = _md5(label.encode()).hexdigest()
color_index = int(hex_digest, 16) % len(_COLOR_NAME_TO_RGB)
color = _COLOR_NAMES[color_index]
if not color:
color = _DEFAULT_COLOR_NAME
if type(color) is not str:
raise TypeError("'color' must be a str")
if color not in _COLOR_NAME_TO_RGB:
msg = "'color' must be one of " + ", ".join(_COLOR_NAME_TO_RGB)
raise ValueError(msg)
colors = [_rgb_to_bgr(item) for item in _COLOR_NAME_TO_RGB[color]]
color, color_text = colors
_cv2.rectangle(image, (left, top), (right, bottom), color, 2)
if label:
_, image_width, _ = image.shape
label_image = _get_label_image(label, color_text, color, font)
label_height, label_width, _ = label_image.shape
rectangle_height, rectangle_width = 1 + label_height, 1 + label_width
rectangle_bottom = top
rectangle_left = max(0, min(left - 1, image_width - rectangle_width))
rectangle_top = rectangle_bottom - rectangle_height
rectangle_right = rectangle_left + rectangle_width
label_top = rectangle_top + 1
if rectangle_top < 0:
rectangle_top = top
rectangle_bottom = rectangle_top + label_height + 1
label_top = rectangle_top
label_left = rectangle_left + 1
label_bottom = label_top + label_height
label_right = label_left + label_width
rec_left_top = (rectangle_left, rectangle_top)
rec_right_bottom = (rectangle_right, rectangle_bottom)
_cv2.rectangle(image, rec_left_top, rec_right_bottom, color, -1)
image[label_top:label_bottom, label_left:label_right, :] = label_image