diff --git a/engine/detection.py b/engine/detection.py index 3fd01f3..d172659 100644 --- a/engine/detection.py +++ b/engine/detection.py @@ -102,6 +102,53 @@ def match_boxes( ) +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 图像并返回）。 diff --git a/engine/tiled_inference.py b/engine/tiled_inference.py new file mode 100644 index 0000000..187d7c8 --- /dev/null +++ b/engine/tiled_inference.py @@ -0,0 +1,85 @@ +"""用 ultralytics 做切片推理的公共封装。 + +三个调用方（可视化、批量推送预测、ML backend）都要"切片 -> 逐片推理 -> 拼回原图"， +逻辑一致，放在一处免得各写一份走样。 + +几何在 `engine.tiling` 里，与推理框架无关且单独可测；这里只负责把 ultralytics 接上去。 +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import numpy as np + +from engine.detection import Box +from engine.tiling import tiled_predict + + +def predict_image( + model, + image: np.ndarray, + size: int, + overlap: int = 0, + conf: float = 0.25, + device: str | int = "cpu", + batch: int = 16, + iou_threshold: float = 0.5, + whole_image_imgsz: int | None = None, +) -> list[Box]: + """对一张任意尺寸的图做切片推理，返回整图上的归一化框。 + + `size <= 0` 时退化为整图推理，交给 ultralytics 按 `whole_image_imgsz` 做 letterbox + ——那会缩放，只用于和旧模型对比。 + """ + if size <= 0: + return _predict_batch( + model, [image], whole_image_imgsz, conf=conf, device=device, batch=batch + )[0] + + def run(patches: Sequence[np.ndarray]) -> list[list[Box]]: + return _predict_batch(model, patches, size, conf, device, batch) + + return tiled_predict(run, image, size, overlap, iou_threshold) + + +def _predict_batch( + model, + images: Sequence[np.ndarray], + size: int | None, + conf: float, + device: str | int, + batch: int, +) -> list[list[Box]]: + """分批送进 ultralytics，返回每张图的归一化框。 + + 必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch， + 片数一多就会撑爆显存。 + """ + results = [] + for start in range(0, len(images), batch): + chunk = list(images[start:start + batch]) + kwargs = {"conf": conf, "device": device, "verbose": False} + if size: + kwargs["imgsz"] = size + results.extend(model.predict(source=chunk, **kwargs)) + + out = [] + for result in results: + out.append([ + Box(int(class_id), *(float(v) for v in xyxyn), float(score)) + for xyxyn, class_id, score in zip( + result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(), + result.boxes.conf.tolist(), strict=True, + ) + ]) + return out + + +def load_image(path) -> np.ndarray: + """读图为 numpy 数组，保持原始通道数（灰度图不强行转 RGB，免得多占三倍内存）。""" + from PIL import Image + + with Image.open(path) as handle: + image = handle.convert("L") if handle.mode in ("L", "1", "I;16") else handle.convert("RGB") + return np.asarray(image) diff --git a/engine/tiling.py b/engine/tiling.py index d6d5ef5..c957017 100644 --- a/engine/tiling.py +++ b/engine/tiling.py @@ -132,6 +132,65 @@ def remap_boxes( 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。 diff --git a/ml_backend/README.md b/ml_backend/README.md index c8cd7ce..6cb7965 100644 --- a/ml_backend/README.md +++ b/ml_backend/README.md @@ -65,6 +65,8 @@ Label Studio 里：Settings → **Model** → Connect Model，URL 填 `http://lo | `YOLO_IMGSZ` | `640` | 推理尺寸，与训练保持一致 | | `YOLO_DEVICE` | `cuda` | `cuda` / `cpu` / `mps` | | `YOLO_BATCH` | `16` | 一次推理几张。Retrieve Predictions 会一次发来上千个任务，不分批会撑爆显存 | +| `YOLO_TILE_SIZE` | `160` | 模型输入边长。按它切片推理再拼回，全程不缩放；0 = 整图推理（会 letterbox） | +| `YOLO_TILE_OVERLAP` | `32` | 切片重叠像素。大于缺陷尺寸时，总有一片能完整看到它 | | `RETRAIN_AFTER` | `0` | 累计多少条新标注触发重训，0 = 不触发 | | `RETRAIN_COMMAND` | 空 | 触发时执行的命令 | diff --git a/ml_backend/model.py b/ml_backend/model.py index c3912dd..4a3bac3 100644 --- a/ml_backend/model.py +++ b/ml_backend/model.py @@ -8,6 +8,9 @@ YOLO_DEVICE cuda / cpu / mps，默认 cuda（不可用时 ultralytics 自己退回） YOLO_BATCH 一次推理几张，默认 16。Retrieve Predictions 会一次发来成百上千个 任务，不分批会直接把显存撑爆 + YOLO_TILE_SIZE 模型输入边长，默认 160。按它把图切片推理再拼回，全程不缩放； + 0 = 整图推理，交给 ultralytics letterbox（会缩放） + YOLO_TILE_OVERLAP 切片重叠像素，默认 32 RETRAIN_AFTER 累计多少条新标注后自动触发重训，默认 0=只靠手动触发 RETRAIN_COMMAND 重训命令；不设则只记数不执行 RETRAIN_LOG 重训命令的输出写到哪，默认 /tmp/retrain.log @@ -19,6 +22,7 @@ import logging import os import shlex import subprocess +import sys import threading from pathlib import Path @@ -26,6 +30,10 @@ from convert import to_prediction from label_studio_ml.model import LabelStudioMLBase from label_studio_ml.response import ModelResponse +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from engine.tiled_inference import load_image, predict_image # noqa: E402 + logger = logging.getLogger(__name__) @@ -59,8 +67,6 @@ class FilmDefectDetector(LabelStudioMLBase): def predict(self, tasks: list[dict], context: dict | None = None, **kwargs) -> ModelResponse: """对一批任务推理，返回 Label Studio 格式的预测。""" - from PIL import Image - # 控件名从项目的 labeling config 里读，不写死：改了配置这里要跟着变 from_name, to_name, value_key = self.get_first_tag_occurence( "RectangleLabels", "Image" @@ -68,42 +74,29 @@ class FilmDefectDetector(LabelStudioMLBase): model = self._load() classes = model.names conf = float(os.environ.get("YOLO_CONF", 0.25)) - imgsz = int(os.environ.get("YOLO_IMGSZ", 640)) device = os.environ.get("YOLO_DEVICE", "cuda") + batch = int(os.environ.get("YOLO_BATCH", 16)) + tile_size = int(os.environ.get("YOLO_TILE_SIZE", 160)) + overlap = int(os.environ.get("YOLO_TILE_OVERLAP", 32)) + imgsz = int(os.environ.get("YOLO_IMGSZ", 640)) - paths, sizes = [], [] + model_version = self.get("model_version") + predictions = [] for task in tasks: # 本地存储的图要靠 LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT 解析； # 上传或云存储的图则会被下载到本地缓存 path = self.get_local_path(task["data"][value_key], task_id=task.get("id")) - with Image.open(path) as image: - sizes.append(image.size) - paths.append(path) - - model_version = self.get("model_version") - if not paths: - return ModelResponse(model_version=model_version, predictions=[]) - - # 必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch。 - # Label Studio 的 Retrieve Predictions 一次能发来上千个任务，不分批必然 OOM - batch = int(os.environ.get("YOLO_BATCH", 16)) - results = [] - for start in range(0, len(paths), batch): - results.extend(model.predict( - source=paths[start:start + batch], imgsz=imgsz, conf=conf, - device=device, verbose=False, - )) - - predictions = [] - for result, (width, height) in zip(results, sizes, strict=True): + image = load_image(path) + height, width = image.shape[:2] + # 切片推理：模型的输入尺寸是固定的，整图直接送进去会被 letterbox 缩放。 + # predict_image 内部也会分批，Retrieve Predictions 一次发来上千个任务也不会 OOM + detections = predict_image( + model, image, size=tile_size, overlap=overlap, conf=conf, + device=device, batch=batch, whole_image_imgsz=imgsz, + ) boxes = [ - (classes[int(class_id)], xyxyn, float(score)) - for xyxyn, class_id, score in zip( - result.boxes.xyxyn.tolist(), - result.boxes.cls.tolist(), - result.boxes.conf.tolist(), - strict=True, - ) + (classes[box.class_id], (box.x1, box.y1, box.x2, box.y2), box.score) + for box in detections ] predictions.append( to_prediction(boxes, width, height, from_name, to_name, model_version) diff --git a/scripts/push_predictions.py b/scripts/push_predictions.py index b3d92b4..f1b4693 100644 --- a/scripts/push_predictions.py +++ b/scripts/push_predictions.py @@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "ml_backend")) from convert import to_prediction # noqa: E402 +from engine.tiled_inference import load_image, predict_image # noqa: E402 from utils import setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 @@ -186,7 +187,12 @@ def main() -> int: ) parser.add_argument("--model-version", default=None, help="缺省按权重文件名与修改时间生成") parser.add_argument("--conf", type=float, default=0.25) - parser.add_argument("--imgsz", type=int, default=640) + parser.add_argument( + "--tile-size", type=int, default=160, + help="模型输入边长；按它切片推理再拼回，全程不缩放。0=整图推理（会 letterbox）", + ) + parser.add_argument("--tile-overlap", type=int, default=32, help="切片重叠像素") + parser.add_argument("--imgsz", type=int, default=640, help="仅 --tile-size 0 时用") parser.add_argument("--device", default="auto") parser.add_argument("--batch", type=int, default=16, help="每次推理几张") parser.add_argument("--chunk", type=int, default=500, help="每次上传几条预测") @@ -256,12 +262,15 @@ def main() -> int: removed = client.delete_predictions(args.project, model_version) logger.info("删除同版本（%s）旧预测 %d 条", model_version, removed) - from PIL import Image from ultralytics import YOLO model = YOLO(str(weights)) classes = model.names device_arg = device.index or 0 if device.type == "cuda" else device.type + if args.tile_size > 0: + logger.info("切片推理 %dx%d，重叠 %d（不缩放）", args.tile_size, args.tile_size, args.tile_overlap) + else: + logger.warning("整图推理 imgsz=%d —— ultralytics 会 letterbox 缩放", args.imgsz) pending: list[dict] = [] uploaded = with_boxes = total_boxes = 0 @@ -275,35 +284,30 @@ def main() -> int: uploaded += len(pending) pending = [] - for start in range(0, len(pairs), args.batch): - chunk = pairs[start:start + args.batch] - results = model.predict( - source=[str(path) for _, path in chunk], imgsz=args.imgsz, conf=args.conf, - device=device_arg, verbose=False, + for processed, (task_id, path) in enumerate(pairs, start=1): + image = load_image(path) + height, width = image.shape[:2] + detections = predict_image( + model, image, size=args.tile_size, overlap=args.tile_overlap, conf=args.conf, + device=device_arg, batch=args.batch, whole_image_imgsz=args.imgsz, ) - for (task_id, path), result in zip(chunk, results, strict=True): - with Image.open(path) as image: - width, height = image.size - boxes = [ - (classes[int(class_id)], xyxyn, float(score)) - for xyxyn, class_id, score in zip( - result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(), - result.boxes.conf.tolist(), strict=True, - ) - ] - prediction = to_prediction(boxes, width, height, "label", "image", model_version) - prediction["task"] = task_id - pending.append(prediction) - total_boxes += len(boxes) - with_boxes += bool(boxes) + boxes = [ + (classes[box.class_id], (box.x1, box.y1, box.x2, box.y2), box.score) + for box in detections + ] + prediction = to_prediction(boxes, width, height, "label", "image", model_version) + prediction["task"] = task_id + pending.append(prediction) + total_boxes += len(boxes) + with_boxes += bool(boxes) + if len(pending) >= args.chunk: try: flush() except RuntimeError as exc: logger.error("%s", exc) return 1 - processed = start + len(chunk) - if processed % 1000 < args.batch: + if processed % 1000 == 0: logger.info("处理 [%d/%d]，已上传 %d", processed, len(pairs), uploaded) try: flush() diff --git a/scripts/visualize_detections.py b/scripts/visualize_detections.py index 4af07d9..d4f1b22 100644 --- a/scripts/visualize_detections.py +++ b/scripts/visualize_detections.py @@ -25,6 +25,9 @@ train / val 不分目录，但 `index.csv` 里有 `split` 列——训练集的 同样地，"命中 / 漏检 / 误检"记在 `index.csv` 的 `outcome` 列里，按它排序筛选即可。 `--only errors` 只输出有漏检或误检的图，扫起来快得多。 + +模型的输入固定为 `--tile-size`（默认 160，与训练一致）时，图会按同样规则切片推理再拼回， +全程不缩放。`--tile-size 0` 退回整图推理（ultralytics 自己 letterbox），用于和旧模型对比。 """ from __future__ import annotations @@ -43,6 +46,7 @@ import yaml # noqa: E402 from datasets.yolo import IMAGE_SUFFIXES # noqa: E402 from engine.detection import Box, draw_boxes, match_boxes, parse_label # noqa: E402 +from engine.tiled_inference import load_image, predict_image # noqa: E402 from utils import setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 @@ -124,6 +128,11 @@ def main() -> int: parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值") parser.add_argument("--iou-threshold", type=float, default=0.5, help="判定命中的 IoU") parser.add_argument("--batch", type=int, default=16) + parser.add_argument( + "--tile-size", type=int, default=160, + help="模型输入边长；按它切片推理再拼回，全程不缩放。0=整图推理（会 letterbox）", + ) + parser.add_argument("--tile-overlap", type=int, default=32, help="切片重叠像素") parser.add_argument( "--only", choices=("all", "errors"), default="all", help="errors=只输出有漏检或误检的图，统计仍按全量计", @@ -150,7 +159,16 @@ def main() -> int: "数据集: %s（%d 类，%s）", args.data, len(classes), "，".join(f"{name} {len(images)} 张" for name, images in splits.items()), ) - logger.info("conf=%.2f IoU=%.2f imgsz=%d", args.conf, args.iou_threshold, args.imgsz) + if args.tile_size > 0: + logger.info( + "conf=%.2f 命中 IoU=%.2f 切片 %dx%d 重叠 %d（不缩放）", + args.conf, args.iou_threshold, args.tile_size, args.tile_size, args.tile_overlap, + ) + else: + logger.warning( + "conf=%.2f 命中 IoU=%.2f 整图推理 imgsz=%d —— ultralytics 会 letterbox 缩放", + args.conf, args.iou_threshold, args.imgsz, + ) if out_dir.exists() and any(out_dir.iterdir()): if not args.overwrite: @@ -172,26 +190,13 @@ def main() -> int: rows = [] processed = 0 - def predict_in_batches(images): - """自己按批切分：ultralytics 的 predict 不理会 batch 参数， - 一次性把整个列表当成一个 batch，上万张图必然 OOM。""" - for start in range(0, len(images), args.batch): - chunk = images[start:start + args.batch] - results = model.predict( - source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf, - device=device_arg, verbose=False, - ) - yield from zip(chunk, results, strict=True) - for split, images in splits.items(): - for image_path, result in predict_in_batches(images): - predictions = [ - Box(int(cls), *(float(v) for v in xyxy), float(score)) - for xyxy, cls, score in zip( - result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(), - result.boxes.conf.tolist(), strict=True, - ) - ] + for image_path in images: + predictions = predict_image( + model, load_image(image_path), size=args.tile_size, overlap=args.tile_overlap, + conf=args.conf, device=device_arg, batch=args.batch, + whole_image_imgsz=args.imgsz, + ) label_file = label_path_for(image_path) truth = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else [] match = match_boxes(predictions, truth, args.iou_threshold) diff --git a/tests/test_detection.py b/tests/test_detection.py index 4e1e4f1..f8b9296 100644 --- a/tests/test_detection.py +++ b/tests/test_detection.py @@ -13,6 +13,7 @@ from engine.detection import ( # noqa: E402 from_yolo, iou, match_boxes, + nms, parse_label, ) @@ -129,3 +130,41 @@ def test_draw_handles_boxes_at_the_border(): prediction = Box(0, 0.0, 0.0, 0.5, 0.05, 0.9) draw_boxes(image, [], [prediction], ["blackdot"], match_boxes([prediction], [])) assert COLOR_UNMATCHED in _colors(image) + + +# ---------------- NMS ---------------- + + +def test_nms_removes_duplicates_of_the_same_defect(): + high = Box(0, 0.30, 0.30, 0.50, 0.50, 0.9) + duplicate = Box(0, 0.31, 0.31, 0.51, 0.51, 0.7) + assert nms([duplicate, high]) == [1] + + +def test_nms_keeps_neighbours_below_the_threshold(): + """真正分开的两个缺陷，轻微相邻也要都留住。""" + a = Box(0, 0.0, 0.0, 0.5, 0.5, 0.9) + b = Box(0, 0.4, 0.4, 0.9, 0.9, 0.8) # IoU 约 0.02 + assert nms([a, b]) == [0, 1] + + +def test_nms_is_per_class_by_default(): + """不同类别在同一位置的分歧是有意义的，不该互相压制。""" + a = Box(0, 0.3, 0.3, 0.6, 0.6, 0.9) + b = Box(4, 0.31, 0.31, 0.61, 0.61, 0.8) + assert nms([a, b]) == [0, 1] + assert nms([a, b], per_class=False) == [0] + + +def test_nms_handles_empty_input(): + assert nms([]) == [] + + +def test_nms_can_suppress_contained_fragments(): + """切片接缝上的残框：与完整框 IoU 很小，却几乎整个落在里面。""" + full = Box(0, 0.30, 0.50, 0.40, 0.60, 0.9) + fragment = Box(0, 0.39, 0.50, 0.40, 0.60, 0.8) # 只剩一条边 + assert iou(full, fragment) < 0.15 + + assert nms([full, fragment]) == [0, 1] # 只看 IoU 压不掉 + assert nms([full, fragment], contained_threshold=0.8) == [0] # 按包含比例能压掉 diff --git a/tests/test_tiling.py b/tests/test_tiling.py index 21acdeb..2e37986 100644 --- a/tests/test_tiling.py +++ b/tests/test_tiling.py @@ -184,3 +184,106 @@ def test_score_survives_both_directions(): mapped, = remap_boxes([box], tile, 100, 100, SIZE) assert mapped.score == 0.83 assert to_source(mapped, tile, SIZE).score == 0.83 + + +# ---------------- 切片推理的拼回 ---------------- + + +def test_merge_deduplicates_across_the_seam(): + """重叠区里同一处缺陷会被两片各检出一次，合并后应只剩一个。""" + from engine.tiling import merge_tiles + + width = height = 300 + tiles = plan_tiles(width, height, SIZE, overlap=64) + # 让每片都"检出"同一个真实位置的框（像素 150..170） + per_tile = [] + for tile in tiles: + px1, py1, px2, py2 = 150.0, 150.0, 170.0, 170.0 + bx1 = (px1 - tile.x + tile.pad_left) / SIZE + by1 = (py1 - tile.y + tile.pad_top) / SIZE + bx2 = (px2 - tile.x + tile.pad_left) / SIZE + by2 = (py2 - tile.y + tile.pad_top) / SIZE + inside = 0 <= bx1 and bx2 <= 1 and 0 <= by1 and by2 <= 1 + per_tile.append([Box(0, bx1, by1, bx2, by2, 0.9)] if inside else []) + + assert sum(len(b) for b in per_tile) >= 2 # 确实有多片检出 + merged = merge_tiles(per_tile, tiles, SIZE, width, height) + assert len(merged) == 1 + assert merged[0].x1 * width == pytest.approx(150, abs=1e-6) + assert merged[0].x2 * width == pytest.approx(170, abs=1e-6) + + +def test_merge_keeps_genuinely_separate_defects(): + from engine.tiling import merge_tiles + + width = height = 320 + tiles = plan_tiles(width, height, SIZE, overlap=0) + per_tile = [[] for _ in tiles] + per_tile[0] = [Box(0, 0.1, 0.1, 0.2, 0.2, 0.9)] # 左上片 + per_tile[3] = [Box(0, 0.8, 0.8, 0.9, 0.9, 0.8)] # 右下片 + assert len(merge_tiles(per_tile, tiles, SIZE, width, height)) == 2 + + +def test_merge_clips_boxes_that_reach_into_mirrored_padding(): + """小图填充后，模型可能在镜像区出框；那里不是真实区域，要裁掉。""" + from engine.tiling import merge_tiles + + width = height = 80 + tiles = plan_tiles(width, height, SIZE) + # 一个横跨整片的框，两端都落在镜像填充里 + merged = merge_tiles([[Box(0, 0.0, 0.4, 1.0, 0.6, 0.9)]], tiles, SIZE, width, height) + + assert len(merged) == 1 + assert merged[0].x1 == pytest.approx(0.0) + assert merged[0].x2 == pytest.approx(1.0) # 裁到原图边界，不超出 + assert 0.0 <= merged[0].y1 <= merged[0].y2 <= 1.0 + + +def test_merge_drops_boxes_entirely_inside_the_mirror(): + from engine.tiling import merge_tiles + + width = height = 80 + tiles = plan_tiles(width, height, SIZE) # pad 40，原图占 40..120 + ghost = Box(0, 0.05, 0.05, 0.2, 0.2, 0.9) # 像素 8..32，全在镜像区 + assert merge_tiles([[ghost]], tiles, SIZE, width, height) == [] + + +def test_tiled_predict_wires_geometry_to_the_model(): + """端到端：切片 -> 逐片推理 -> 拼回，位置要落在真实缺陷上。""" + from engine.tiling import tiled_predict + + image = np.zeros((300, 400), dtype=np.uint8) + image[200:220, 100:130] = 255 # 一块亮斑 + + def predict_batch(patches): + """假模型：把片里最亮的连通区域当成检出。""" + out = [] + for patch in patches: + ys, xs = np.nonzero(patch == 255) + if len(xs) == 0: + out.append([]) + continue + out.append([Box(0, xs.min() / SIZE, ys.min() / SIZE, + (xs.max() + 1) / SIZE, (ys.max() + 1) / SIZE, 0.9)]) + return out + + merged = tiled_predict(predict_batch, image, SIZE, overlap=32) + assert len(merged) == 1 + assert merged[0].x1 * 400 == pytest.approx(100, abs=1) + assert merged[0].y1 * 300 == pytest.approx(200, abs=1) + assert merged[0].x2 * 400 == pytest.approx(130, abs=1) + assert merged[0].y2 * 300 == pytest.approx(220, abs=1) + + +def test_tiled_predict_on_a_small_image_uses_one_tile(): + from engine.tiling import tiled_predict + + seen = [] + + def predict_batch(patches): + seen.append(len(patches)) + assert all(p.shape[:2] == (SIZE, SIZE) for p in patches) + return [[] for _ in patches] + + tiled_predict(predict_batch, np.zeros((83, 116), dtype=np.uint8), SIZE) + assert seen == [1]