diff --git a/README.md b/README.md index 04e78e6..e681063 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,7 @@ + +Fork from https://github.com/nalepae/bounding-box/tree/master. I've set the line width and label size proportional to the image size. Helpfull for large images. I've opened a pull request to the original repo, but it appears to be non-active. Hence I'll use my own fork. Original documentation below. + + # Bounding Box **Bounding Box** is a library to plot pretty bounding boxes with a simple Python API. diff --git a/bounding_box/bounding_box.py b/bounding_box/bounding_box.py index d5f8850..2a8046b 100644 --- a/bounding_box/bounding_box.py +++ b/bounding_box/bounding_box.py @@ -8,6 +8,7 @@ from PIL import ImageFont import numpy as _np from hashlib import md5 as _md5 +import string as _string _LOC = _path.realpath(_path.join(_os.getcwd(),_path.dirname(__file__))) @@ -36,8 +37,26 @@ _DEFAULT_COLOR_NAME = "green" _FONT_PATH = _os.path.join(_LOC, "Ubuntu-B.ttf") -_FONT_HEIGHT = 15 -_FONT = ImageFont.truetype(_FONT_PATH, _FONT_HEIGHT) + +def _hex_to_rgb(color): + if type(color) is not str: + return None + color = color.strip() + if color.startswith("#"): + color = color[1:] + if len(color) == 3: + color = "".join(ch * 2 for ch in color) + if len(color) != 6 or any(ch not in _string.hexdigits for ch in color): + return None + try: + return tuple(int(color[i:i + 2], 16) for i in (0, 2, 4)) + except ValueError: + return None + +def _get_contrasting_rgb(rgb): + r, g, b = rgb + luminance = 0.299 * r + 0.587 * g + 0.114 * b + return (255, 255, 255) if luminance < 186 else (0, 0, 0) def _rgb_to_bgr(color): return list(reversed(color)) @@ -45,8 +64,8 @@ def _rgb_to_bgr(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): - text_image = _FONT.getmask(text) +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) @@ -58,7 +77,7 @@ def _get_label_image(text, 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): +def add(image, left, top, right, bottom, label=None, color=None, size=2): if type(image) is not _np.ndarray: raise TypeError("'image' parameter must be a numpy.ndarray") try: @@ -79,20 +98,53 @@ def add(image, left, top, right, bottom, label=None, color=None): 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.strip() + + if color in _COLOR_NAME_TO_RGB: + box_rgb, text_rgb = _COLOR_NAME_TO_RGB[color] + else: + rgb_color = _hex_to_rgb(color) + if rgb_color is None: + msg = "'color' must be a recognised name or hex value such as #RRGGBB" + raise ValueError(msg) + box_rgb = rgb_color + text_rgb = _get_contrasting_rgb(rgb_color) + + colors = [_rgb_to_bgr(item) for item in (box_rgb, text_rgb)] color, color_text = colors - _cv2.rectangle(image, (left, top), (right, bottom), color, 2) + image_height, image_width, _ = image.shape + + # The line width of the box and the font size are calculated based on the image size and the vis_size + # vis_size is set by the user (0=XS, 1=S, 2=M, 3=L, 4=XL) + if size == 0: # XS + width_factor = 1 / 1200 + font_factor = width_factor * 10 + elif size == 1: # S + width_factor = 1 / 600 + font_factor = width_factor * 8 + elif size == 2: # M + width_factor = 1 / 300 + font_factor = width_factor * 6 + elif size == 3: # L + width_factor = 1 / 150 + font_factor = width_factor * 4 + elif size == 4: # XL + width_factor = 1 / 75 + font_factor = width_factor * 2 + + # calculate the line width of the box + line_width_box = max(round(min(image_height, image_width) * width_factor), 1) + _cv2.rectangle(image, (left, top), (right, bottom), color, line_width_box) if label: - _, image_width, _ = image.shape + + # caluculate the font size + font_height = round(min(image_height, image_width) * font_factor) + font = ImageFont.truetype(_FONT_PATH, font_height) + + label_image = _get_label_image(label, color_text, color, font) - label_image = _get_label_image(label, color_text, color) label_height, label_width, _ = label_image.shape rectangle_height, rectangle_width = 1 + label_height, 1 + label_width