File size: 4,305 Bytes
e18f153 a84ccd1 e18f153 a84ccd1 e18f153 a84ccd1 e18f153 a84ccd1 e18f153 a84ccd1 e18f153 |
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 |
# 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
DEFAULTS = {
'width': 2,
'bbox_color': "red",
'label_text_color': "black",
'label_rectangle_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)
def visualize_bboxes_on_image(
image: Image.Image,
bboxes: List[List[int]],
titles: List[str] = None,
width = DEFAULTS["width"],
bbox_color = DEFAULTS["bbox_color"],
label_text_color = DEFAULTS["label_text_color"],
label_rectangle_color = DEFAULTS["label_rectangle_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
titles: Titles of the bounding boxes
width: Width of the bounding box
bbox_color: Color of the bounding box
label_text_color: Color of the label text
label_rectangle_color: Color of the label rectangle
convert_to_x0y0x1y1: Function to convert bounding box to x0y0x1y1 format
label_text_padding: Padding of the label text
label_rectangle_left_padding: Left padding of the label rectangle
label_rectangle_top_padding: Top padding of the label rectangle
label_text_size: Font size of the label text
Returns:
Image: Image with bounding boxes
'''
image = image.copy().convert("RGB")
draw = ImageDraw.Draw(image)
font = get_font(size = label_text_size)
titles = (titles or []) + np.full(len(bboxes) - len(titles or []), None).tolist()
for bbox, title in zip(bboxes, titles):
x0, y0, x1, y1 = convert_to_x0y0x1y1(bbox) if convert_to_x0y0x1y1 is not None else bbox
draw.rectangle([x0, y0, x1, y1], outline=bbox_color, width=width)
if title is not None:
draw_text_on_image(
draw,
[x0, y0],
title,
label_text_color,
label_rectangle_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],
title: str,
label_text_color = DEFAULTS["label_text_color"],
label_rectangle_color = DEFAULTS["label_rectangle_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_left, text_bbox_top, text_bbox_right, text_bbox_bottom = draw.textbbox(text_position, title, 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_rectangle_color)
draw.text(text_position, title, font=font, fill=label_text_color)
return image |