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 containment(inner: Box, outer: Box) -> float:
    """inner 有多大比例落在 outer 里（0~1）。面积悬殊时 IoU 会很小，但这个值接近 1。"""
    left, top = max(inner.x1, outer.x1), max(inner.y1, outer.y1)
    right, bottom = min(inner.x2, outer.x2), min(inner.y2, outer.y2)
    if right <= left or bottom <= top:
        return 0.0
    area = max(inner.x2 - inner.x1, 0) * max(inner.y2 - inner.y1, 0)
    return (right - left) * (bottom - top) / area if area > 0 else 0.0


def nms(
    boxes: Sequence[Box],
    iou_threshold: float = 0.5,
    per_class: bool = True,
    contained_threshold: float = 0.0,
) -> list[int]:
    """非极大值抑制，返回保留下来的下标（按输入顺序）。

    用于把切片推理的结果拼回整图：相邻片在重叠区会把同一处缺陷各检出一次，
    合并后必须去重。注意这与"禁止框重叠"是两回事——这里只压掉重复，
    真正分开的两个缺陷即使轻微相邻也都保留。

    Args:
        per_class: 只在同类之间抑制。不同类别的框重叠往往是有意义的分歧，不该互相压制。
        contained_threshold: 额外按**包含比例**抑制，0 = 关闭。切片时缺陷被切割线切开，
            相邻片只看到一角，会给出一个小残框；它与完整框的 IoU 很小（面积悬殊），
            光靠 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]
        duplicate = False
        for other in kept:
            if per_class and boxes[other].class_id != box.class_id:
                continue
            if iou(box, boxes[other]) > iou_threshold:
                duplicate = True
                break
            if contained_threshold and containment(box, boxes[other]) >= contained_threshold:
                duplicate = True
                break
        if not duplicate:
            kept.append(index)
    return sorted(kept)


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