src/visualization.py

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