"""把数据集切成固定边长的方片。""" 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()