"""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