src/visualization.py

from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw, ImageFont

from .obstacle_dataset import YoloBox


def draw_yolo_with_heatmap(
    image: Image.Image,
    yolo_boxes: list[YoloBox],
    heatmap: np.ndarray,
    saliency_threshold: float = 0.7,
) -> Image.Image:
    canvas = overlay_heatmap(image, heatmap, saliency_threshold=saliency_threshold)
    draw = ImageDraw.Draw(canvas)
    font = ImageFont.load_default()

    for box in yolo_boxes:
        xyxy = box.to_xyxy(*canvas.size)
        draw.rectangle(xyxy, outline="lime", width=4)
        draw_label(draw, xyxy, f"YOLO: {box.class_name}", "lime", font)

    return canvas


def overlay_heatmap(
    image: Image.Image,
    heatmap: np.ndarray,
    saliency_threshold: float = 0.7,
) -> Image.Image:
    base = image.convert("RGB")
    heat = np.clip(heatmap, 0, 1)
    heat_img = Image.fromarray((heat * 255).astype(np.uint8), mode="L").resize(base.size)
    heat_values = np.asarray(heat_img).astype(np.float32) / 255.0
    alpha_values = np.where(
        heat_values >= saliency_threshold,
        ((heat_values - saliency_threshold) / max(1e-6, 1 - saliency_threshold)) * 210,
        0,
    )
    alpha = Image.fromarray(alpha_values.astype(np.uint8), mode="L")

    color = Image.new("RGBA", base.size, (255, 20, 0, 0))
    color.putalpha(alpha)
    return Image.alpha_composite(base.convert("RGBA"), color).convert("RGB")


def draw_detection_overlay(
    image: Image.Image,
    yolo_boxes: list[YoloBox],
    ijepa_box: tuple[int, int, int, int] | None = None,
    ijepa_candidate_boxes: list[tuple[int, int, int, int]] | None = None,
) -> Image.Image:
    canvas = image.convert("RGB").copy()
    draw = ImageDraw.Draw(canvas)
    font = ImageFont.load_default()

    for box in yolo_boxes:
        xyxy = box.to_xyxy(*canvas.size)
        draw.rectangle(xyxy, outline="lime", width=4)
        draw_label(draw, xyxy, f"YOLO: {box.class_name}", "lime", font)

    for index, candidate_box in enumerate(ijepa_candidate_boxes or [], start=1):
        draw.rectangle(candidate_box, outline="red", width=3)
        draw_label(draw, candidate_box, f"I-JEPA {index}", "red", font)

    if ijepa_box is not None:
        draw.rectangle(ijepa_box, outline="red", width=4)
        draw_label(draw, ijepa_box, "I-JEPA estimate", "red", font)

    return canvas


def draw_label(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], text: str, color: str, font) -> None:
    x1, y1, _, _ = box
    text_box = draw.textbbox((x1, y1), text, font=font)
    text_h = text_box[3] - text_box[1]
    label_box = (x1, max(0, y1 - text_h - 6), x1 + text_box[2] - text_box[0] + 8, max(text_h + 6, y1))
    draw.rectangle(label_box, fill=color)
    draw.text((label_box[0] + 4, label_box[1] + 3), text, fill="black", font=font)