ModelTrainEval/engine/detection.py

"""检测结果的匹配与绘制：把预测框和真值框放在一起看。

匹配用标准的贪心策略——预测按置信度从高到低，each 认领 IoU 最高且未被认领的真值框。
这与 mAP 的匹配规则一致，所以图上看到的"漏检/误检"和指标算的是同一回事。
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import NamedTuple

#: 真值绿、命中的预测蓝、未命中的预测红。红绿对比最刺眼，留给"错"的那一类
COLOR_GT = (0, 220, 0)
COLOR_MATCHED = (0, 160, 255)
COLOR_UNMATCHED = (255, 40, 40)


class Box(NamedTuple):
    """归一化的 xyxy 框。`score` 仅预测框有。"""

    class_id: int
    x1: float
    y1: float
    x2: float
    y2: float
    score: float | None = None


def from_yolo(class_id: int, cx: float, cy: float, width: float, height: float) -> Box:
    """YOLO 的中心点记法 -> xyxy。"""
    return Box(class_id, cx - width / 2, cy - height / 2, cx + width / 2, cy + height / 2)


def parse_label(text: str) -> list[Box]:
    """读一个 YOLO 标签文件的内容；空文件返回空列表（负样本）。"""
    boxes = []
    for line in text.splitlines():
        parts = line.split()
        if len(parts) == 5:
            boxes.append(from_yolo(int(parts[0]), *(float(v) for v in parts[1:])))
    return boxes


def iou(a: Box, b: Box) -> float:
    """两个框的交并比。不相交时为 0。"""
    left, top = max(a.x1, b.x1), max(a.y1, b.y1)
    right, bottom = min(a.x2, b.x2), min(a.y2, b.y2)
    if right <= left or bottom <= top:
        return 0.0
    overlap = (right - left) * (bottom - top)
    area_a = max(a.x2 - a.x1, 0) * max(a.y2 - a.y1, 0)
    area_b = max(b.x2 - b.x1, 0) * max(b.y2 - b.y1, 0)
    union = area_a + area_b - overlap
    return overlap / union if union > 0 else 0.0


class Match(NamedTuple):
    """一张图的匹配结果。下标分别指向 predictions 与 ground_truth。"""

    pairs: list[tuple[int, int]]        # (预测下标, 真值下标)
    unmatched_predictions: list[int]    # 误检
    unmatched_ground_truth: list[int]   # 漏检


def match_boxes(
    predictions: Sequence[Box],
    ground_truth: Sequence[Box],
    iou_threshold: float = 0.5,
    same_class: bool = True,
) -> Match:
    """贪心匹配，与 mAP 的规则一致。

    Args:
        iou_threshold: 低于它不算命中。
        same_class: 是否要求类别一致。设 False 可以单看"位置对不对"——
            定位准但分类错，和完全没找到是两种问题，值得分开看。
    """
    order = sorted(range(len(predictions)), key=lambda i: -(predictions[i].score or 0.0))
    taken: set[int] = set()
    pairs = []
    unmatched_predictions = []
    for index in order:
        prediction = predictions[index]
        best, best_iou = None, iou_threshold
        for candidate, truth in enumerate(ground_truth):
            if candidate in taken:
                continue
            if same_class and truth.class_id != prediction.class_id:
                continue
            overlap = iou(prediction, truth)
            if overlap >= best_iou:
                best, best_iou = candidate, overlap
        if best is None:
            unmatched_predictions.append(index)
        else:
            taken.add(best)
            pairs.append((index, best))
    return Match(
        sorted(pairs),
        sorted(unmatched_predictions),
        [i for i in range(len(ground_truth)) if i not in taken],
    )


def suppress_overlaps(
    boxes: Sequence[Box], max_iou: float = 0.0, contained: bool = True
) -> list[int]:
    """贪心地挑出互不重叠的框，返回保留下来的下标（按输入顺序）。

    按置信度从高到低考察，与任一已保留框重叠超过 `max_iou` 就丢弃。

    为什么需要这一步：NMS 只在**同类**内、且按固定 IoU 阈值抑制，所以同一处缺陷被两个
    类别各框一次（跨类重复），或两个框 IoU 低于 NMS 阈值时，都会留下重叠框。而人工标注的
    真值里一处缺陷只有一个框，零重叠——推理侧不加这层约束，指标和图上都会出现真值里
    不可能有的形态。

    Args:
        max_iou: 允许的最大 IoU。0 表示碰到就算重叠，完全不允许。
        contained: 是否额外丢弃被已保留框完全包住的框。整图大框套小框时 IoU 可能很小
            （面积悬殊），但那显然是同一处目标的重复表达。
    """
    order = sorted(range(len(boxes)), key=lambda i: -(boxes[i].score or 0.0))
    kept: list[int] = []
    for index in order:
        box = boxes[index]
        if any(
            iou(box, boxes[other]) > max_iou or (contained and _inside(box, boxes[other]))
            for other in kept
        ):
            continue
        kept.append(index)
    return sorted(kept)


def _inside(inner: Box, outer: Box) -> bool:
    """inner 是否完全落在 outer 里（含边界）。"""
    return (
        outer.x1 <= inner.x1 and outer.y1 <= inner.y1
        and outer.x2 >= inner.x2 and outer.y2 >= inner.y2
    )


def draw_boxes(image, ground_truth: Sequence[Box], predictions: Sequence[Box],
               classes: Sequence[str], match: Match, width: int = 2):
    """把真值与预测画在图上（就地修改传入的 PIL 图像并返回）。

    真值绿色、命中的预测蓝色、未命中的预测红色；预测标注类别与置信度。
    """
    from PIL import ImageDraw

    draw = ImageDraw.Draw(image)
    w, h = image.size
    matched = {index for index, _ in match.pairs}

    def rectangle(box: Box, color, label: str | None) -> None:
        xy = [box.x1 * w, box.y1 * h, box.x2 * w, box.y2 * h]
        draw.rectangle(xy, outline=color, width=width)
        if label:
            # 文字画在框上方，贴边时改画到框内，免得被裁掉
            y = xy[1] - 11 if xy[1] > 12 else xy[1] + 1
            draw.text((xy[0] + 1, y), label, fill=color)

    for box in ground_truth:
        rectangle(box, COLOR_GT, None)
    for index, box in enumerate(predictions):
        color = COLOR_MATCHED if index in matched else COLOR_UNMATCHED
        name = classes[box.class_id] if box.class_id < len(classes) else str(box.class_id)
        rectangle(box, color, f"{name} {box.score:.2f}" if box.score is not None else name)
    return image