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>
230 lines
8.8 KiB
Python
230 lines
8.8 KiB
Python
"""把 Label Studio 的 YOLO 导出装配成带 train/val 划分的数据集。"""
|
|
|
|
import json
|
|
import sys
|
|
import zipfile
|
|
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 build_yolo_dataset as byd # noqa: E402
|
|
|
|
#: 与 Label Studio 导出一致：类别按字母序，标签里的 class_id 就是这个顺序的下标
|
|
CLASSES = ["blackdot", "grease", "normal", "wrinkle"]
|
|
|
|
#: {文件名主干: 标签内容}。normal 是空标签（负样本），wrinkle_0001 有两个不同类别的框
|
|
LABELS = {
|
|
"blackdot_0001": "0 0.5 0.5 0.1 0.1\n",
|
|
"blackdot_0002": "0 0.2 0.3 0.1 0.1\n",
|
|
"blackdot_0003": "0 0.6 0.6 0.2 0.2\n",
|
|
"grease_0001": "1 0.4 0.4 0.3 0.3\n",
|
|
"grease_0002": "1 0.5 0.5 0.2 0.2\n",
|
|
"wrinkle_0001": "3 0.5 0.5 0.4 0.4\n0 0.1 0.1 0.05 0.05\n3 0.8 0.8 0.1 0.1\n",
|
|
"wrinkle_0002": "3 0.3 0.3 0.3 0.3\n",
|
|
"normal_0001": "",
|
|
"normal_0002": "",
|
|
"normal_0003": "",
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def export(tmp_path):
|
|
"""一个 zip 形式的导出 + 与之配套的图片目录。"""
|
|
images = tmp_path / "label_studio" / "images"
|
|
images.mkdir(parents=True)
|
|
rng = np.random.default_rng(0)
|
|
for stem in LABELS:
|
|
Image.fromarray(rng.integers(0, 255, (32, 40, 3), dtype=np.uint8)).save(
|
|
images / f"{stem}.png"
|
|
)
|
|
|
|
archive = tmp_path / "export.zip"
|
|
with zipfile.ZipFile(archive, "w") as zipped:
|
|
zipped.writestr("classes.txt", "".join(f"{c}\n" for c in CLASSES))
|
|
zipped.writestr("notes.json", "{}")
|
|
for stem, text in LABELS.items():
|
|
zipped.writestr(f"labels/{stem}.txt", text)
|
|
return archive, images
|
|
|
|
|
|
def run(export, out, *extra):
|
|
archive, images = export
|
|
argv = ["build_yolo_dataset.py", "--export", str(archive), "--images", str(images),
|
|
"--out", str(out), *extra]
|
|
original, sys.argv = sys.argv, argv
|
|
try:
|
|
return byd.main()
|
|
finally:
|
|
sys.argv = original
|
|
|
|
|
|
def outputs(out, kind):
|
|
return {p.stem: p for split in ("train", "val") for p in (out / kind / split).iterdir()}
|
|
|
|
|
|
def test_builds_a_split_dataset(export, tmp_path):
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out) == 0
|
|
|
|
images, labels = outputs(out, "images"), outputs(out, "labels")
|
|
assert set(images) == set(LABELS) == set(labels)
|
|
for stem, path in labels.items():
|
|
assert path.read_text(encoding="utf-8") == LABELS[stem]
|
|
# 图与标签必须落在同一个 split，否则验证集会混进训练图
|
|
assert path.parent.name == images[stem].parent.name
|
|
|
|
|
|
def test_class_order_comes_from_the_export(export, tmp_path):
|
|
"""类别顺序必须照抄导出的 classes.txt：标签里的 class_id 是按它写的。"""
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, 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
|
|
assert Path(meta["path"]).is_absolute()
|
|
assert meta["train"] == "images/train" and meta["val"] == "images/val"
|
|
|
|
|
|
def test_images_are_hard_linked(export, tmp_path):
|
|
_, images_dir = export
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out) == 0
|
|
|
|
linked = outputs(out, "images")["grease_0001"]
|
|
assert linked.stat().st_ino == (images_dir / "grease_0001.png").stat().st_ino
|
|
assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["link_modes"] == {
|
|
"hard": len(LABELS)
|
|
}
|
|
|
|
|
|
def test_small_classes_appear_in_both_splits(export, tmp_path):
|
|
"""分层切分：三张图的类别也要在验证集里留下一张。"""
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out) == 0
|
|
per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class_images"]
|
|
for name in ("blackdot", "(无框)"):
|
|
assert per_class["train"][name] >= 1 and per_class["val"][name] == 1
|
|
|
|
|
|
def test_multi_class_image_is_stratified_by_its_dominant_class(export, tmp_path):
|
|
"""一张图可能有多个类别的框（人工复核后常见），按框最多的类别分层。"""
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out) == 0
|
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
|
# wrinkle_0001 有 2 个 wrinkle 框 + 1 个 blackdot 框 -> 计入 wrinkle
|
|
assert summary["per_class_images"]["train"].get("wrinkle", 0) + summary[
|
|
"per_class_images"
|
|
]["val"].get("wrinkle", 0) == 2
|
|
# 框数按每个框各自的类别统计
|
|
assert summary["per_class_boxes"] == {"blackdot": 4, "grease": 2, "wrinkle": 3}
|
|
assert summary["num_boxes"] == 9
|
|
|
|
|
|
def test_empty_labels_kept_as_negatives_by_default(export, tmp_path):
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out) == 0
|
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
|
assert summary["empty_labels"] == 3
|
|
assert summary["num_images"] == 10
|
|
assert outputs(out, "labels")["normal_0001"].read_text(encoding="utf-8") == ""
|
|
|
|
dropped = tmp_path / "data" / "dropped"
|
|
assert run(export, dropped, "--drop-empty") == 0
|
|
assert json.loads((dropped / "summary.json").read_text(encoding="utf-8"))["num_images"] == 7
|
|
assert "normal_0001" not in outputs(dropped, "images")
|
|
|
|
|
|
def test_reads_an_extracted_directory_too(export, tmp_path):
|
|
"""zip 与解压后的目录应当等价。"""
|
|
archive, images = export
|
|
extracted = tmp_path / "extracted"
|
|
with zipfile.ZipFile(archive) as zipped:
|
|
zipped.extractall(extracted)
|
|
|
|
from_zip, from_dir = tmp_path / "a", tmp_path / "b"
|
|
assert run(export, from_zip) == 0
|
|
assert run((extracted, images), from_dir) == 0
|
|
assert sorted(outputs(from_zip, "images")) == sorted(outputs(from_dir, "images"))
|
|
|
|
|
|
def test_missing_image_is_an_error(export, tmp_path, capsys):
|
|
"""标签找不到图就报错，而不是悄悄少装配一张。"""
|
|
_, images = export
|
|
(images / "grease_0001.png").unlink()
|
|
assert run(export, tmp_path / "out") != 0
|
|
assert "找不到对应的图片" in capsys.readouterr().out
|
|
|
|
|
|
def test_extra_images_are_reported(export, tmp_path, capsys):
|
|
"""图片多于标签时告警——多半是 --images 指错了目录或导出不完整。"""
|
|
_, images = export
|
|
Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(images / "stray_0001.png")
|
|
assert run(export, tmp_path / "out") == 0
|
|
assert "没有对应的标签" in capsys.readouterr().out
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"text,expected",
|
|
[
|
|
("0 0.5 0.5 0.1\n", "需要 5 个"), # 字段数不对
|
|
("9 0.5 0.5 0.1 0.1\n", "超出"), # 类别下标越界
|
|
("0 0.5 1.5 0.1 0.1\n", "不在 [0,1]"), # 坐标越界
|
|
("x 0.5 0.5 0.1 0.1\n", "无法解析"),
|
|
],
|
|
)
|
|
def test_malformed_labels_are_rejected(export, tmp_path, capsys, text, expected):
|
|
"""坏标签在训练时不会崩，只会悄悄拉低指标，所以这里必须拦住。"""
|
|
archive, images = export
|
|
broken = tmp_path / "broken.zip"
|
|
with zipfile.ZipFile(archive) as source, zipfile.ZipFile(broken, "w") as target:
|
|
for item in source.namelist():
|
|
data = source.read(item)
|
|
if item == "labels/grease_0001.txt":
|
|
data = text.encode("utf-8")
|
|
target.writestr(item, data)
|
|
|
|
assert run((broken, images), tmp_path / "out") != 0
|
|
assert expected in capsys.readouterr().out
|
|
|
|
|
|
def test_not_a_label_studio_export(tmp_path, capsys):
|
|
empty = tmp_path / "empty"
|
|
(empty / "labels").mkdir(parents=True)
|
|
assert run((empty, tmp_path), tmp_path / "out") != 0
|
|
assert "classes.txt" in capsys.readouterr().out
|
|
|
|
|
|
def test_split_is_reproducible_and_ratio_respected(export, tmp_path):
|
|
first, second = tmp_path / "a", tmp_path / "b"
|
|
assert run(export, first, "--seed", "7") == 0
|
|
assert run(export, second, "--seed", "7") == 0
|
|
assert sorted(outputs(first, "images")) == sorted(outputs(second, "images"))
|
|
for split in ("train", "val"):
|
|
assert sorted(p.name for p in (first / "images" / split).iterdir()) == sorted(
|
|
p.name for p in (second / "images" / split).iterdir()
|
|
)
|
|
|
|
# 比例是按每个分层分别取整的：blackdot 3 张 -> 2 张进验证集，所以合计不是对半
|
|
other = tmp_path / "c"
|
|
assert run(export, other, "--val-ratio", "0.5") == 0
|
|
sizes = json.loads((other / "summary.json").read_text(encoding="utf-8"))["split_sizes"]
|
|
assert sizes["val"] == 6 and sizes["train"] == 4
|
|
|
|
|
|
def test_refuses_to_clobber_and_dry_run_writes_nothing(export, tmp_path):
|
|
out = tmp_path / "data" / "film"
|
|
assert run(export, out, "--dry-run") == 0
|
|
assert not out.exists()
|
|
|
|
assert run(export, out) == 0
|
|
assert run(export, out) != 0 # 已存在且非空
|
|
assert run(export, out, "--overwrite") == 0
|