ModelTrainEval/engine/tiling.py

"""把任意尺寸的图切成固定边长的方片，只用裁剪与镜像填充，不做任何缩放。

产线要求送进模型的图固定为 N×N（当前 160），且**禁止 resize 这类拉伸**——缺陷的绝对尺度
是判别依据之一，拉伸会把"一个 5 像素的黑点"变成别的东西。所以每个维度只有两种处理：

- 边长 < N：**镜像填充**到 N（`reflect`，真镜像，不重复边缘像素），原图居中；
- 边长 ≥ N：按步长**裁剪**成若干片，最后一片贴右/下边缘，保证整图被覆盖。

两种操作都不改变像素的尺度与形状，框只是平移与裁剪。

切片后坐标系变了，框要跟着重映射：与片相交的部分保留，露出太少的丢弃——半个缺陷
既不像缺陷也不像背景，留着只会教坏模型。

推理时用同一套 `plan_tiles` 切图，把每片的框按 `tile.x/y/pad` 反向平移回原图即可拼回。
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import NamedTuple

import numpy as np

from engine.detection import Box


class Tile(NamedTuple):
    """一个片的取法。

    先从原图取 `source[y : y + crop_height, x : x + crop_width]`，
    再镜像填充成 `size × size`，原图内容落在 `(pad_left, pad_top)` 处。

    维度足够大时 `pad = 0`（纯裁剪）；维度不足时 `crop` 就是整条边（纯填充）。
    """

    x: int
    y: int
    crop_width: int
    crop_height: int
    pad_left: int
    pad_top: int


def axis_starts(length: int, size: int, overlap: int) -> list[int]:
    """一个维度上的裁剪起点。

    步长 = size - overlap；最后一片强制贴边（起点 = length - size），
    这样末尾那条不足一步的区域不会被漏掉，代价是它与前一片重叠更多。
    """
    if length <= size:
        return [0]
    stride = max(1, size - overlap)
    starts = list(range(0, length - size + 1, stride))
    if starts[-1] != length - size:
        starts.append(length - size)
    return starts


def plan_tiles(width: int, height: int, size: int, overlap: int = 0) -> list[Tile]:
    """把 `width × height` 切成若干 `size × size` 的片。"""
    if size <= 0:
        raise ValueError(f"size 必须为正，得到 {size}")
    if overlap < 0 or overlap >= size:
        raise ValueError(f"overlap 需在 [0, size) 内，得到 {overlap}")

    crop_width, crop_height = min(width, size), min(height, size)
    # 不足一片时居中放置：缺陷贴边时两侧都有镜像出来的上下文，比靠角落更自然
    pad_left = (size - crop_width) // 2
    pad_top = (size - crop_height) // 2
    return [
        Tile(x, y, crop_width, crop_height, pad_left, pad_top)
        for y in axis_starts(height, size, overlap)
        for x in axis_starts(width, size, overlap)
    ]


def cut(image: np.ndarray, tile: Tile, size: int) -> np.ndarray:
    """按 `tile` 取出一片并镜像填充到 `size × size`。

    `image` 是 HxW 或 HxWxC。填充用 `reflect`：填充量超过边长时 numpy 会反复折返，
    小图（最小的只有 83 像素高）也能填到 160。
    """
    patch = image[tile.y:tile.y + tile.crop_height, tile.x:tile.x + tile.crop_width]
    pad_bottom = size - tile.crop_height - tile.pad_top
    pad_right = size - tile.crop_width - tile.pad_left
    if not (tile.pad_top or tile.pad_left or pad_bottom or pad_right):
        return patch
    widths = [(tile.pad_top, pad_bottom), (tile.pad_left, pad_right)]
    widths += [(0, 0)] * (image.ndim - 2)
    return np.pad(patch, widths, mode="reflect")


def remap_boxes(
    boxes: Sequence[Box],
    tile: Tile,
    width: int,
    height: int,
    size: int,
    min_visibility: float = 0.3,
) -> list[Box]:
    """把原图上的归一化框映射到某一片上，返回该片内的归一化框。

    Args:
        boxes: 相对原图归一化的 xyxy 框。
        width, height: 原图尺寸，用于还原成像素。
        min_visibility: 框落在片内的面积占比下限。低于它就丢弃——只露出一角的缺陷
            既不像缺陷也不像背景。
    """
    kept = []
    for box in boxes:
        x1, y1 = box.x1 * width, box.y1 * height
        x2, y2 = box.x2 * width, box.y2 * height
        area = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
        if area <= 0:
            continue

        # 平移到片的坐标系：先减去裁剪原点，再加上填充偏移
        tx1 = x1 - tile.x + tile.pad_left
        ty1 = y1 - tile.y + tile.pad_top
        tx2 = x2 - tile.x + tile.pad_left
        ty2 = y2 - tile.y + tile.pad_top

        # 只保留落在片内的部分。镜像出来的区域里也有缺陷的倒影，但那是复制品，
        # 不该再标一个框——标了等于告诉模型"这里真有一个缺陷"
        cx1, cy1 = max(tx1, 0.0), max(ty1, 0.0)
        cx2 = min(tx2, float(tile.pad_left + tile.crop_width))
        cy2 = min(ty2, float(tile.pad_top + tile.crop_height))
        clipped = max(cx2 - cx1, 0.0) * max(cy2 - cy1, 0.0)
        if clipped <= 0 or clipped / area < min_visibility:
            continue
        kept.append(Box(box.class_id, cx1 / size, cy1 / size, cx2 / size, cy2 / size, box.score))
    return kept


def merge_tiles(
    per_tile: Sequence[Sequence[Box]],
    tiles: Sequence[Tile],
    size: int,
    width: int,
    height: int,
    iou_threshold: float = 0.5,
    contained_threshold: float = 0.8,
) -> list[Box]:
    """把每片各自的框拼回整图，返回相对原图归一化的框。

    三步：映射回原图像素 → 裁掉落在镜像填充区里的部分（那里是复制品，不是真实区域）
    → 跨片去重。

    去重同时看 IoU 与包含比例：缺陷被切割线切开时，相邻片只看到一角，给出的小残框
    与完整框的 IoU 很小（面积悬殊），但几乎整个落在完整框里。只看 IoU 会把残框留下来。

    残框与完整框互不包含的情况（缺陷正好被切成两半、两片各见一半）只能靠 `overlap`
    避免——重叠大于缺陷尺寸时，总有一片能完整看到它。
    """
    from engine.detection import nms

    collected = []
    for tile, boxes in zip(tiles, per_tile, strict=True):
        for box in boxes:
            mapped = to_source(box, tile, size)
            x1, y1 = max(mapped.x1, 0.0), max(mapped.y1, 0.0)
            x2, y2 = min(mapped.x2, float(width)), min(mapped.y2, float(height))
            if x2 <= x1 or y2 <= y1:
                continue
            collected.append(
                Box(box.class_id, x1 / width, y1 / height, x2 / width, y2 / height, box.score)
            )
    return [
        collected[i]
        for i in nms(collected, iou_threshold, contained_threshold=contained_threshold)
    ]


def tiled_predict(
    predict_batch,
    image: np.ndarray,
    size: int,
    overlap: int = 0,
    iou_threshold: float = 0.5,
) -> list[Box]:
    """按固定边长切片推理再拼回，返回整图上的归一化框。

    Args:
        predict_batch: 接收一批 `size × size` 的图、返回每张图的归一化框列表。
            推理框架留在调用方，这里只管几何，便于单测。
        image: HxW 或 HxWxC 的原图。
    """
    height, width = image.shape[:2]
    tiles = plan_tiles(width, height, size, overlap)
    patches = [cut(image, tile, size) for tile in tiles]
    return merge_tiles(predict_batch(patches), tiles, size, width, height, iou_threshold)


def to_source(box: Box, tile: Tile, size: int) -> Box:
    """反向：把片上的归一化框映射回原图的**像素** xyxy。

    推理时每片各自出框，靠这个拼回原图坐标再做一次 NMS。
    """
    return Box(
        box.class_id,
        box.x1 * size - tile.pad_left + tile.x,
        box.y1 * size - tile.pad_top + tile.y,
        box.x2 * size - tile.pad_left + tile.x,
        box.y2 * size - tile.pad_top + tile.y,
        box.score,
    )