# 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