#!/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())