Spaces:
Running
Running
File size: 4,018 Bytes
23db765 cf388f7 23db765 cf388f7 23db765 | 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 | """Visualization utilities for drawing bounding boxes on images."""
from __future__ import annotations
from typing import TYPE_CHECKING
from PIL import Image, ImageDraw, ImageFont
if TYPE_CHECKING:
from src.parsing import BBox
BOX_COLORS = [
"#FF0000", "#00FF00", "#0000FF", "#FFFF00", "#FF00FF", "#00FFFF",
"#FF8800", "#8800FF", "#00FF88", "#FF0088", "#88FF00", "#0088FF",
]
MIN_BOX_SIZE = 4
def _get_font(size: int = 14) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
"""Try to load a reasonable font, fall back to default."""
try:
return ImageFont.truetype("arial.ttf", size)
except OSError:
try:
return ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", size)
except OSError:
return ImageFont.load_default()
def _hex_to_rgba(hex_color: str, alpha: int = 80) -> tuple[int, int, int, int]:
"""Convert hex color to RGBA tuple."""
h = hex_color.lstrip("#")
r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
return (r, g, b, alpha)
def draw_boxes(
image: Image.Image,
boxes: list[BBox],
labels: list[str] | None = None,
show_confidence: bool = True,
line_width: int = 3,
font_size: int = 14,
) -> Image.Image:
"""Draw bounding boxes with labels on an image.
Args:
image: Source PIL image.
boxes: List of BBox objects in pixel coordinates.
labels: Optional per-box labels. If None, uses box.label or index.
show_confidence: Whether to show confidence score in label.
line_width: Width of bounding box outlines.
font_size: Font size for labels.
Returns:
New image with drawn overlays.
"""
if not boxes:
return image.copy()
img = image.copy().convert("RGBA")
overlay = Image.new("RGBA", img.size, (0, 0, 0, 0))
draw_overlay = ImageDraw.Draw(overlay)
draw_text = ImageDraw.Draw(img)
font = _get_font(font_size)
img_w, img_h = img.size
for i, box in enumerate(boxes):
color_hex = BOX_COLORS[i % len(BOX_COLORS)]
fill_rgba = _hex_to_rgba(color_hex, alpha=50)
outline_rgb = color_hex
bx1, by1 = max(0, box.x1), max(0, box.y1)
bx2, by2 = min(img_w, box.x2), min(img_h, box.y2)
if (bx2 - bx1) < MIN_BOX_SIZE or (by2 - by1) < MIN_BOX_SIZE:
cx, cy = (bx1 + bx2) / 2, (by1 + by2) / 2
half = MIN_BOX_SIZE
bx1, by1 = cx - half, cy - half
bx2, by2 = cx + half, cy + half
draw_overlay.rectangle([bx1, by1, bx2, by2], fill=fill_rgba, outline=outline_rgb, width=line_width)
label = labels[i] if labels and i < len(labels) else (box.label or f"#{i+1}")
if show_confidence and box.confidence > 0:
label = f"{label} ({box.confidence:.0%})"
text_bbox = draw_text.textbbox((0, 0), label, font=font)
text_w = text_bbox[2] - text_bbox[0]
text_h = text_bbox[3] - text_bbox[1]
text_y = by1 - text_h - 4 if by1 - text_h - 4 > 0 else by1 + 4
text_x = max(0, bx1)
draw_text.rectangle(
[text_x, text_y, text_x + text_w + 6, text_y + text_h + 4],
fill=color_hex,
)
draw_text.text((text_x + 3, text_y + 2), label, fill="white", font=font)
img = Image.alpha_composite(img, overlay).convert("RGB")
return img
def create_no_detection_overlay(image: Image.Image, message: str = "No detections found") -> Image.Image:
"""Create an overlay indicating no objects were detected."""
img = image.copy()
draw = ImageDraw.Draw(img)
font = _get_font(18)
text_bbox = draw.textbbox((0, 0), message, font=font)
text_w = text_bbox[2] - text_bbox[0]
text_h = text_bbox[3] - text_bbox[1]
img_w, img_h = img.size
x = (img_w - text_w) / 2
y = img_h - text_h - 20
draw.rectangle([x - 10, y - 5, x + text_w + 10, y + text_h + 5], fill=(0, 0, 0, 180))
draw.text((x, y), message, fill="yellow", font=font)
return img
|