feat(detect): 把 Label Studio 的 YOLO 导出装配成带 train/val 划分的数据集
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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
65bc6b2045
commit
998e6601d4
@@ -0,0 +1,295 @@
|
||||
#!/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())
|
||||
@@ -55,7 +55,6 @@ from __future__ import annotations
|
||||
import argparse
|
||||
import contextlib
|
||||
import csv
|
||||
import errno
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
@@ -73,7 +72,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
import yaml # noqa: E402
|
||||
|
||||
from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt # noqa: E402
|
||||
from utils import setup_logging # noqa: E402
|
||||
from utils import link_or_copy, setup_logging # noqa: E402
|
||||
|
||||
logger = None
|
||||
|
||||
@@ -116,21 +115,6 @@ def materialize_png(source: Path, destination: Path) -> str:
|
||||
return "convert"
|
||||
|
||||
|
||||
def link_or_copy(source: Path, destination: Path) -> str:
|
||||
"""硬链接；跨文件系统时退回复制。"""
|
||||
if destination.exists() or destination.is_symlink():
|
||||
destination.unlink()
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
os.link(source, destination)
|
||||
return "hard"
|
||||
except OSError as exc:
|
||||
if exc.errno != errno.EXDEV:
|
||||
raise
|
||||
shutil.copy2(source, destination)
|
||||
return "copy"
|
||||
|
||||
|
||||
class Sample(NamedTuple):
|
||||
"""一张图及其标签。
|
||||
|
||||
|
||||
Reference in New Issue
Block a user