Shaun ZhouandClaude Opus 5
23562448ba
原先把"有一边小于 160"记为填充、否则记为裁剪，但 400x100 这种既要横向裁剪 又要纵向填充，会被算成纯填充。改成两个独立计数并说明可同时发生。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
229 lines
9.6 KiB
Python
229 lines
9.6 KiB
Python
#!/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)
|
||
# 两者可以同时发生：400x100 的图既要横向裁剪又要纵向填充
|
||
if width < args.size or height < args.size:
|
||
sizes_seen["需要填充"] += 1
|
||
if width > args.size or height > args.size:
|
||
sizes_seen["需要裁剪"] += 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 张有维度超出需裁剪，两者可同时发生）",
|
||
stats["源图"], sizes_seen["需要填充"], sizes_seen["需要裁剪"],
|
||
)
|
||
logger.info(
|
||
"产出片 %d（有框 %d，空 %d；另按配额丢弃空片 %d）",
|
||
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())
|