Shaun ZhouandClaude Opus 5
af0897f8c7
prepare_yolo_dataset.py 的产出先导入 Label Studio 复核，划分留到复核之后，于是：
- 去掉按类别分层的 train/val 切分以及 --val-ratio / --seed，images/ 与 labels/
下直接是文件；name_map.csv 去掉 split 列，summary.json 的逐类数量改为单列。
- label_studio/ 子目录的内容上移一层：扁平的数据集本身就是 label-studio-converter
要求的 <input>/{images,labels,classes.txt}，tasks.json 与 labeling config 直接
生成在根目录，不再需要额外的硬链接视图，DOCUMENT_ROOT 即数据集根目录。
- data.yaml 的 train 与 val 只能指向同一个 images/。ultralytics 能读，但验证指标
算的是训练集本身，mAP 虚高、早停与 best.pt 的选择失真。脚本结束时告警，
detect.py 检测到 train == val 时也告警。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
450 lines
19 KiB
Python
450 lines
19 KiB
Python
"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""
|
||
|
||
import csv
|
||
import json
|
||
import sys
|
||
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 prepare_yolo_dataset as pyd # noqa: E402
|
||
|
||
CLASSES = ["grease", "normal", "whitedot"]
|
||
|
||
|
||
@pytest.fixture
|
||
def source(tmp_path):
|
||
"""模拟 generate_boxes.py 的产出：images/<类别>/ + labels/<类别>/ + data.yaml。"""
|
||
root = tmp_path / "runs" / "boxes"
|
||
rng = np.random.default_rng(0)
|
||
# whitedot 只有 2 张（小类），normal 里有 1 张无框（负样本）
|
||
for class_id, (name, count) in enumerate(zip(CLASSES, (6, 5, 2), strict=True)):
|
||
image_dir = root / "images" / name
|
||
label_dir = root / "labels" / name
|
||
image_dir.mkdir(parents=True)
|
||
label_dir.mkdir(parents=True)
|
||
for i in range(count):
|
||
stem = f"{name}_{i}"
|
||
Image.fromarray(rng.integers(0, 255, (32, 32, 3), dtype=np.uint8)).save(
|
||
image_dir / f"{stem}.png"
|
||
)
|
||
empty = name == "normal" and i == 0
|
||
(label_dir / f"{stem}.txt").write_text(
|
||
"" if empty else f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8"
|
||
)
|
||
(root / "data.yaml").write_text(
|
||
yaml.safe_dump({"nc": len(CLASSES), "names": dict(enumerate(CLASSES))}), encoding="utf-8"
|
||
)
|
||
return root
|
||
|
||
|
||
def output_of(out, source_stem):
|
||
"""通过 name_map.csv 找源文件对应的输出图片——重命名后只能这么找。"""
|
||
with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle:
|
||
for row in csv.DictReader(handle):
|
||
if Path(row["source"]).stem == source_stem:
|
||
return out / "images" / row["name"]
|
||
return None
|
||
|
||
|
||
def run(source, out, *extra):
|
||
argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra]
|
||
original, sys.argv = sys.argv, argv
|
||
try:
|
||
return pyd.main()
|
||
finally:
|
||
sys.argv = original
|
||
|
||
|
||
def test_flat_layout_without_split(source, tmp_path):
|
||
"""类别目录被拍平、不划分 train/val：images/ 与 labels/ 下直接是文件，一一对应。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
|
||
for kind in ("images", "labels"):
|
||
assert not [p for p in (out / kind).iterdir() if p.is_dir()]
|
||
images = sorted(p.stem for p in (out / "images").iterdir())
|
||
labels = sorted(p.stem for p in (out / "labels").iterdir())
|
||
assert images == labels
|
||
assert len(images) == 13
|
||
assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["split"] is False
|
||
|
||
|
||
def test_hard_links_share_inode(source, tmp_path):
|
||
"""用硬链接而不是复制，不占额外磁盘。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
original = source / "images" / "grease" / "grease_0.png"
|
||
linked = output_of(out, "grease_0")
|
||
assert linked.stat().st_ino == original.stat().st_ino
|
||
|
||
|
||
def test_per_class_counts_in_summary(source, tmp_path):
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class"]
|
||
assert per_class == {"grease": 6, "normal": 4, "(无框)": 1, "whitedot": 2}
|
||
|
||
|
||
def test_warns_that_data_is_not_split(source, tmp_path, capsys):
|
||
"""train 与 val 指向同一目录，直接训练验证指标会虚高，必须提醒。"""
|
||
assert run(source, tmp_path / "data" / "boxes") == 0
|
||
assert "未划分" in capsys.readouterr().out
|
||
|
||
|
||
def test_empty_labels_kept_as_negatives_by_default(source, tmp_path):
|
||
"""无框的图默认保留——YOLO 里它们是压低误检的负样本。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||
assert summary["empty_labels"] == {"normal": {"policy": "keep", "count": 1}}
|
||
assert summary["num_images"] == 13
|
||
|
||
|
||
def test_empty_policy_drop_removes_the_image(source, tmp_path):
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--empty-policy", "drop") == 0
|
||
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||
assert summary["num_images"] == 12
|
||
assert summary["empty_labels"]["normal"]["policy"] == "drop"
|
||
assert output_of(out, "normal_0") is None
|
||
|
||
|
||
def test_empty_policy_is_per_class(source, tmp_path):
|
||
"""缺陷类的空标签丢掉，normal 的留作负样本——两种处置能共存。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--empty-policy", "drop", "normal=keep") == 0
|
||
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||
assert summary["empty_labels"]["normal"]["policy"] == "keep"
|
||
assert summary["num_images"] == 13
|
||
|
||
|
||
def test_empty_policy_full_writes_a_whole_image_box(source, tmp_path):
|
||
"""整图属性的类别退化成整图框：标签要重写，不能链接源文件（内容不同）。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--empty-policy", "normal=full") == 0
|
||
|
||
image = output_of(out, "normal_0")
|
||
label = out / "labels" / f"{image.stem}.txt"
|
||
assert label.read_text(encoding="utf-8").split() == ["1", "0.5", "0.5", "1.0", "1.0"]
|
||
# 重写而不是硬链接，源标签必须还是空的
|
||
assert (source / "labels" / "normal" / "normal_0.txt").read_text(encoding="utf-8") == ""
|
||
|
||
# 有了确定类别之后，计入 normal 而不是"(无框)"
|
||
per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class"]
|
||
assert per_class["normal"] == 5
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"tokens,expected",
|
||
[
|
||
(["drop", "nromal=keep"], "不在 data.yaml"), # 类别拼错
|
||
(["normal=kep"], "未知策略"), # 策略拼错
|
||
(["dorp"], "看不懂"), # 全局默认拼错
|
||
],
|
||
)
|
||
def test_empty_policy_rejects_typos(source, tmp_path, capsys, tokens, expected):
|
||
"""静默忽略拼错的类别名会让策略看起来生效了、其实没有。"""
|
||
assert run(source, tmp_path / "out", "--empty-policy", *tokens) != 0
|
||
assert expected in capsys.readouterr().out
|
||
|
||
|
||
def test_data_yaml_uses_absolute_path_and_source_names(source, tmp_path):
|
||
"""data.yaml 要能被 ultralytics 从任意工作目录读到，类别顺序沿用来源。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
|
||
assert Path(meta["path"]).is_absolute()
|
||
assert meta["nc"] == 3
|
||
assert [meta["names"][i] for i in sorted(meta["names"])] == CLASSES
|
||
# 不划分：两者都指向扁平的 images/
|
||
assert meta["train"] == meta["val"] == "images"
|
||
assert (Path(meta["path"]) / meta["train"]).is_dir()
|
||
|
||
|
||
def test_name_collision_after_flattening_is_reported(source, tmp_path, capsys):
|
||
"""保留原名时，不同类别目录下的同名文件拍平后会互相覆盖——必须报错而不是静默丢数据。"""
|
||
duplicate = source / "images" / "normal" / "grease_0.png"
|
||
Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(duplicate)
|
||
(source / "labels" / "normal" / "grease_0.txt").write_text("1 0.5 0.5 0.2 0.2\n", encoding="utf-8")
|
||
|
||
assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0
|
||
assert "冲突" in capsys.readouterr().out
|
||
|
||
|
||
def test_refuses_to_clobber_existing_output(source, tmp_path):
|
||
"""输出目录非空时要求显式 --overwrite。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
out.mkdir(parents=True)
|
||
(out / "keep.txt").write_text("x", encoding="utf-8")
|
||
assert run(source, out) != 0
|
||
assert (out / "keep.txt").exists()
|
||
assert run(source, out, "--overwrite") == 0
|
||
assert not (out / "keep.txt").exists()
|
||
|
||
|
||
def test_missing_label_is_an_error(source, tmp_path):
|
||
"""图片多于标签说明来源目录不完整，不能悄悄跳过。"""
|
||
(source / "labels" / "grease" / "grease_0.txt").unlink()
|
||
assert run(source, tmp_path / "out") != 0
|
||
|
||
|
||
def test_dry_run_writes_nothing(source, tmp_path):
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--dry-run") == 0
|
||
assert not out.exists()
|
||
|
||
|
||
def _all_outputs(out, kind):
|
||
return list((out / kind).iterdir())
|
||
|
||
|
||
def _add_sample(source, group, stem, suffix, class_id, *, fmt=None, mode="RGB"):
|
||
"""往来源目录里加一张指定格式的图 + 一个非空标签。"""
|
||
size = (32, 24)
|
||
array = np.random.default_rng(1).integers(0, 255, (size[1], size[0], 3), dtype=np.uint8)
|
||
image = Image.fromarray(array)
|
||
if mode != "RGB":
|
||
image = image.convert(mode)
|
||
path = source / "images" / group / f"{stem}{suffix}"
|
||
image.save(path, format=fmt)
|
||
(source / "labels" / group / f"{stem}.txt").write_text(
|
||
f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8"
|
||
)
|
||
return path, image
|
||
|
||
|
||
def test_all_output_images_are_png(source, tmp_path):
|
||
"""BMP / JPEG 来源都要落成真正的 PNG——后缀和文件头都得是。"""
|
||
_add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
|
||
_add_sample(source, "normal", "jpg_img", ".jpg", 1, fmt="JPEG")
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
|
||
images = _all_outputs(out, "images")
|
||
assert len(images) == 15
|
||
assert {p.suffix for p in images} == {".png"}
|
||
assert all(pyd.is_png(p) for p in images)
|
||
# 标签跟着改名后的图片走，一一对应
|
||
assert sorted(p.stem for p in images) == sorted(p.stem for p in _all_outputs(out, "labels"))
|
||
|
||
|
||
def test_bmp_conversion_is_lossless(source, tmp_path):
|
||
"""转换不能缩放或有损压缩，否则归一化的框就不再对应原来的像素。"""
|
||
_, original = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
|
||
converted = output_of(out, "bmp_img")
|
||
with Image.open(converted) as image:
|
||
assert image.size == original.size
|
||
assert np.array_equal(np.asarray(image), np.asarray(original))
|
||
|
||
|
||
def test_png_sources_are_still_hard_linked(source, tmp_path):
|
||
"""只有非 PNG 才转换；PNG 照旧硬链接，不占额外磁盘。"""
|
||
bmp, _ = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
|
||
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||
assert summary["link_modes"] == {"hard": 13, "convert": 1}
|
||
converted = output_of(out, "bmp_img")
|
||
assert converted.stat().st_ino != bmp.stat().st_ino
|
||
assert bmp.read_bytes()[:2] == b"BM" # 源文件没被动过
|
||
|
||
|
||
def test_png_suffix_with_non_png_content_is_converted(source, tmp_path):
|
||
"""按文件头判断而不是后缀：写着 .png 其实是 JPEG 的也要转。"""
|
||
fake, _ = _add_sample(source, "grease", "fake", ".png", 0, fmt="JPEG")
|
||
assert not pyd.is_png(fake)
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
output = output_of(out, "fake")
|
||
assert pyd.is_png(output)
|
||
|
||
|
||
@pytest.mark.parametrize("mode", ["L", "RGBA", "CMYK"])
|
||
def test_conversion_handles_pixel_modes(source, tmp_path, mode):
|
||
"""灰度 / 带透明度原样保留；PNG 存不下的 CMYK 转成 RGB。"""
|
||
_add_sample(source, "grease", "moded", ".tif", 0, fmt="TIFF", mode=mode)
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
output = output_of(out, "moded")
|
||
with Image.open(output) as image:
|
||
assert image.format == "PNG"
|
||
assert image.mode == ("RGB" if mode == "CMYK" else mode)
|
||
|
||
|
||
def test_same_stem_different_suffix_collides(source, tmp_path, capsys):
|
||
"""grease_0.png 与另一目录的 grease_0.bmp 统一后缀后会撞成同一个文件，必须报错。"""
|
||
_add_sample(source, "normal", "grease_0", ".bmp", 1, fmt="BMP")
|
||
assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0
|
||
assert "冲突" in capsys.readouterr().out
|
||
|
||
|
||
def test_dry_run_reports_conversion_count(source, tmp_path, capsys):
|
||
_add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP")
|
||
assert run(source, tmp_path / "out", "--dry-run") == 0
|
||
assert "1 张不是 PNG" in capsys.readouterr().out
|
||
|
||
|
||
def test_classes_txt_matches_data_yaml(source, tmp_path):
|
||
"""classes.txt 的行号就是 class_id，必须与 data.yaml 的 names 严格一致。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
|
||
lines = (out / "classes.txt").read_text(encoding="utf-8").splitlines()
|
||
meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
|
||
assert lines == CLASSES
|
||
assert lines == [meta["names"][i] for i in range(meta["nc"])]
|
||
assert (out / "classes.txt").read_text(encoding="utf-8").endswith("\n")
|
||
|
||
|
||
def test_classes_txt_rejects_names_that_break_line_numbering(tmp_path):
|
||
from datasets.yolo import write_classes_txt
|
||
|
||
with pytest.raises(ValueError, match="换行"):
|
||
write_classes_txt(tmp_path / "classes.txt", ["ok", "bad\nname"])
|
||
with pytest.raises(ValueError, match="为空"):
|
||
write_classes_txt(tmp_path / "classes.txt", ["ok", " "])
|
||
|
||
|
||
def _stems(out):
|
||
return sorted(p.stem for p in _all_outputs(out, "images"))
|
||
|
||
|
||
def test_files_renamed_by_class(source, tmp_path):
|
||
"""默认按 <类别>_<序号> 重新编号，序号在类内从 1 连续编到该类总数。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
assert _stems(out) == (
|
||
[f"grease_{i:04d}" for i in range(1, 7)]
|
||
+ [f"normal_{i:04d}" for i in range(1, 6)]
|
||
+ [f"whitedot_{i:04d}" for i in range(1, 3)]
|
||
)
|
||
assert _stems(out) == sorted(p.stem for p in _all_outputs(out, "labels"))
|
||
|
||
|
||
def test_renamed_file_class_matches_its_label(source, tmp_path):
|
||
"""文件名里的类别必须与标签第一列一致，否则改名反倒制造了误导。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
for label in _all_outputs(out, "labels"):
|
||
text = label.read_text(encoding="utf-8").split()
|
||
if text:
|
||
assert label.stem.rsplit("_", 1)[0] == CLASSES[int(text[0])]
|
||
|
||
|
||
def test_name_map_traces_back_to_source(source, tmp_path):
|
||
"""name_map.csv 覆盖每一张输出，且源路径真实存在、内容一致。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle:
|
||
rows = list(csv.DictReader(handle))
|
||
assert len(rows) == 13
|
||
for row in rows:
|
||
output = out / "images" / row["name"]
|
||
assert output.read_bytes() == Path(row["source"]).read_bytes()
|
||
assert row["name"].startswith(row["class"] + "_")
|
||
|
||
|
||
def test_names_stable_across_reruns(source, tmp_path):
|
||
"""重跑同一份来源，同一张源图拿到同一个名字，已做的复核记录才对得上。"""
|
||
first, second = tmp_path / "a", tmp_path / "b"
|
||
assert run(source, first) == 0
|
||
assert run(source, second) == 0
|
||
for stem in ("grease_0", "grease_5", "normal_3", "whitedot_1"):
|
||
assert output_of(first, stem).name == output_of(second, stem).name
|
||
|
||
|
||
def test_keep_names_preserves_source_stems(source, tmp_path):
|
||
"""--keep-names 保留原名，只是后缀统一成 .png。"""
|
||
_add_sample(source, "grease", "原名_中文", ".bmp", 0, fmt="BMP")
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--keep-names") == 0
|
||
assert "原名_中文" in _stems(out)
|
||
assert output_of(out, "原名_中文").name == "原名_中文.png"
|
||
|
||
|
||
def test_kept_empty_labels_named_after_source_folder(source, tmp_path):
|
||
"""空标签没有 class_id，按来源目录名编号；与同类有框的图共用序号，不会撞名。"""
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
assert output_of(out, "normal_0").stem.startswith("normal_")
|
||
assert len([s for s in _stems(out) if s.startswith("normal_")]) == 5
|
||
|
||
|
||
@pytest.fixture
|
||
def label_studio_out(source, tmp_path):
|
||
pytest.importorskip("label_studio_sdk.converter.imports.yolo")
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out, "--empty-policy", "normal=keep", "--label-studio") == 0
|
||
return out
|
||
|
||
|
||
def _tasks(out):
|
||
return json.loads((out / "tasks.json").read_text(encoding="utf-8"))
|
||
|
||
|
||
def test_label_studio_files_at_dataset_root(label_studio_out):
|
||
"""tasks.json 与 labeling config 直接放在数据集根目录，不再有 label_studio/ 子目录。"""
|
||
assert (label_studio_out / "tasks.json").is_file()
|
||
assert (label_studio_out / "tasks.label_config.xml").is_file()
|
||
assert not (label_studio_out / "label_studio").exists()
|
||
|
||
|
||
def test_label_studio_tasks_are_predictions_with_urls(label_studio_out):
|
||
"""每张图一个任务；框作为 predictions，图片 URL 相对 DOCUMENT_ROOT。"""
|
||
tasks = _tasks(label_studio_out)
|
||
assert len(tasks) == 13
|
||
for task in tasks:
|
||
assert "annotations" not in task
|
||
assert task["data"]["image"].startswith("/data/local-files/?d=images/")
|
||
name = task["data"]["image"].rsplit("/", 1)[1]
|
||
assert (label_studio_out / "images" / name).exists()
|
||
|
||
|
||
def test_label_studio_boxes_round_trip_to_yolo(label_studio_out):
|
||
"""Label Studio 的左上角百分比坐标反算回 YOLO 中心点坐标，必须与标签一致。"""
|
||
view = label_studio_out
|
||
checked = 0
|
||
for task in _tasks(label_studio_out):
|
||
name = Path(task["data"]["image"].rsplit("/", 1)[1]).stem
|
||
lines = (view / "labels" / f"{name}.txt").read_text(encoding="utf-8").split()
|
||
results = task["predictions"][0]["result"]
|
||
if not lines:
|
||
assert results == [] # 保留下来的空标签：任务里没有框
|
||
continue
|
||
value = results[0]["value"]
|
||
class_id, cx, cy, w, h = lines
|
||
assert value["rectanglelabels"] == [CLASSES[int(class_id)]]
|
||
assert value["x"] / 100 + value["width"] / 200 == pytest.approx(float(cx))
|
||
assert value["y"] / 100 + value["height"] / 200 == pytest.approx(float(cy))
|
||
assert value["width"] / 100 == pytest.approx(float(w))
|
||
checked += 1
|
||
assert checked == 12
|
||
|
||
|
||
def test_label_studio_config_and_summary(label_studio_out):
|
||
config = (label_studio_out / "tasks.label_config.xml").read_text(encoding="utf-8")
|
||
for name in CLASSES:
|
||
assert f'value="{name}"' in config
|
||
summary = json.loads((label_studio_out / "summary.json").read_text(encoding="utf-8"))
|
||
assert summary["label_studio"]["tasks"] == 13
|
||
assert Path(summary["label_studio"]["document_root"]) == label_studio_out.resolve()
|
||
|
||
|
||
def test_label_studio_not_built_by_default(source, tmp_path):
|
||
out = tmp_path / "data" / "boxes"
|
||
assert run(source, out) == 0
|
||
assert not (out / "tasks.json").exists()
|
||
assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["label_studio"] is None
|