inference/visualizer.py

from __future__ import annotations

import random
from PIL import Image, ImageDraw, ImageFont


def _get_color(label: str) -> tuple[int, int, int]:
    rng = random.Random(hash(label) & 0xFFFFFFFF)
    h = rng.random()
    # HSV → RGB with S=0.7, V=0.9
    import colorsys
    r, g, b = colorsys.hsv_to_rgb(h, 0.70, 0.90)
    return int(r * 255), int(g * 255), int(b * 255)


def _load_font(size: int = 14) -> ImageFont.ImageFont:
    for path in [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    ]:
        try:
            return ImageFont.truetype(path, size)
        except OSError:
            continue
    return ImageFont.load_default()


def draw_detections(
    image: Image.Image,
    detections: list[dict],
    line_width: int = 2,
    font_size: int = 14,
) -> Image.Image:
    """Overlay bounding boxes, labels, and confidence scores on *image*.

    *detections* is the list returned by RTDetrPredictor.predict().
    Returns a new RGB image (original is not mutated).
    """
    out = image.convert("RGB").copy()
    draw = ImageDraw.Draw(out, "RGBA")
    font = _load_font(font_size)

    for det in detections:
        label: str = det["label"]
        score: float = det["score"]
        box = det["box"]
        x1, y1, x2, y2 = box["x1"], box["y1"], box["x2"], box["y2"]

        color = _get_color(label)
        fill_rgba = (*color, 40)  # translucent fill

        draw.rectangle([x1, y1, x2, y2], outline=color, fill=fill_rgba, width=line_width)

        text = f"{label} {score:.2f}"
        bbox = draw.textbbox((x1, y1), text, font=font)
        text_w = bbox[2] - bbox[0]
        text_h = bbox[3] - bbox[1]

        # Background pill behind text
        pad = 3
        draw.rectangle(
            [x1, y1 - text_h - pad * 2, x1 + text_w + pad * 2, y1],
            fill=(*color, 220),
        )
        draw.text((x1 + pad, y1 - text_h - pad), text, fill=(255, 255, 255), font=font)

    return out