Shaun ZhouandClaude Opus 5
998e6601d4
scripts/build_yolo_dataset.py：导出（zip 或目录）+ 复核材料里的 images/ -> data/ 下 可训练的数据集。Label Studio 的 YOLO 导出只有标签，本地存储的图片不会打包进去 （get_local_path 对 local-files 直接返回磁盘路径，YOLO_WITH_IMAGES 也一样）， 所以图片从准备复核材料的目录按文件名取，硬链接过去。 - 类别顺序照抄导出的 classes.txt：转换器把标签收进 set 再排序编号，与 labeling config 的顺序无关，标签里的 class_id 是按它写的。 - 按类别分层切分，分层用"框最多的那个类别"——人工复核后一张图可能有多个类别的框。 - 标签逐行校验：字段数、类别下标、坐标范围。坏标签训练时不会崩，只会悄悄拉低指标。 - 空标签默认保留为负样本；标签找不到图直接报错，图多于标签则告警。 顺带把 link_or_copy 提到 utils/fs.py 共用（此前 generate_boxes 与 prepare_label_studio 各有一份）。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
296 lines
12 KiB
Python
296 lines
12 KiB
Python
#!/usr/bin/env python
|
||
"""把 Label Studio 的 YOLO 导出装配成可训练的数据集，划分 train/val 放到 data/ 下。
|
||
|
||
uv run python scripts/build_yolo_dataset.py \\
|
||
--export data/film_yolo_shap.zip \\
|
||
--images runs/boxes_shap_label_studio/images \\
|
||
--out data/film_yolo
|
||
|
||
Label Studio 的 YOLO 导出**只有标签**：本地存储的图片它不会打包进去（`get_local_path` 对
|
||
local-files 直接返回磁盘路径，不复制，`YOLO_WITH_IMAGES` 也一样）。所以图片从我们自己准备
|
||
复核材料的那个目录取——文件名是一一对应的。
|
||
|
||
```
|
||
<export>.zip runs/<...>_label_studio/images/
|
||
├── classes.txt └── <类别>_<序号>.png
|
||
├── labels/<类别>_<序号>.txt
|
||
└── images/ ← 空的
|
||
```
|
||
|
||
装配结果：
|
||
|
||
```
|
||
data/<name>/
|
||
├── images/{train,val}/<类别>_<序号>.png # 硬链接，不占额外磁盘
|
||
├── labels/{train,val}/<类别>_<序号>.txt
|
||
├── classes.txt # 原样沿用导出的顺序
|
||
├── data.yaml
|
||
└── summary.json
|
||
```
|
||
|
||
**类别顺序以导出的 classes.txt 为准。** Label Studio 的转换器把标签收进 set 再排序编号，
|
||
与 labeling config 里写的顺序无关，也与我们 `generate_boxes.py` 时的顺序无关。标签里的
|
||
class_id 是按导出的顺序写的，所以 data.yaml 必须照抄，不能另排。
|
||
|
||
**空标签保留为负样本。** 导出后仍为空的图（本例是 normal，预测已在 Label Studio 里批量删掉）
|
||
表示"这张图没有缺陷"，在检测训练里是压低误检的负样本。`--drop-empty` 可以丢掉它们。
|
||
注意：没复核、也没把预测转成标注的任务同样会导成空标签，那种"空"是信息缺失而不是负样本，
|
||
装配前应先在 Label Studio 里用 Create Annotations From Predictions 补上。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
import shutil
|
||
import sys
|
||
import zipfile
|
||
from collections import Counter, defaultdict
|
||
from pathlib import Path
|
||
from typing import NamedTuple
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt, write_data_yaml # noqa: E402
|
||
from utils import link_or_copy, setup_logging # noqa: E402
|
||
|
||
logger = None
|
||
|
||
EMPTY = "(无框)"
|
||
|
||
|
||
class Sample(NamedTuple):
|
||
"""一张图、它的标签文本，以及分层划分用的类别。"""
|
||
|
||
image: Path
|
||
label_text: str
|
||
stratum: str
|
||
|
||
|
||
class ExportError(RuntimeError):
|
||
"""导出内容不自洽，装配无法继续。"""
|
||
|
||
|
||
def read_export(path: Path) -> tuple[list[str], dict[str, str]]:
|
||
"""读 Label Studio 的 YOLO 导出（zip 或已解压目录）。
|
||
|
||
Returns:
|
||
(类别名列表, {文件名主干: 标签文本})
|
||
"""
|
||
if path.is_dir():
|
||
classes_file = path / "classes.txt"
|
||
if not classes_file.exists():
|
||
raise ExportError(f"{path} 下没有 classes.txt，这不像 Label Studio 的 YOLO 导出")
|
||
classes = classes_file.read_text(encoding="utf-8").split()
|
||
labels = {
|
||
p.stem: p.read_text(encoding="utf-8")
|
||
for p in sorted((path / "labels").glob("*.txt"))
|
||
}
|
||
else:
|
||
with zipfile.ZipFile(path) as archive:
|
||
names = archive.namelist()
|
||
if "classes.txt" not in names:
|
||
raise ExportError(f"{path} 里没有 classes.txt，这不像 Label Studio 的 YOLO 导出")
|
||
classes = archive.read("classes.txt").decode("utf-8").split()
|
||
labels = {
|
||
Path(name).stem: archive.read(name).decode("utf-8")
|
||
for name in sorted(names)
|
||
if name.startswith("labels/") and name.endswith(".txt")
|
||
}
|
||
if not classes:
|
||
raise ExportError("classes.txt 是空的")
|
||
if not labels:
|
||
raise ExportError("导出里没有任何标签文件")
|
||
return classes, labels
|
||
|
||
|
||
def parse_boxes(text: str, stem: str, num_classes: int) -> list[int]:
|
||
"""校验一个标签文件，返回其中所有框的类别下标。
|
||
|
||
坐标越界或类别越界都直接报错：这类问题在训练时不会崩，只会悄悄拉低指标。
|
||
"""
|
||
class_ids = []
|
||
for number, line in enumerate(text.splitlines(), start=1):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
parts = line.split()
|
||
if len(parts) != 5:
|
||
raise ExportError(f"{stem}.txt 第 {number} 行有 {len(parts)} 个字段，YOLO 需要 5 个：{line!r}")
|
||
try:
|
||
class_id = int(parts[0])
|
||
coordinates = [float(v) for v in parts[1:]]
|
||
except ValueError as exc:
|
||
raise ExportError(f"{stem}.txt 第 {number} 行无法解析：{line!r}") from exc
|
||
if not 0 <= class_id < num_classes:
|
||
raise ExportError(f"{stem}.txt 第 {number} 行的类别下标 {class_id} 超出 0~{num_classes - 1}")
|
||
if not all(0.0 <= v <= 1.0 for v in coordinates):
|
||
raise ExportError(f"{stem}.txt 第 {number} 行的坐标不在 [0,1] 内：{coordinates}")
|
||
class_ids.append(class_id)
|
||
return class_ids
|
||
|
||
|
||
def collect(
|
||
classes: list[str], labels: dict[str, str], images_dir: Path
|
||
) -> tuple[list[Sample], list[str]]:
|
||
"""把标签与图片配对，返回 (样本列表, 多余的图片名)。"""
|
||
samples = []
|
||
missing = []
|
||
for stem, text in labels.items():
|
||
candidates = [
|
||
images_dir / f"{stem}{suffix}" for suffix in (".png", *IMAGE_SUFFIXES)
|
||
]
|
||
image = next((p for p in candidates if p.exists()), None)
|
||
if image is None:
|
||
missing.append(stem)
|
||
continue
|
||
class_ids = parse_boxes(text, stem, len(classes))
|
||
# 分层用"框最多的那个类别"：一张图可能有多个类别的框，人工复核后尤其常见
|
||
stratum = classes[Counter(class_ids).most_common(1)[0][0]] if class_ids else EMPTY
|
||
samples.append(Sample(image, text, stratum))
|
||
|
||
if missing:
|
||
raise ExportError(
|
||
f"有 {len(missing)} 个标签找不到对应的图片，例如 {missing[0]}；"
|
||
f"--images 是否指向了准备复核材料时的 images/ 目录（当前 {images_dir}）"
|
||
)
|
||
known = {s.image.stem for s in samples}
|
||
extra = sorted(
|
||
p.stem for p in images_dir.iterdir()
|
||
if p.suffix.lower() in IMAGE_SUFFIXES and p.stem not in known
|
||
)
|
||
return samples, extra
|
||
|
||
|
||
def split(samples: list[Sample], val_ratio: float, seed: int) -> dict[str, list[Sample]]:
|
||
"""按类别分层切分。小类只有几十张，随机整体切很容易让它在验证集里挂零。"""
|
||
grouped: dict[str, list[Sample]] = defaultdict(list)
|
||
for sample in samples:
|
||
grouped[sample.stratum].append(sample)
|
||
|
||
rng = random.Random(seed)
|
||
splits: dict[str, list[Sample]] = {"train": [], "val": []}
|
||
for _stratum, group in sorted(grouped.items()):
|
||
order = sorted(group, key=lambda s: s.image.name) # 先定序，结果才可复现
|
||
rng.shuffle(order)
|
||
n_val = max(1, round(len(order) * val_ratio)) if len(order) > 1 else 0
|
||
splits["val"].extend(order[:n_val])
|
||
splits["train"].extend(order[n_val:])
|
||
return splits
|
||
|
||
|
||
def main() -> int:
|
||
global logger
|
||
parser = argparse.ArgumentParser(
|
||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||
)
|
||
parser.add_argument("--export", required=True, help="Label Studio 的 YOLO 导出（zip 或目录）")
|
||
parser.add_argument(
|
||
"--images", required=True, help="图片目录，即准备复核材料时的 <...>_label_studio/images",
|
||
)
|
||
parser.add_argument("--out", required=True, help="输出目录，建议放在 data/ 下")
|
||
parser.add_argument("--val-ratio", type=float, default=0.2)
|
||
parser.add_argument("--seed", type=int, default=42)
|
||
parser.add_argument(
|
||
"--drop-empty", action="store_true",
|
||
help="丢掉空标签的图。默认保留——它们是压低误检的负样本",
|
||
)
|
||
parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
|
||
parser.add_argument("--dry-run", action="store_true", help="只打印计划")
|
||
args = parser.parse_args()
|
||
|
||
out_dir = Path(args.out)
|
||
logger = setup_logging()
|
||
|
||
try:
|
||
classes, labels = read_export(Path(args.export))
|
||
samples, extra = collect(classes, labels, Path(args.images))
|
||
except (ExportError, FileNotFoundError) as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
|
||
logger.info("导出: %s（%d 个标签，%d 类）", args.export, len(labels), len(classes))
|
||
if extra:
|
||
logger.warning(
|
||
"%s 里有 %d 张图片没有对应的标签，不会进入数据集（例如 %s）",
|
||
args.images, len(extra), extra[0],
|
||
)
|
||
|
||
empty = [s for s in samples if s.stratum == EMPTY]
|
||
if args.drop_empty:
|
||
samples = [s for s in samples if s.stratum != EMPTY]
|
||
logger.info(
|
||
"空标签 %d 张（%.1f%%）：%s",
|
||
len(empty), 100 * len(empty) / max(len(labels), 1),
|
||
"已丢弃" if args.drop_empty else "保留为负样本",
|
||
)
|
||
|
||
splits = split(samples, args.val_ratio, args.seed)
|
||
counts = {name: Counter(s.stratum for s in items) for name, items in splits.items()}
|
||
boxes = Counter()
|
||
for sample in samples:
|
||
for line in sample.label_text.splitlines():
|
||
if line.strip():
|
||
boxes[classes[int(line.split()[0])]] += 1
|
||
|
||
logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"]))
|
||
logger.info("%-12s %8s %8s %8s", "类别", "train", "val", "框数")
|
||
for name in [*classes, EMPTY]:
|
||
if counts["train"][name] or counts["val"][name]:
|
||
logger.info(
|
||
"%-12s %8d %8d %8d",
|
||
name, counts["train"][name], counts["val"][name], boxes[name],
|
||
)
|
||
|
||
if args.dry_run:
|
||
logger.info("--dry-run：未写入任何文件")
|
||
return 0
|
||
|
||
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)
|
||
|
||
modes: Counter[str] = Counter()
|
||
for name, items in splits.items():
|
||
for sample in items:
|
||
modes[link_or_copy(sample.image, out_dir / "images" / name / sample.image.name)] += 1
|
||
label_path = out_dir / "labels" / name / f"{sample.image.stem}.txt"
|
||
label_path.parent.mkdir(parents=True, exist_ok=True)
|
||
label_path.write_text(sample.label_text, encoding="utf-8")
|
||
|
||
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)
|
||
(out_dir / "summary.json").write_text(
|
||
json.dumps(
|
||
{
|
||
"export": str(Path(args.export).resolve()),
|
||
"images": str(Path(args.images).resolve()),
|
||
"num_images": len(samples),
|
||
"num_boxes": sum(boxes.values()),
|
||
"empty_labels": len(empty),
|
||
"dropped_empty": args.drop_empty,
|
||
"images_without_label": len(extra),
|
||
"val_ratio": args.val_ratio,
|
||
"seed": args.seed,
|
||
"split_sizes": {k: len(v) for k, v in splits.items()},
|
||
"per_class_images": {k: dict(v) for k, v in counts.items()},
|
||
"per_class_boxes": dict(boxes),
|
||
"link_modes": dict(modes),
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
logger.info("图片落盘方式: %s", dict(modes))
|
||
logger.info("完成 -> %s（%d 个框）", out_dir, sum(boxes.values()))
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|