detection_utils.py

"""Rendering and summary helpers for street-scene object detection."""

from __future__ import annotations

import csv
from collections import defaultdict
from pathlib import Path
from typing import Iterable

from PIL import Image, ImageDraw, ImageFont


STREET_OBJECT_GROUPS: dict[str, set[str]] = {
    "people": {"person"},
    "active_mobility": {"person", "bicycle"},
    "motor_vehicles": {"car", "motorcycle", "bus", "truck", "train"},
    "all_transport": {
        "bicycle",
        "car",
        "motorcycle",
        "bus",
        "truck",
        "train",
    },
}


def _load_font(size: int) -> ImageFont.ImageFont:
    candidates = (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
        "/System/Library/Fonts/Helvetica.ttc",
    )
    for candidate in candidates:
        try:
            return ImageFont.truetype(candidate, size)
        except (OSError, IOError):
            continue
    return ImageFont.load_default()


def detection_color(class_id: int) -> tuple[int, int, int]:
    """Return a stable, high-contrast color for a COCO class ID."""
    return (
        int((67 * class_id + 37) % 190 + 40),
        int((97 * class_id + 71) % 190 + 40),
        int((43 * class_id + 113) % 190 + 40),
    )


def _boxes_overlap(
    first: tuple[float, float, float, float],
    second: tuple[float, float, float, float],
) -> bool:
    return not (
        first[2] <= second[0]
        or second[2] <= first[0]
        or first[3] <= second[1]
        or second[3] <= first[1]
    )


def _label_box(
    object_box: list[float],
    label_width: float,
    label_height: float,
    image_size: tuple[int, int],
    occupied: list[tuple[float, float, float, float]],
) -> tuple[float, float, float, float]:
    """Place a label near its object while avoiding earlier labels."""
    image_width, image_height = image_size
    center_x = (object_box[0] + object_box[2]) / 2
    x = max(0.0, min(center_x - label_width / 2, image_width - label_width))

    candidates: list[tuple[float, float, float, float]] = []
    for slot in range(7):
        bottom = object_box[1] - 4 - slot * (label_height + 4)
        top = bottom - label_height
        if top >= 0:
            candidates.append((x, top, x + label_width, bottom))
    for slot in range(4):
        top = object_box[3] + 4 + slot * (label_height + 4)
        bottom = top + label_height
        if bottom <= image_height:
            candidates.append((x, top, x + label_width, bottom))

    for candidate in candidates:
        if not any(_boxes_overlap(candidate, used) for used in occupied):
            return candidate
    return candidates[0] if candidates else (x, 0.0, x + label_width, label_height)


def render_detection(
    image: Image.Image,
    detections: Iterable[dict[str, object]],
) -> Image.Image:
    """Draw labeled bounding boxes on a copy of the input image."""
    rendered = image.convert("RGB").copy()
    draw = ImageDraw.Draw(rendered)
    short_side = min(rendered.size)
    line_width = max(2, round(short_side / 320))
    font = _load_font(max(13, min(24, round(short_side / 55))))
    occupied_labels: list[tuple[float, float, float, float]] = []

    ordered_detections = sorted(
        detections,
        key=lambda item: (float(item["x1"]), float(item["y1"])),
    )
    for detection in ordered_detections:
        class_id = int(detection["class_id"])
        color = detection_color(class_id)
        box = [
            float(detection["x1"]),
            float(detection["y1"]),
            float(detection["x2"]),
            float(detection["y2"]),
        ]
        label = (
            f"{detection['class_name']} "
            f"{float(detection['confidence']):.2f}"
        )
        draw.rectangle(box, outline=color, width=line_width)
        text_box = draw.textbbox((0, 0), label, font=font)
        text_height = text_box[3] - text_box[1]
        text_width = text_box[2] - text_box[0]
        background = _label_box(
            box,
            text_width + 8,
            text_height + 8,
            rendered.size,
            occupied_labels,
        )
        occupied_labels.append(background)
        draw.rectangle(background, fill=color)
        draw.text(
            (background[0] + 4, background[1] + 4),
            label,
            fill="white",
            font=font,
        )

    return rendered


def build_detection_summary(
    detections: Iterable[dict[str, object]],
) -> list[list[object]]:
    """Aggregate detection counts and confidence by class."""
    grouped: dict[str, list[float]] = defaultdict(list)
    for detection in detections:
        grouped[str(detection["class_name"])].append(
            float(detection["confidence"])
        )

    rows = [
        [class_name, len(confidences), round(sum(confidences) / len(confidences), 3), round(max(confidences), 3)]
        for class_name, confidences in grouped.items()
    ]
    rows.sort(key=lambda row: (-int(row[1]), str(row[0])))
    return rows


def build_detection_table(
    detections: Iterable[dict[str, object]],
) -> list[list[object]]:
    """Build one exportable row per bounding box."""
    rows: list[list[object]] = []
    for index, detection in enumerate(detections, start=1):
        rows.append(
            [
                index,
                str(detection["class_name"]),
                round(float(detection["confidence"]), 3),
                round(float(detection["x1"]), 1),
                round(float(detection["y1"]), 1),
                round(float(detection["x2"]), 1),
                round(float(detection["y2"]), 1),
            ]
        )
    return rows


def build_street_indicators(
    detections: Iterable[dict[str, object]],
) -> dict[str, int]:
    """Derive transparent street-scene counts from visible COCO objects."""
    class_names = [str(detection["class_name"]) for detection in detections]
    return {
        name: sum(class_name in members for class_name in class_names)
        for name, members in STREET_OBJECT_GROUPS.items()
    }


def write_detection_csv(
    path: Path,
    rows: Iterable[Iterable[object]],
) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["object_id", "class", "confidence", "x1", "y1", "x2", "y2"])
        writer.writerows(rows)