RE_UPLOAD-REBUILD-RESTART

utils/visualize_bboxes_on_image.py

@@ -0,0 +1,177 @@
+# 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, Callable
+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 parse_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 draw_bounding_box(
+        image: Image.Image, 
+        bbox_outline_width: int, 
+        bbox_fill_color: str | list[tuple | str], 
+        bbox_outline_color: str | list[tuple | str], 
+        bbox: List[List[int]],
+        label_rotate_angle: int = 0,
+        mask_callback: Callable[[ImageDraw.ImageDraw], None] = None) -> Image.Image:
+    options = {
+        'xy': bbox,
+        'fill': parse_color(bbox_fill_color) if bbox_fill_color else None,
+        'outline': parse_color(bbox_outline_color) if bbox_outline_color else None,
+        'width': bbox_outline_width
+    }
+    options = {k: v for k, v in options.items() if v is not None}
+    rectangle_image = Image.new('RGBA', image.size)
+    rectangle_image_draw = ImageDraw.Draw(rectangle_image)
+    rectangle_image_draw.rectangle(**options)
+    if mask_callback:
+        mask_callback(rectangle_image_draw)
+    rectangle_image = rectangle_image.rotate(label_rotate_angle,  expand=1)
+    image.paste(im=rectangle_image, mask=rectangle_image) 
+    # draw.bitmap((100, 100), rectangle_image)
+    return image
+
+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,
+        label_rotate_angle: int = 0) -> 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
+      label_rotate_angle: Angle to rotate the label text
+    Returns:
+      Image: Image annotated with bounding boxes
+    '''
+    image = image.copy().convert("RGB")
+    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
+        
+        image = draw_bounding_box(
+            image = image, 
+            bbox_outline_width = bbox_outline_width, 
+            bbox_fill_color = _bbox_fill_color, 
+            bbox_outline_color = _bbox_outline_color, 
+            bbox = [x0, y0, x1, y1])
+
+        if label is not None:
+            image = draw_text_on_image(
+                image = image,
+                text_position_xy = [x0, y0],
+                label = label,
+                label_text_color = label_text_color,
+                label_fill_color = label_fill_color,
+                label_text_padding = label_text_padding,
+                label_rectangle_left_margin = label_rectangle_left_margin,
+                label_rectangle_top_margin = label_rectangle_top_margin,
+                label_text_size = label_text_size,
+                font = font,
+                label_rotate_angle = label_rotate_angle)
+    return image
+
+def draw_text_on_image(
+        image: Image.Image,
+        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,
+        label_rotate_angle: int = 0) -> Image.Image:
+    image = image.copy().convert("RGB")
+    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)
+    _, _, 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
+    ]
+    image = draw_bounding_box(
+        image = image, 
+        bbox_outline_width = 0,
+        bbox_fill_color = label_fill_color, 
+        bbox_outline_color = None,
+        bbox = xy,
+        label_rotate_angle = label_rotate_angle,
+        mask_callback = lambda mask_draw: mask_draw.text(text_position, label, font=font, fill=label_text_color))
+    return image