From 2a9227531722acacfff31ff78945cea8e5a414ff Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 23 Sep 2026 12:27:57 +0800 Subject: [PATCH] =?UTF-8?q?feat(detect):=20=E5=9B=BA=E5=AE=9A=20160x160=20?= =?UTF-8?q?=E8=BE=93=E5=85=A5=EF=BC=8C=E5=8F=AA=E7=94=A8=E8=A3=81=E5=89=AA?= =?UTF-8?q?=E4=B8=8E=E9=95=9C=E5=83=8F=E5=A1=AB=E5=85=85=EF=BC=8C=E7=A6=81?= =?UTF-8?q?=E6=AD=A2=E7=BC=A9=E6=94=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 产线要求送进模型的图固定为 N×N 且不许 resize——缺陷的绝对尺度是判别依据， 拉伸会把"5 像素的黑点"变成别的东西。 - engine/tiling.py：切片几何。边长 < N 镜像填充（reflect，真镜像，原图居中）； 边长 >= N 按步长裁剪、末片贴边保证全覆盖。框跟着平移裁剪，露出面积低于 min_visibility 的丢弃；to_source 反向映射供推理时拼回原图。 - scripts/tile_dataset.py：整套数据集转换。有框图切出的空片会爆量 （524x959 能切 24 片），用 --max-empty-ratio 限量；原本无框的真负样本不受限。 - 训练配置 imgsz 改 160，并显式关掉 scale 与 translate——ultralytics 的默认值是 0.5 / 0.1 而不是 0，不写出来等于一直在做仿射缩放，违反约束。 尺寸正好等于 imgsz 时 letterbox 退化为空操作（r=1、dw=dh=0），核对过源码。 合成数据上验证：所有片正好 160x160、框区域与原图逐像素一致、填充是真镜像。 Co-Authored-By: Claude Opus 5 --- configs/detect/film.yaml | 18 ++- engine/tiling.py | 147 ++++++++++++++++++++++++ scripts/tile_dataset.py | 221 +++++++++++++++++++++++++++++++++++++ tests/test_tile_dataset.py | 183 ++++++++++++++++++++++++++++++ tests/test_tiling.py | 186 +++++++++++++++++++++++++++++++ 5 files changed, 753 insertions(+), 2 deletions(-) create mode 100644 engine/tiling.py create mode 100644 scripts/tile_dataset.py create mode 100644 tests/test_tile_dataset.py create mode 100644 tests/test_tiling.py diff --git a/configs/detect/film.yaml b/configs/detect/film.yaml index 1d87720..5923423 100644 --- a/configs/detect/film.yaml +++ b/configs/detect/film.yaml @@ -17,7 +17,10 @@ seed: 42 device: auto data: - yaml: data/film_yolo/data.yaml + # 切成固定 160x160 的片：产线要求模型输入固定尺寸且禁止 resize。 + # uv run python scripts/tile_dataset.py --data data/film_yolo/data.yaml \ + # --out data/film_yolo_160 --size 160 + yaml: data/film_yolo_160/data.yaml model: # CLAUDE.md 的"只用 convolution 模型"只约束聚类与数据集生成阶段，检测这一步不受限， @@ -27,7 +30,11 @@ model: train: epochs: 500 - imgsz: 640 + + # 必须与数据集的片大小一致。数据已经是 160x160，ultralytics 的 letterbox 这时 + # 缩放比 r = min(160/160, 160/160) = 1、填充 dw=dh=0，退化成空操作，不会碰像素。 + # 改尺寸时：重跑 tile_dataset.py，再同步改这里。 + imgsz: 160 batch: 16 workers: 8 optimizer: auto @@ -39,8 +46,15 @@ train: # 数据增强。翻转按聚类阶段的结论开着（缺陷没有固定朝向）； # 但 mosaic/mixup 关掉——它们会把四张图拼在一起，而我们每张图只有一个弱监督框， # 拼接后小目标被缩得更小，对 blackdot / whitedot 这类本来就只占千分之几面积的类别是伤害。 + # 只留镜像翻转：翻转不改变尺度，符合"禁止拉伸"的约束 fliplr: 0.5 flipud: 0.5 + + # 以下几何增强都会重采样像素，全部关掉。 + # 注意 scale 与 translate 的 ultralytics 默认值是 0.5 / 0.1（不是 0）， + # 不显式写出来就等于一直在做仿射缩放平移。 + scale: 0.0 + translate: 0.0 degrees: 0.0 mosaic: 0.0 mixup: 0.0 diff --git a/engine/tiling.py b/engine/tiling.py new file mode 100644 index 0000000..d6d5ef5 --- /dev/null +++ b/engine/tiling.py @@ -0,0 +1,147 @@ +"""把任意尺寸的图切成固定边长的方片，只用裁剪与镜像填充，不做任何缩放。 + +产线要求送进模型的图固定为 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 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, + ) diff --git a/scripts/tile_dataset.py b/scripts/tile_dataset.py new file mode 100644 index 0000000..cb900d9 --- /dev/null +++ b/scripts/tile_dataset.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python +"""把 YOLO 数据集切成固定边长的方片，只用裁剪与镜像填充，不做任何缩放。 + + uv run python scripts/tile_dataset.py \\ + --data data/film_yolo/data.yaml --out data/film_yolo_160 --size 160 + +产线要求送进模型的图固定为 N×N 且禁止 resize。这一步把训练数据转成同样的形态： + +- 边长 < N 的图**镜像填充**到 N，原图居中； +- 边长 ≥ N 的图按步长**裁剪**成多片，末片贴边，整图都被覆盖； +- 框跟着平移/裁剪，露出面积低于 `--min-visibility` 的丢弃。 + +尺寸改了（比如将来 160 → 192）重跑一次即可，训练配置里的 `imgsz` 要同步改。 + +**空片会爆量**。一张 524×959 的图能切出 24 片，其中大部分没有缺陷。全留下会让负样本 +淹没正样本，所以用 `--max-empty-ratio` 限制：空片数最多是有框片数的这个倍数，超出的 +按固定种子随机丢弃。原本就没有框的图（normal 这类真负样本）不受这个限制，它们本来 +就是要留的。 +""" + +from __future__ import annotations + +import argparse +import csv +import json +import random +import shutil +import sys +from collections import Counter +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import numpy as np # noqa: E402 +import yaml # noqa: E402 + +from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt, write_data_yaml # noqa: E402 +from engine.detection import parse_label # noqa: E402 +from engine.tiling import cut, plan_tiles, remap_boxes # noqa: E402 +from utils import setup_logging # noqa: E402 + +logger = None + + +def to_yolo_lines(boxes) -> str: + """归一化 xyxy -> YOLO 的 `class cx cy w h`。""" + lines = [] + for box in boxes: + cx, cy = (box.x1 + box.x2) / 2, (box.y1 + box.y2) / 2 + width, height = box.x2 - box.x1, box.y2 - box.y1 + lines.append(f"{box.class_id} {cx:.6f} {cy:.6f} {width:.6f} {height:.6f}") + return "".join(f"{line}\n" for line in lines) + + +def label_path_for(image: Path) -> Path: + """YOLO 约定：路径里的 images/ 换成 labels/，后缀换成 .txt。""" + parts = list(image.parts) + for index in range(len(parts) - 1, -1, -1): + if parts[index] == "images": + parts[index] = "labels" + break + return Path(*parts).with_suffix(".txt") + + +def main() -> int: + global logger + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--data", required=True, help="源数据集的 data.yaml") + parser.add_argument("--out", required=True, help="输出目录") + parser.add_argument("--size", type=int, default=160, help="片的边长，默认 160") + parser.add_argument( + "--overlap", type=int, default=0, + help="相邻片的重叠像素。大于 0 可减少缺陷被切断的概率，代价是片数变多", + ) + parser.add_argument( + "--min-visibility", type=float, default=0.3, + help="框落在片内的面积占比下限，低于它就丢弃这个框", + ) + parser.add_argument( + "--max-empty-ratio", type=float, default=1.0, + help="有框图切出的空片，最多保留到有框片数的这个倍数；<0 表示全留", + ) + parser.add_argument("--seed", type=int, default=42, help="丢弃空片时的随机种子") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--dry-run", action="store_true", help="只统计，不写文件") + args = parser.parse_args() + + logger = setup_logging() + data_yaml = Path(args.data) + meta = yaml.safe_load(data_yaml.read_text(encoding="utf-8")) + names = meta["names"] + classes = [names[i] for i in sorted(names)] if isinstance(names, dict) else list(names) + root = Path(meta.get("path", data_yaml.parent)) + if not root.is_absolute(): + root = (data_yaml.parent / root).resolve() + out_dir = Path(args.out) + + from PIL import Image + + splits = { + split: sorted( + p for p in (root / meta[split]).iterdir() if p.suffix.lower() in IMAGE_SUFFIXES + ) + for split in ("train", "val") + if meta.get(split) and (root / meta[split]).is_dir() + } + if not splits: + logger.error("%s 里没有可用的 train / val", data_yaml) + return 1 + + logger.info( + "源数据集 %s：%s；目标边长 %d，重叠 %d", + data_yaml, "，".join(f"{k} {len(v)} 张" for k, v in splits.items()), args.size, args.overlap, + ) + + if not args.dry_run: + if out_dir.exists() and any(out_dir.iterdir()): + if not args.overwrite: + logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir) + return 1 + shutil.rmtree(out_dir) + + rng = random.Random(args.seed) + rows = [] + stats: Counter[str] = Counter() + per_class: Counter[str] = Counter() + sizes_seen: Counter[str] = Counter() + + for split, images in splits.items(): + for source in images: + with Image.open(source) as handle: + image = handle.convert("L") if handle.mode == "L" else handle.convert("RGB") + width, height = image.size + array = np.asarray(image) + + label_file = label_path_for(source) + boxes = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else [] + tiles = plan_tiles(width, height, args.size, args.overlap) + sizes_seen["填充" if width < args.size or height < args.size else "裁剪"] += 1 + stats["源图"] += 1 + + positives, empties = [], [] + for index, tile in enumerate(tiles): + mapped = remap_boxes(boxes, tile, width, height, args.size, args.min_visibility) + (positives if mapped else empties).append((index, tile, mapped)) + + # 原本就没有框的图是真负样本，整张都要留；有框的图切出的空片才限量 + if boxes and args.max_empty_ratio >= 0: + allowed = int(len(positives) * args.max_empty_ratio) + if len(empties) > allowed: + stats["丢弃的空片"] += len(empties) - allowed + empties = rng.sample(empties, allowed) + + for index, tile, mapped in sorted(positives + empties): + name = f"{source.stem}_t{index:02d}" + stats["产出片"] += 1 + stats["有框片" if mapped else "空片"] += 1 + for box in mapped: + per_class[classes[box.class_id]] += 1 + rows.append({ + "split": split, "tile": f"{name}.png", "source": source.name, + "x": tile.x, "y": tile.y, + "pad_left": tile.pad_left, "pad_top": tile.pad_top, + "crop_width": tile.crop_width, "crop_height": tile.crop_height, + "boxes": len(mapped), + }) + if args.dry_run: + continue + patch = cut(array, tile, args.size) + image_path = out_dir / "images" / split / f"{name}.png" + image_path.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(patch).save(image_path) + label_out = out_dir / "labels" / split / f"{name}.txt" + label_out.parent.mkdir(parents=True, exist_ok=True) + label_out.write_text(to_yolo_lines(mapped), encoding="utf-8") + + if stats["源图"] % 2000 == 0: + logger.info("处理 [%d] 已产出 %d 片", stats["源图"], stats["产出片"]) + + logger.info( + "源图 %d（其中 %d 张靠镜像填充、%d 张需裁剪）-> 片 %d（有框 %d，空 %d，丢弃空片 %d）", + stats["源图"], sizes_seen["填充"], sizes_seen["裁剪"], stats["产出片"], + stats["有框片"], stats["空片"], stats["丢弃的空片"], + ) + logger.info("逐类框数: %s", dict(sorted(per_class.items()))) + + if args.dry_run: + logger.info("--dry-run：未写入任何文件") + return 0 + + write_data_yaml(out_dir / "data.yaml", classes, root=out_dir.resolve(), + train="images/train", val="images/val") + write_classes_txt(out_dir / "classes.txt", classes) + with (out_dir / "tiles.csv").open("w", encoding="utf-8-sig", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + (out_dir / "summary.json").write_text( + json.dumps( + { + "source": str(data_yaml.resolve()), + "size": args.size, "overlap": args.overlap, + "min_visibility": args.min_visibility, + "max_empty_ratio": args.max_empty_ratio, "seed": args.seed, + "counts": dict(stats), "per_class_boxes": dict(per_class), + "source_images_padded": sizes_seen["填充"], + "source_images_cropped": sizes_seen["裁剪"], + }, + ensure_ascii=False, indent=2, + ), + encoding="utf-8", + ) + logger.info("完成 -> %s（tiles.csv 记录每片对应的原图与偏移）", out_dir) + logger.info("训练时把 configs/detect 里的 imgsz 改成 %d，并确认 scale/translate 为 0", args.size) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_tile_dataset.py b/tests/test_tile_dataset.py new file mode 100644 index 0000000..09ba2b7 --- /dev/null +++ b/tests/test_tile_dataset.py @@ -0,0 +1,183 @@ +"""把数据集切成固定边长的方片。""" + +import csv +import json +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("PIL") + +import numpy as np # noqa: E402 +import yaml # noqa: E402 +from PIL import Image # noqa: E402 + +sys.path.insert(0, "scripts") +import tile_dataset as td # noqa: E402 + +from engine.detection import parse_label # noqa: E402 +from engine.tiling import Tile, to_source # noqa: E402 + +SIZE = 160 +CLASSES = ["blackdot", "grease"] + + +@pytest.fixture +def dataset(tmp_path): + """含小图（要填充）、大图（要切片）与真负样本的最小数据集。""" + root = tmp_path / "src" + rng = np.random.default_rng(0) + for split in ("train", "val"): + (root / "images" / split).mkdir(parents=True) + (root / "labels" / split).mkdir(parents=True) + + def put(split, name, width, height, lines): + Image.fromarray(rng.integers(0, 255, (height, width, 3), dtype=np.uint8)).save( + root / "images" / split / f"{name}.png" + ) + (root / "labels" / split / f"{name}.txt").write_text("".join(lines), encoding="utf-8") + + put("train", "blackdot_0001", 128, 128, ["0 0.5 0.5 0.1 0.1\n"]) # 小图 + put("train", "grease_0001", 400, 300, ["1 0.25 0.3 0.1 0.1\n"]) # 大图 + put("train", "normal_0001", 116, 83, []) # 真负样本 + put("val", "blackdot_0002", 320, 160, ["0 0.25 0.5 0.1 0.2\n"]) + + (root / "data.yaml").write_text( + yaml.safe_dump({"path": str(root), "train": "images/train", "val": "images/val", + "nc": 2, "names": dict(enumerate(CLASSES))}), + encoding="utf-8", + ) + return root + + +def run(dataset, out, *extra): + argv = ["tile_dataset.py", "--data", str(dataset / "data.yaml"), "--out", str(out), + "--size", str(SIZE), *extra] + original, sys.argv = sys.argv, argv + try: + return td.main() + finally: + sys.argv = original + + +def tiles_of(out): + with (out / "tiles.csv").open(encoding="utf-8-sig", newline="") as handle: + return list(csv.DictReader(handle)) + + +def test_every_tile_is_exactly_the_target_size(dataset, tmp_path): + """这是整件事的目的：模型输入固定 160x160。""" + out = tmp_path / "out" + assert run(dataset, out) == 0 + images = list((out / "images").rglob("*.png")) + assert images + for path in images: + with Image.open(path) as image: + assert image.size == (SIZE, SIZE) + + +def test_pixels_are_never_resampled(dataset, tmp_path): + """片上框住的像素必须与原图逐像素相同——有任何缩放都会露馅。""" + out = tmp_path / "out" + assert run(dataset, out) == 0 + + checked = 0 + for row in tiles_of(out): + if int(row["boxes"]) == 0: + continue + tile = Tile(int(row["x"]), int(row["y"]), int(row["crop_width"]), + int(row["crop_height"]), int(row["pad_left"]), int(row["pad_top"])) + patch = np.asarray(Image.open(out / "images" / row["split"] / row["tile"]).convert("RGB")) + source = np.asarray( + Image.open(dataset / "images" / row["split"] / row["source"]).convert("RGB") + ) + label = out / "labels" / row["split"] / f"{Path(row['tile']).stem}.txt" + for box in parse_label(label.read_text(encoding="utf-8")): + back = to_source(box, tile, SIZE) + region = patch[ + round(box.y1 * SIZE):round(box.y2 * SIZE), + round(box.x1 * SIZE):round(box.x2 * SIZE), + ] + original = source[round(back.y1):round(back.y2), round(back.x1):round(back.x2)] + assert np.array_equal(region, original) + checked += 1 + assert checked >= 3 + + +def test_small_image_is_mirror_padded_not_stretched(dataset, tmp_path): + out = tmp_path / "out" + assert run(dataset, out) == 0 + row = next(r for r in tiles_of(out) if r["source"] == "blackdot_0001.png") + tile = Tile(int(row["x"]), int(row["y"]), int(row["crop_width"]), int(row["crop_height"]), + int(row["pad_left"]), int(row["pad_top"])) + + patch = np.asarray(Image.open(out / "images" / row["split"] / row["tile"]).convert("RGB")) + source = np.asarray(Image.open(dataset / "images/train/blackdot_0001.png").convert("RGB")) + inner = patch[tile.pad_top:tile.pad_top + 128, tile.pad_left:tile.pad_left + 128] + assert np.array_equal(inner, source) # 原图原样居中 + # 填充是镜像而不是复制边缘：左侧第一个填充列等于原图第 1 列（不是第 0 列） + assert np.array_equal(patch[tile.pad_top, tile.pad_left - 1], source[0, 1]) + + +def test_large_image_is_cut_into_covering_tiles(dataset, tmp_path): + out = tmp_path / "out" + assert run(dataset, out, "--max-empty-ratio", "-1") == 0 + rows = [r for r in tiles_of(out) if r["source"] == "grease_0001.png"] + # 400x300 -> x 起点 [0,160,240]，y 起点 [0,140] + assert len(rows) == 6 + assert {int(r["x"]) for r in rows} == {0, 160, 240} + assert {int(r["y"]) for r in rows} == {0, 140} + assert all(int(r["pad_left"]) == 0 and int(r["pad_top"]) == 0 for r in rows) + + +def test_true_negative_images_keep_all_their_tiles(dataset, tmp_path): + """原本就没框的图是真负样本，不受空片配额限制。""" + out = tmp_path / "out" + assert run(dataset, out, "--max-empty-ratio", "0") == 0 + assert [r for r in tiles_of(out) if r["source"] == "normal_0001.png"] + + +def test_empty_tile_quota_limits_background_from_positive_images(dataset, tmp_path): + """有框图切出的空片会爆量，要能限量。""" + strict = tmp_path / "strict" + assert run(dataset, strict, "--max-empty-ratio", "0") == 0 + loose = tmp_path / "loose" + assert run(dataset, loose, "--max-empty-ratio", "-1") == 0 + + def empties(out): + return [r for r in tiles_of(out) if r["source"] == "grease_0001.png" and r["boxes"] == "0"] + + assert empties(strict) == [] + assert len(empties(loose)) == 5 # 6 片里 1 片有框 + + +def test_box_split_by_a_cut_line_is_dropped_when_barely_visible(dataset, tmp_path): + """min_visibility > 0.5 时，一个框最多只会出现在一片里。""" + out = tmp_path / "out" + assert run(dataset, out, "--min-visibility", "0.6") == 0 + # val 里的 320x160：框中心在 x=80，完整落在第一片内 + rows = [r for r in tiles_of(out) if r["source"] == "blackdot_0002.png"] + assert sum(int(r["boxes"]) for r in rows) == 1 + + +def test_outputs_are_a_valid_yolo_dataset(dataset, tmp_path): + out = tmp_path / "out" + assert run(dataset, out) == 0 + + meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8")) + assert [meta["names"][i] for i in range(meta["nc"])] == CLASSES + assert (out / "classes.txt").read_text(encoding="utf-8").split() == CLASSES + for row in tiles_of(out): + assert (out / "images" / row["split"] / row["tile"]).exists() + assert (out / "labels" / row["split"] / f"{Path(row['tile']).stem}.txt").exists() + + summary = json.loads((out / "summary.json").read_text(encoding="utf-8")) + assert summary["size"] == SIZE + assert summary["counts"]["源图"] == 4 + + +def test_dry_run_writes_nothing(dataset, tmp_path): + out = tmp_path / "out" + assert run(dataset, out, "--dry-run") == 0 + assert not out.exists() diff --git a/tests/test_tiling.py b/tests/test_tiling.py new file mode 100644 index 0000000..21acdeb --- /dev/null +++ b/tests/test_tiling.py @@ -0,0 +1,186 @@ +"""固定边长切片：只裁剪与镜像填充，不缩放。""" + +import pytest + +np = pytest.importorskip("numpy") + +from engine.detection import Box, from_yolo # noqa: E402 +from engine.tiling import axis_starts, cut, plan_tiles, remap_boxes, to_source # noqa: E402 + +SIZE = 160 + + +# ---------------- 切片网格 ---------------- + + +@pytest.mark.parametrize( + "length,overlap,expected", + [ + (160, 0, [0]), # 正好一片 + (128, 0, [0]), # 不足，填充 + (320, 0, [0, 160]), # 整两片 + (200, 0, [0, 40]), # 最后一片贴边，与前片重叠 120 + (480, 32, [0, 128, 256, 320]), # 步长 128，末片贴边 + ], +) +def test_axis_starts(length, overlap, expected): + assert axis_starts(length, SIZE, overlap) == expected + + +def test_last_tile_always_reaches_the_edge(): + """末尾那条不足一步的区域不能被漏掉。""" + for length in range(161, 700): + starts = axis_starts(length, SIZE, 16) + assert starts[-1] + SIZE == length + + +def test_plan_tiles_covers_a_large_image(): + tiles = plan_tiles(524, 959, SIZE, overlap=0) + assert len(tiles) == len(axis_starts(524, SIZE, 0)) * len(axis_starts(959, SIZE, 0)) + assert all(t.pad_left == 0 and t.pad_top == 0 for t in tiles) # 大图纯裁剪 + + +def test_small_image_is_centred_and_padded(): + tile, = plan_tiles(116, 83, SIZE) + assert (tile.crop_width, tile.crop_height) == (116, 83) + assert tile.pad_left == (160 - 116) // 2 == 22 + assert tile.pad_top == (160 - 83) // 2 == 38 + + +@pytest.mark.parametrize("bad", [0, -1]) +def test_rejects_invalid_size(bad): + with pytest.raises(ValueError, match="size"): + plan_tiles(100, 100, bad) + + +@pytest.mark.parametrize("bad", [-1, 160, 200]) +def test_rejects_invalid_overlap(bad): + with pytest.raises(ValueError, match="overlap"): + plan_tiles(300, 300, SIZE, bad) + + +# ---------------- 取片 ---------------- + + +def test_cut_crops_without_resizing(): + """大图取片必须是原像素的子块，一个插值都不能有。""" + image = np.arange(300 * 400, dtype=np.uint8).reshape(300, 400) % 251 + tile = plan_tiles(400, 300, SIZE, overlap=0)[1] # 第二片，x 有偏移 + patch = cut(image, tile, SIZE) + + assert patch.shape == (SIZE, SIZE) + expected = image[tile.y:tile.y + SIZE, tile.x:tile.x + SIZE] + assert np.array_equal(patch, expected) + + +def test_cut_mirrors_small_images(): + image = np.arange(1, 5, dtype=np.uint8).reshape(1, 4) # 1x4 + tile, = plan_tiles(4, 1, 4) + patch = cut(image, tile, 4) + assert patch.shape == (4, 4) + # 竖直方向镜像：原始那一行在中间，上下是它的倒影（reflect 不重复边缘行） + assert np.array_equal(patch[tile.pad_top], image[0]) + assert patch.dtype == image.dtype + + +def test_cut_keeps_channels(): + image = np.zeros((50, 60, 3), dtype=np.uint8) + tile, = plan_tiles(60, 50, SIZE) + assert cut(image, tile, SIZE).shape == (SIZE, SIZE, 3) + + +def test_cut_of_an_exact_size_image_is_the_image_itself(): + image = np.random.default_rng(0).integers(0, 255, (SIZE, SIZE), dtype=np.uint8) + tile, = plan_tiles(SIZE, SIZE, SIZE) + assert np.array_equal(cut(image, tile, SIZE), image) + + +# ---------------- 框的重映射 ---------------- + + +def test_box_shifts_with_padding(): + """小图居中填充后，框只是整体平移。""" + width, height = 80, 80 + box = from_yolo(3, 0.5, 0.5, 0.25, 0.25) # 原图中心 20x20 的框 + tile, = plan_tiles(width, height, SIZE) + mapped, = remap_boxes([box], tile, width, height, SIZE) + + assert mapped.class_id == 3 + # 原图中心落在片中心 + assert (mapped.x1 + mapped.x2) / 2 == pytest.approx(0.5) + # 尺寸不变：原来 20 像素，片上还是 20 像素（占 160 的 1/8） + assert (mapped.x2 - mapped.x1) * SIZE == pytest.approx(20) + + +def test_box_outside_the_tile_is_dropped(): + width = height = 400 + box = from_yolo(0, 0.9, 0.9, 0.05, 0.05) # 右下角 + first = plan_tiles(width, height, SIZE, overlap=0)[0] # 左上片 + assert remap_boxes([box], first, width, height, SIZE) == [] + + +def test_partially_visible_box_is_clipped_or_dropped(): + """跨片的框：露得多就裁剪保留，露得少就丢弃。""" + width = height = 320 + # 框中心正好在 x=160 的切割线上，左右各露一半 + box = from_yolo(1, 160 / width, 0.25, 40 / width, 40 / height) + left, right = plan_tiles(width, height, SIZE, overlap=0)[:2] + + kept_left, = remap_boxes([box], left, width, height, SIZE, min_visibility=0.4) + assert kept_left.x2 == pytest.approx(1.0) # 裁到片的右边缘 + assert (kept_left.x2 - kept_left.x1) * SIZE == pytest.approx(20) + + # 阈值提到 0.6 时，只露一半就不够了 + assert remap_boxes([box], left, width, height, SIZE, min_visibility=0.6) == [] + assert remap_boxes([box], right, width, height, SIZE, min_visibility=0.6) == [] + + +def test_mirrored_region_gets_no_box(): + """镜像出来的是复制品，不能再给它标框——否则等于告诉模型那里真有缺陷。""" + width = height = 80 + box = from_yolo(0, 0.5, 0.5, 0.5, 0.5) + tile, = plan_tiles(width, height, SIZE) + mapped = remap_boxes([box], tile, width, height, SIZE) + + assert len(mapped) == 1 # 只有一个，不是四个倒影 + # 框整体落在原图内容的范围内 + assert mapped[0].x1 * SIZE >= tile.pad_left - 1e-6 + assert mapped[0].x2 * SIZE <= tile.pad_left + tile.crop_width + 1e-6 + + +def test_zero_area_box_is_ignored(): + tile, = plan_tiles(100, 100, SIZE) + degenerate = Box(0, 0.5, 0.5, 0.5, 0.5) + assert remap_boxes([degenerate], tile, 100, 100, SIZE) == [] + + +# ---------------- 反向映射 ---------------- + + +@pytest.mark.parametrize("size_wh", [(80, 60), (400, 300), (160, 160), (524, 959)]) +def test_round_trip_back_to_source(size_wh): + """片上的框映射回原图，必须落回原来的像素位置。 + + 框放在某一片裁剪区域的正中，确保它完整落在该片内——无重叠切片时， + 原图正中的框反而可能横跨切割线，不适合用来做往返校验。 + """ + width, height = size_wh + for tile in plan_tiles(width, height, SIZE, overlap=0): + cx = tile.x + tile.crop_width / 2 + cy = tile.y + tile.crop_height / 2 + half_w = min(tile.crop_width, 20) / 4 + half_h = min(tile.crop_height, 20) / 4 + px = [cx - half_w, cy - half_h, cx + half_w, cy + half_h] + box = Box(2, px[0] / width, px[1] / height, px[2] / width, px[3] / height) + + mapped, = remap_boxes([box], tile, width, height, SIZE, min_visibility=0.99) + back = to_source(mapped, tile, SIZE) + assert [back.x1, back.y1, back.x2, back.y2] == pytest.approx(px, abs=1e-6) + + +def test_score_survives_both_directions(): + tile, = plan_tiles(100, 100, SIZE) + box = Box(1, 0.4, 0.4, 0.6, 0.6, 0.83) + mapped, = remap_boxes([box], tile, 100, 100, SIZE) + assert mapped.score == 0.83 + assert to_source(mapped, tile, SIZE).score == 0.83