Files
ModelTrainEval/tests/test_prepare_label_studio.py
T
Shaun ZhouandClaude Opus 5 891f930de2 refactor: prepare_label_studio.py 的内容改动(补上 a3e08a9 漏掉的部分)
a3e08a9 只提交了 git mv 的改名,内容修改因 git add 的路径报错没有进入暂存区。
本提交补上:

- 缺省输出 runs/<source>_label_studio/,不再写到 data/;该目录同时是 Label Studio
  的 DOCUMENT_ROOT,打印的导入步骤里提醒复核期间不要删除或 --overwrite。
- 总是生成 tasks.json 与 labeling config,去掉 --label-studio 开关。
- 不再写 data.yaml:复核材料不是训练集,没有 data.yaml 就不会被误拿去训练,
  原先"train 与 val 相同"的告警也随之移除(detect.py 里的同类告警保留)。
- detect.py、configs/detect/film.yaml、README 里对旧脚本的引用改为新流程。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-15 17:26:22 +08:00

454 lines
19 KiB
Python

"""把 generate_boxes.py 的产出整理成 Label Studio 复核材料。"""
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_label_studio as pls # 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_label_studio.py", "--source", str(source), "--out", str(out), *extra]
original, sys.argv = sys.argv, argv
try:
return pls.main()
finally:
sys.argv = original
def test_flat_layout(source, tmp_path):
"""类别目录被拍平:images/ 与 labels/ 下直接是文件,一一对应(converter 不递归)。"""
out = tmp_path / "runs" / "boxes_label_studio"
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
def test_hard_links_share_inode(source, tmp_path):
"""用硬链接而不是复制,不占额外磁盘。"""
out = tmp_path / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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_empty_labels_kept_as_negatives_by_default(source, tmp_path):
"""无框的图默认保留——YOLO 里它们是压低误检的负样本。"""
out = tmp_path / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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_no_data_yaml_so_it_cannot_be_trained_on_by_mistake(source, tmp_path):
"""复核材料不是训练集:不写 data.yaml,免得没修过的弱监督框被直接拿去训练。"""
out = tmp_path / "runs" / "boxes_label_studio"
assert run(source, out) == 0
assert not (out / "data.yaml").exists()
assert not (out / "images" / "train").exists()
def test_default_output_goes_under_runs(source, tmp_path, monkeypatch):
"""缺省输出是 runs/<source>_label_studio,不再写到 data/。"""
pytest.importorskip("label_studio_sdk.converter.imports.yolo")
monkeypatch.chdir(tmp_path)
original, sys.argv = sys.argv, ["prepare_label_studio.py", "--source", str(source)]
try:
assert pls.main() == 0
finally:
sys.argv = original
assert (tmp_path / "runs" / "boxes_label_studio" / "tasks.json").is_file()
assert not (tmp_path / "data").exists()
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 / "runs" / "boxes_label_studio", "--keep-names") != 0
assert "冲突" in capsys.readouterr().out
def test_refuses_to_clobber_existing_output(source, tmp_path):
"""输出目录非空时要求显式 --overwrite。"""
out = tmp_path / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
assert run(source, out) == 0
images = _all_outputs(out, "images")
assert len(images) == 15
assert {p.suffix for p in images} == {".png"}
assert all(pls.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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 pls.is_png(fake)
out = tmp_path / "runs" / "boxes_label_studio"
assert run(source, out) == 0
output = output_of(out, "fake")
assert pls.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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio", "--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_source_data_yaml(source, tmp_path):
"""classes.txt 的行号就是 class_id,必须与来源 data.yaml 的 names 严格一致。"""
out = tmp_path / "runs" / "boxes_label_studio"
assert run(source, out) == 0
lines = (out / "classes.txt").read_text(encoding="utf-8").splitlines()
meta = yaml.safe_load((source / "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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
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 / "runs" / "boxes_label_studio"
assert run(source, out, "--empty-policy", "normal=keep") == 0
return out
def _tasks(out):
return json.loads((out / "tasks.json").read_text(encoding="utf-8"))
def test_label_studio_files_at_output_root(label_studio_out):
"""tasks.json 与 labeling config 总是生成,直接放在输出根目录(即 DOCUMENT_ROOT)。"""
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_prints_import_steps_with_real_paths(label_studio_out, capsys):
"""converter 自带提示漏了环境变量,脚本要给出带实际路径的完整步骤。"""
capsys.readouterr()
source = label_studio_out.parent / "boxes"
assert run(source, label_studio_out, "--overwrite") == 0
out = capsys.readouterr().out
assert f"LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT={label_studio_out.resolve()}" in out
assert "不要点 Sync" in out