annotate.py

"""Draws forensic evidence markers on the original screenshot."""

from __future__ import annotations

from PIL import Image, ImageDraw, ImageFont

from detective import CaseFile, Evidence

SEVERITY_COLORS = {
    "capital": (192, 57, 43),    # blood red
    "high":    (231, 76, 60),    # bright red
    "medium":  (230, 126, 34),   # orange
    "low":     (241, 196, 15),   # evidence yellow
}

DEFAULT_COLOR = (231, 76, 60)


def _load_font(size: int) -> ImageFont.ImageFont:
    # Try a few common fonts; fall back to PIL default.
    for name in ("arialbd.ttf", "Arial Bold.ttf", "DejaVuSans-Bold.ttf", "arial.ttf"):
        try:
            return ImageFont.truetype(name, size=size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


def _draw_marker(
    draw: ImageDraw.ImageDraw,
    *,
    cx: int,
    cy: int,
    number: int,
    color: tuple[int, int, int],
    radius: int,
    font: ImageFont.ImageFont,
) -> None:
    # Outer white halo for readability on dark UIs.
    draw.ellipse(
        [cx - radius - 3, cy - radius - 3, cx + radius + 3, cy + radius + 3],
        fill=(255, 255, 255),
    )
    draw.ellipse(
        [cx - radius, cy - radius, cx + radius, cy + radius],
        fill=color,
        outline=(20, 20, 20),
        width=2,
    )
    text = str(number)
    bbox = draw.textbbox((0, 0), text, font=font)
    tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
    draw.text(
        (cx - tw / 2 - bbox[0], cy - th / 2 - bbox[1]),
        text,
        fill=(255, 255, 255),
        font=font,
    )


def annotate(image: Image.Image, case: CaseFile) -> Image.Image:
    """Return a new image with bold evidence rectangles + numbered markers.
    Markers are nudged apart so overlapping boxes don't hide each other."""
    img = image.convert("RGB").copy()
    W, H = img.size

    # Clean, legible weight (the heavy double-casing looked muddy).
    stroke = max(3, min(W, H) // 240)
    marker_radius = max(15, min(W, H) // 58)
    font_size = max(17, marker_radius + 3)
    font = _load_font(font_size)

    draw = ImageDraw.Draw(img, "RGBA")
    placed: list[tuple[int, int]] = []

    def _nudge(cx: int, cy: int) -> tuple[int, int]:
        """Push a marker off any already-placed one (keeps every number visible)."""
        gap = int(marker_radius * 2.15)
        for _ in range(12):
            clash = any((cx - px) ** 2 + (cy - py) ** 2 < gap ** 2 for px, py in placed)
            if not clash:
                break
            cx += gap
            if cx > W - marker_radius:
                cx = marker_radius + 2
                cy += gap
        cx = max(marker_radius + 2, min(W - marker_radius - 2, cx))
        cy = max(marker_radius + 2, min(H - marker_radius - 2, cy))
        placed.append((cx, cy))
        return cx, cy

    for ev in case.evidence:
        x1, y1, x2, y2 = ev.bbox
        color = SEVERITY_COLORS.get(ev.severity, DEFAULT_COLOR)

        # A grease-pencil style ellipse circling the evidence (matches the
        # animated draw-on circles in the live app). Pad it out a touch so it
        # encircles the element rather than clipping it.
        px = max(4, (x2 - x1) // 22)
        py = max(4, (y2 - y1) // 22)
        ex1, ey1, ex2, ey2 = x1 - px, y1 - py, x2 + px, y2 + py
        draw.ellipse([ex1 - 1, ey1 - 1, ex2 + 1, ey2 + 1], outline=(10, 10, 10, 180), width=1)
        draw.ellipse([ex1, ey1, ex2, ey2], outline=color + (255,), width=stroke)

        cx, cy = _nudge(max(marker_radius + 2, x1), max(marker_radius + 2, y1))
        _draw_marker(draw, cx=cx, cy=cy, number=ev.id, color=color,
                     radius=marker_radius, font=font)

    return img