diff --git a/README.md b/README.md index 32850f4..ea68a9c 100644 --- a/README.md +++ b/README.md @@ -621,8 +621,38 @@ data/boxes_gradcam/ 按类别分层（`--val-ratio` 默认 0.2，`--seed` 默认 42，两次运行结果一致）。分层是必需的： `black` 只有 53 张、`others` 69 张，随机切很容易让某个小类在验证集里挂零，指标就成了空的。 -无框的图（Grad-CAM 没框出东西的那 339 张，占 3%）**默认保留**——在检测训练里它们是负样本， -对压低误检有用。不要时加 `--drop-empty`。 +### 没有框的图怎么处理 + +Grad-CAM 没框出东西的有 339 张（占 3%）。它们**默认保留为负样本**，对压低误检有用。 + +但要当心：**空标签在 YOLO 里的语义是"这张图什么都没有"**——一个明确的负样本断言， +而不是"有缺陷但没能定位出来"。而这 339 张里只有 69 张真的来自 `normal`，其余 270 张是人工 +标成缺陷、Grad-CAM 却失手的图： + +| 来源类别 | 空标签 | 占该类 | +|---|---|---| +| grease | 95 | 7.8% | +| blackdot | 82 | 11.3% | +| normal | 69 | 2.6% | +| strip_h | 28 | 2.9% | +| zigzag | 23 | 2.4% | +| black | 17 | 24.3% | +| 其余 9 类 | 25 | <3% | + +照搬的话，`blackdot` 有 11% 的正样本、`black` 有 24% 会以"此图无缺陷"的身份进入训练， +等于在教模型抑制真实缺陷。用 `--empty-policy` 按来源类别分别处置： + +```bash +uv run python scripts/prepare_yolo_dataset.py --source runs/boxes_gradcam \ + --empty-policy drop normal=keep black=full edge_clean=full +``` + +- `keep`（默认）——当负样本保留。只有真正无缺陷的类别该用 +- `drop`——丢弃。缺陷类没定位出来的属于**信息缺失**，不该冒充负样本 +- `full`——退化成整图框。适合 `black`/`edge_clean` 这种描述整图属性、缺陷本就铺满全图的类别 + +第一个裸词是全局默认，后面 `类别=策略` 逐类覆盖。类别名或策略写错会直接报错而不是静默忽略。 +没有显式指定策略、又落到 `keep` 的类别会打一条告警，提醒你确认它是否真的无缺陷。 拍平之后不同类别目录下的同名文件会互相覆盖，脚本会直接报错而不是静默丢数据； 真遇上就用 `scripts/export_clusters.py --prefix-parent` 重新导出。 diff --git a/scripts/prepare_yolo_dataset.py b/scripts/prepare_yolo_dataset.py index 943b5e8..7faa536 100644 --- a/scripts/prepare_yolo_dataset.py +++ b/scripts/prepare_yolo_dataset.py @@ -25,7 +25,12 @@ data// └── data.yaml ``` -没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。 +没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。但**空标签的 +语义是"这张图什么都没有"**，不是"有缺陷但没定位出来"。来自缺陷类别的空标签如果照搬， +等于在教模型抑制真实缺陷，所以可以用 `--empty-policy` 按来源类别分别处置： + + --empty-policy drop normal=keep black=full + # 缺陷类没定位出来的直接丢弃；normal 留作负样本；整图属性的类退化成整图框 """ from __future__ import annotations @@ -38,7 +43,9 @@ import random import shutil import sys from collections import Counter, defaultdict +from collections.abc import Sequence from pathlib import Path +from typing import NamedTuple sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -49,6 +56,13 @@ from utils import setup_logging # noqa: E402 logger = None +EMPTY_POLICIES = ("keep", "drop", "full") +EMPTY_POLICY_DESC = { + "keep": "保留为负样本", + "drop": "丢弃", + "full": "退化成整图框", +} + def link_or_copy(source: Path, destination: Path) -> str: """硬链接；跨文件系统时退回复制。""" @@ -65,6 +79,25 @@ def link_or_copy(source: Path, destination: Path) -> str: return "copy" +class Sample(NamedTuple): + """一张图及其标签。 + + Attributes: + image: 图片路径。 + label: 标签路径。 + class_id: 标签里第一个框的类别下标；无框时为 None。 + group: 来源的类别目录名——也就是人工归类给这张图的类别。 + 标签为空时，这是唯一还能知道"这张图本该是什么"的线索。 + full_box: 整图框的类别下标。非 None 时 _emit 会**重写**标签而不是链接原文件。 + """ + + image: Path + label: Path + class_id: int | None + group: str + full_box: int | None = None + + def first_class_id(label_path: Path) -> int | None: """取标签里第一个框的类别下标；空标签（无框）返回 None。""" for line in label_path.read_text(encoding="utf-8").splitlines(): @@ -74,28 +107,92 @@ def first_class_id(label_path: Path) -> int | None: return None -def collect(source: Path) -> list[tuple[Path, Path, int | None]]: - """配对 images/ 与 labels/，返回 (图片, 标签, 类别下标) 列表。""" +def collect(source: Path) -> list[Sample]: + """配对 images/ 与 labels/。""" images_root, labels_root = source / "images", source / "labels" if not images_root.is_dir() or not labels_root.is_dir(): raise FileNotFoundError(f"{source} 下缺少 images/ 或 labels/，先跑 scripts/generate_boxes.py") - pairs = [] + samples = [] missing = [] for image in sorted(p for p in images_root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES): - label = labels_root / image.relative_to(images_root).with_suffix(".txt") + relative = image.relative_to(images_root) + label = labels_root / relative.with_suffix(".txt") if not label.exists(): missing.append(image) continue - pairs.append((image, label, first_class_id(label))) + group = relative.parent.name # 顶层就是图片时为 "" + samples.append(Sample(image, label, first_class_id(label), group)) if missing: raise FileNotFoundError( f"有 {len(missing)} 张图没有对应的标签，例如 {missing[0]}；" "产出目录可能不完整，重新跑一次 generate_boxes.py" ) - if not pairs: + if not samples: raise ValueError(f"{source} 下没有找到任何图片") - return pairs + return samples + + +def parse_empty_policy(tokens: Sequence[str]) -> tuple[str, dict[str, str]]: + """解析 --empty-policy：一个可选的全局默认，加若干 `类别=策略`。 + + --empty-policy drop normal=keep black=full + -> ("drop", {"normal": "keep", "black": "full"}) + """ + default, per_class = "keep", {} + for token in tokens: + if "=" in token: + name, _, policy = token.partition("=") + if policy not in EMPTY_POLICIES: + raise ValueError(f"未知策略 {policy!r}（类别 {name}）；可选 {EMPTY_POLICIES}") + per_class[name] = policy + elif token in EMPTY_POLICIES: + default = token + else: + raise ValueError( + f"看不懂 --empty-policy 的 {token!r}；" + f"要么是全局默认（{'/'.join(EMPTY_POLICIES)}），要么是 `类别=策略`" + ) + return default, per_class + + +def apply_empty_policy( + samples: Sequence[Sample], default: str, per_class: dict[str, str], names: Sequence[str] +) -> tuple[list[Sample], dict[str, tuple[str, int]]]: + """按来源类别处置空标签，返回 (保留下来的样本, {类别: (策略, 数量)})。 + + 空标签在 YOLO 里的语义是"这张图什么都没有"——一个明确的负样本断言，而不是 + "有缺陷但没定位出来"。所以来自缺陷类别的空标签如果照搬，等于在教模型抑制真实缺陷。 + """ + index = {name: i for i, name in enumerate(names)} + unknown = { + name for name in per_class if name not in index + } + if unknown: + raise ValueError( + f"--empty-policy 里的类别 {sorted(unknown)} 不在 data.yaml 中；可选 {list(names)}" + ) + + kept: list[Sample] = [] + stats: dict[str, tuple[str, int]] = {} + for sample in samples: + if sample.class_id is not None: + kept.append(sample) + continue + policy = per_class.get(sample.group, default) + counted = stats.get(sample.group, (policy, 0)) + stats[sample.group] = (policy, counted[1] + 1) + if policy == "drop": + continue + if policy == "full": + if sample.group not in index: + raise ValueError( + f"{sample.group!r} 要用整图框，但它不在 data.yaml 的类别里；可选 {list(names)}" + ) + kept.append(sample._replace(full_box=index[sample.group])) + else: + kept.append(sample) + return kept, stats def main() -> int: @@ -109,8 +206,13 @@ def main() -> int: 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="丢弃没有框的图。默认保留——YOLO 训练里它们是负样本，有助于压低误检", + "--empty-policy", nargs="*", default=[], metavar="POLICY|类别=POLICY", + help=( + "怎么处置没有框的图（标签为空）。可写一个全局默认，再按类别覆盖，如 " + "`--empty-policy drop normal=keep black=full`。" + "keep=当负样本保留（默认）、drop=丢弃、full=退化成整图框。" + "注意空标签在 YOLO 里的意思是『此图无缺陷』，缺陷类的空标签照搬会教模型抑制真实缺陷" + ), ) parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空") parser.add_argument("--dry-run", action="store_true", help="只打印计划") @@ -121,7 +223,7 @@ def main() -> int: logger = setup_logging() try: - pairs = collect(source) + samples = collect(source) except (FileNotFoundError, ValueError) as exc: logger.error("%s", exc) return 1 @@ -133,18 +235,37 @@ def main() -> int: meta = yaml.safe_load(meta_path.read_text(encoding="utf-8")) names = [meta["names"][i] for i in sorted(meta["names"])] - empty = [p for p in pairs if p[2] is None] - if args.drop_empty: - pairs = [p for p in pairs if p[2] is not None] - logger.info("来源: %s（%d 张，%d 类）", source, len(pairs), len(names)) - logger.info( - "无框图片 %d 张：%s", len(empty), "已丢弃" if args.drop_empty else "保留为负样本" - ) + logger.info("来源: %s（%d 张，%d 类）", source, len(samples), len(names)) + try: + default_policy, per_class_policy = parse_empty_policy(args.empty_policy) + samples, empty_stats = apply_empty_policy(samples, default_policy, per_class_policy, names) + except ValueError as exc: + logger.error("%s", exc) + return 1 + + total_empty = sum(count for _, count in empty_stats.values()) + if total_empty: + logger.info("无框图片 %d 张，按来源类别处置：", total_empty) + for group, (policy, count) in sorted(empty_stats.items(), key=lambda kv: -kv[1][1]): + logger.info(" %-14s %5d 张 -> %s", group or "(顶层)", count, EMPTY_POLICY_DESC[policy]) + # 只对"默认落到 keep"的类别告警——显式写了 `类别=keep` 的是用户的判断，不再唠叨 + as_negative = sorted( + g for g, (policy, _) in empty_stats.items() + if policy == "keep" and g not in per_class_policy + ) + if as_negative: + logger.warning( + "以上 %s 的空标签会以『此图无缺陷』的身份进入训练。" + "如果其中某些类其实是有缺陷、只是没能定位出来，它们会教模型抑制真实缺陷——" + "用 --empty-policy 逐类改成 drop 或 full。", + "/".join(as_negative), + ) # 按类别分层切分；无框的图单独作为一组分配，避免它们全落进同一侧 - grouped: dict[int | None, list] = defaultdict(list) - for pair in pairs: - grouped[pair[2]].append(pair) + grouped: dict[int | None, list[Sample]] = defaultdict(list) + for sample in samples: + # 退化成整图框的那些已经有了确定类别，跟着本类一起分层 + grouped[sample.full_box if sample.class_id is None else sample.class_id].append(sample) rng = random.Random(args.seed) splits: dict[str, list] = {"train": [], "val": []} @@ -156,7 +277,12 @@ def main() -> int: splits["train"].extend(order[n_val:]) counts = { - split: Counter(names[c] if c is not None else "(无框)" for _, _, c in items) + split: Counter( + names[s.class_id] if s.class_id is not None + else names[s.full_box] if s.full_box is not None + else "(无框)" + for s in items + ) for split, items in splits.items() } logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"])) @@ -178,7 +304,8 @@ def main() -> int: modes: Counter[str] = Counter() seen: dict[str, Path] = {} for split, items in splits.items(): - for image, label, _ in items: + for sample in items: + image = sample.image # 拍平之后不同类别目录下的同名文件会撞车，显式挡住 if image.name in seen: logger.error( @@ -189,7 +316,13 @@ def main() -> int: return 1 seen[image.name] = image modes[link_or_copy(image, out_dir / "images" / split / image.name)] += 1 - link_or_copy(label, out_dir / "labels" / split / f"{image.stem}.txt") + destination = out_dir / "labels" / split / f"{image.stem}.txt" + if sample.full_box is None: + link_or_copy(sample.label, destination) + else: + # 整图框：内容与源标签不同，只能新写一份，不能链接 + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(f"{sample.full_box} 0.5 0.5 1.0 1.0\n", encoding="utf-8") write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(), train="images/train", val="images/val") @@ -197,9 +330,12 @@ def main() -> int: json.dumps( { "source": str(source.resolve()), - "num_images": len(pairs), - "empty_labels": len(empty), - "dropped_empty": args.drop_empty, + "num_images": len(samples), + "empty_labels": { + group: {"policy": policy, "count": count} + for group, (policy, count) in sorted(empty_stats.items()) + }, + "empty_policy": {"default": default_policy, **per_class_policy}, "val_ratio": args.val_ratio, "seed": args.seed, "split_sizes": {k: len(v) for k, v in splits.items()}, diff --git a/tests/test_prepare_yolo_dataset.py b/tests/test_prepare_yolo_dataset.py index be69f01..0a0ac39 100644 --- a/tests/test_prepare_yolo_dataset.py +++ b/tests/test_prepare_yolo_dataset.py @@ -91,13 +91,60 @@ def test_empty_labels_kept_as_negatives_by_default(source, tmp_path): 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"] == 1 + assert summary["empty_labels"] == {"normal": {"policy": "keep", "count": 1}} assert summary["num_images"] == 13 - out_dropped = tmp_path / "data" / "dropped" - assert run(source, out_dropped, "--drop-empty") == 0 - dropped = json.loads((out_dropped / "summary.json").read_text(encoding="utf-8")) - assert dropped["num_images"] == 12 + +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):