Shaun ZhouandClaude Opus 5
012b3efc73
空标签在 YOLO 里的语义是"这张图什么都没有"，是一个明确的负样本断言，而不是 "有缺陷但没能定位出来"。原先一律保留，于是 runs/boxes_gradcam 的 339 张空标签 里有 270 张是人工标成缺陷、Grad-CAM 却失手的图（blackdot 占该类 11%、black 占 24%），它们正以"此图无缺陷"的身份进入训练，等于在教模型抑制真实缺陷。 --empty-policy 接一个全局默认加若干 `类别=策略`： keep 当负样本保留（默认）/ drop 丢弃 / full 退化成整图框 整图框需要重写标签而不是链接源文件（内容不同），且有了确定类别之后要跟着本类 一起参与分层划分。类别名或策略拼错直接报错——静默忽略会让策略看着生效其实没有。 没显式指定、又落到 keep 的类别会打一条告警。 替换掉原来的 --drop-empty：它是一刀切，会把 normal 那 69 张有价值的负样本一起扔掉。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
204 lines
8.2 KiB
Python
204 lines
8.2 KiB
Python
"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""
|
|
|
|
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 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_flattens_classes_and_splits(source, tmp_path):
|
|
"""类别目录被拍平，train/val 分开，图片与标签一一对应。"""
|
|
out = tmp_path / "data" / "boxes"
|
|
assert run(source, out) == 0
|
|
|
|
for split in ("train", "val"):
|
|
images = sorted(p.stem for p in (out / "images" / split).iterdir())
|
|
labels = sorted(p.stem for p in (out / "labels" / split).iterdir())
|
|
assert images == labels
|
|
total = sum(len(list((out / "images" / s).iterdir())) for s in ("train", "val"))
|
|
assert total == 13
|
|
# 拍平后不应再有 <类别>/ 这一层
|
|
assert not (out / "images" / "train" / "grease").exists()
|
|
|
|
|
|
def test_hard_links_share_inode(source, tmp_path):
|
|
"""用硬链接而不是复制，不占额外磁盘。"""
|
|
out = tmp_path / "data" / "boxes"
|
|
assert run(source, out) == 0
|
|
linked = next((out / "images" / "train").iterdir())
|
|
original = next(p for p in (source / "images").rglob("*.png") if p.name == linked.name)
|
|
assert linked.stat().st_ino == original.stat().st_ino
|
|
|
|
|
|
def test_small_class_present_in_both_splits(source, tmp_path):
|
|
"""只有 2 张的小类也要在验证集里留下一张，否则它的指标是空的。"""
|
|
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["train"]["whitedot"] == 1
|
|
assert per_class["val"]["whitedot"] == 1
|
|
|
|
|
|
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 not (out / "labels" / "train" / "normal_0.txt").exists()
|
|
assert not (out / "images" / "train" / "normal_0.png").exists()
|
|
|
|
|
|
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
|
|
|
|
label = next(
|
|
p for split in ("train", "val")
|
|
for p in (out / "labels" / split).iterdir()
|
|
if p.stem == "normal_0"
|
|
)
|
|
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["train"]["normal"] + per_class["val"]["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
|
|
assert (Path(meta["path"]) / meta["train"]).is_dir()
|
|
assert (Path(meta["path"]) / meta["val"]).is_dir()
|
|
|
|
|
|
def test_reproducible_split(source, tmp_path):
|
|
"""同一个种子两次运行，划分必须一致。"""
|
|
first, second = tmp_path / "a", tmp_path / "b"
|
|
assert run(source, first) == 0
|
|
assert run(source, second) == 0
|
|
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()
|
|
)
|
|
|
|
|
|
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") != 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()
|