refactor: prepare_yolo_dataset.py 改为 prepare_label_studio.py，产出放到 runs/

弱监督框要先在 Label Studio 里人工修正才能训练，这个脚本实际上只负责准备复核材料，名字和产出位置都按这个职责调整： - 改名 scripts/prepare_label_studio.py（git mv 保留历史），测试同步改名。 - 缺省输出 runs/<source>_label_studio/，不再写到 data/；该目录同时是 Label Studio 的 DOCUMENT_ROOT，打印的导入步骤里提醒复核期间不要删除或 --overwrite。 - 总是生成 tasks.json 与 labeling config，去掉 --label-studio 开关。 - 不再写 data.yaml：复核材料不是训练集，没有 data.yaml 就不会被误拿去训练，原先"train 与 val 相同"的告警也随之移除（detect.py 里的同类告警保留）。修正后的标签从 Label Studio 导出、划分 train/val 到 data/ 的脚本另写。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-15 17:26:05 +08:00
parent af0897f8c7
commit a3e08a9bbc
2 changed files with 0 additions and 0 deletions
@@ -0,0 +1,517 @@
+#!/usr/bin/env python
+"""把 generate_boxes.py 的产出整理成扁平的 YOLO 数据集，放到 data/ 下。
+
+    uv run python scripts/prepare_yolo_dataset.py --source runs/boxes_gradcam
+
+generate_boxes.py 的产出是**按类别**分目录的（便于人工核对）：
+
+```
+runs/boxes_gradcam/
+├── images/<类别>/<图名>.png
+└── labels/<类别>/<图名>.txt
+```
+
+而 YOLO 的类别信息在标签文件的第一列里，不靠目录名。这一步做三件事：
+
+1. 把按类别分的目录**拍平**，不划分 train/val——这份数据集先导入 Label Studio 复核，
+   划分留到复核之后；
+2. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式
+   （pearlescent / blue / uv 原始数据都是 BMP）无损转成 PNG。不缩放像素，
+   归一化的框坐标原样有效。
+
+3. 文件**按类别重新编号**：`<类别>_<序号>.png`，序号在类内按源路径排序，重跑结果一致。
+   原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里，不会丢。
+   想保留原名加 `--keep-names`。
+
+```
+data/<name>/
+├── images/<类别>_<序号>.png
+├── labels/<类别>_<序号>.txt
+├── classes.txt              # 每行一个类别名，行号即 class_id
+├── name_map.csv             # 新名 ↔ 源文件对照
+├── data.yaml                # train 与 val 都指向 images/，不能直接用于评估
+├── summary.json
+├── tasks.json               # 仅 --label-studio：框作为 predictions
+└── tasks.label_config.xml   # 仅 --label-studio
+```
+
+这个结构正是 label-studio-converter 的 YOLO 导入要求的 `<input>/{images,labels,classes.txt}`，
+所以 `--label-studio` 直接在数据集根目录上转换，DOCUMENT_ROOT 就是数据集根目录。
+
+没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。但**空标签的
+语义是"这张图什么都没有"**，不是"有缺陷但没定位出来"。来自缺陷类别的空标签如果照搬，
+等于在教模型抑制真实缺陷，所以可以用 `--empty-policy` 按来源类别分别处置：
+
+    --empty-policy drop normal=keep black=full
+    # 缺陷类没定位出来的直接丢弃；normal 留作负样本；整图属性的类退化成整图框
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import csv
+import errno
+import io
+import json
+import logging
+import os
+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]))
+
+import yaml  # noqa: E402
+
+from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt, write_data_yaml  # noqa: E402
+from utils import setup_logging  # noqa: E402
+
+logger = None
+
+EMPTY_POLICIES = ("keep", "drop", "full")
+EMPTY_POLICY_DESC = {
+    "keep": "保留为负样本",
+    "drop": "丢弃",
+    "full": "退化成整图框",
+}
+
+
+PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+# PNG 能原样存下的像素模式；其余（CMYK、YCbCr 等）先转 RGB
+PNG_MODES = frozenset({"1", "L", "LA", "I", "I;16", "P", "RGB", "RGBA"})
+
+
+def is_png(path: Path) -> bool:
+    """看文件头而不是后缀：后缀写着 .png、内容却是 BMP/JPEG 的也要转。"""
+    with path.open("rb") as handle:
+        return handle.read(len(PNG_SIGNATURE)) == PNG_SIGNATURE
+
+
+def materialize_png(source: Path, destination: Path) -> str:
+    """把图片落成 PNG：本来就是 PNG 的硬链接，其余格式无损转换。
+
+    返回 "hard" / "copy" / "convert"。转换不缩放像素，归一化的框坐标原样有效。
+    """
+    if is_png(source):
+        return link_or_copy(source, destination)
+
+    from PIL import Image
+
+    if destination.exists() or destination.is_symlink():
+        destination.unlink()
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    with Image.open(source) as image:
+        if image.mode not in PNG_MODES:
+            image = image.convert("RGB")
+        image.save(destination, format="PNG")
+    return "convert"
+
+
+def link_or_copy(source: Path, destination: Path) -> str:
+    """硬链接；跨文件系统时退回复制。"""
+    if destination.exists() or destination.is_symlink():
+        destination.unlink()
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        os.link(source, destination)
+        return "hard"
+    except OSError as exc:
+        if exc.errno != errno.EXDEV:
+            raise
+        shutil.copy2(source, destination)
+        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():
+        line = line.strip()
+        if line:
+            return int(float(line.split()[0]))
+    return None
+
+
+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")
+
+    samples = []
+    missing = []
+    for image in sorted(p for p in images_root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES):
+        relative = image.relative_to(images_root)
+        label = labels_root / relative.with_suffix(".txt")
+        if not label.exists():
+            missing.append(image)
+            continue
+        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 samples:
+        raise ValueError(f"{source} 下没有找到任何图片")
+    return samples
+
+
+def class_of(sample: Sample, names: Sequence[str]) -> str:
+    """样本在输出里归属的类别名：优先看标签内容，空标签才退回来源目录名。
+
+    以标签为准，重命名后的文件名就永远不会与标签里的 class_id 矛盾。
+    """
+    if sample.class_id is not None:
+        return names[sample.class_id]
+    if sample.full_box is not None:
+        return names[sample.full_box]
+    return sample.group or "background"
+
+
+def assign_names(samples: Sequence[Sample], names: Sequence[str]) -> dict[Path, str]:
+    """按类别重新编号，返回 {源图片: 新文件名主干}，如 `blackdot_0042`。
+
+    序号在类内按源路径排序：同一份来源重跑，同一张图拿到的名字不变，
+    已经做过的人工复核记录还能对得上。
+    """
+    by_class: dict[str, list[Sample]] = defaultdict(list)
+    for sample in samples:
+        by_class[class_of(sample, names)].append(sample)
+    width = max(4, len(str(max((len(v) for v in by_class.values()), default=0))))
+    mapping = {}
+    for name, items in by_class.items():
+        for number, sample in enumerate(sorted(items, key=lambda s: str(s.image)), start=1):
+            mapping[sample.image] = f"{name}_{number:0{width}d}"
+    return mapping
+
+
+def build_label_studio_tasks(out_dir: Path, num_images: int) -> dict:
+    """在数据集根目录上调用 label-studio-converter，生成 tasks.json 与 labeling config。
+
+    converter 的 YOLO 导入只认 `<input>/{images,labels,classes.txt}`，且 os.listdir 不递归——
+    这正是本脚本产出的扁平结构。图片是硬链接而不是软链接也有讲究：Label Studio 本地存储
+    会检查文件是否在 DOCUMENT_ROOT 之下，软链接解析后可能落到外面。
+
+    框以 predictions 导入：它们是 SHAP / Grad-CAM 生成的，不是人工标注。审核员修改提交
+    之后才成为 annotation，Label Studio 导出时也只导出 annotation，没审过的框不会混进去。
+    """
+    try:
+        from label_studio_sdk.converter.imports.yolo import convert_yolo_to_ls
+    except ImportError as exc:
+        raise RuntimeError("--label-studio 需要 label-studio（uv sync 后可用）") from exc
+
+    tasks_file = out_dir / "tasks.json"
+    # converter 会往 stdout 打一份缺了环境变量的导入说明，还会刷一串 root logger 的 INFO；
+    # 静音它们，由 main() 打印带实际路径的完整步骤
+    root_logger = logging.getLogger()
+    previous_level = root_logger.level
+    root_logger.setLevel(logging.WARNING)
+    try:
+        with contextlib.redirect_stdout(io.StringIO()):
+            # 不传 image_dims：图片尺寸不统一，传了会把同一个宽高写进每一条标注
+            convert_yolo_to_ls(
+                input_dir=str(out_dir),
+                out_file=str(tasks_file),
+                out_type="predictions",
+                image_ext=".png",
+            )
+    finally:
+        root_logger.setLevel(previous_level)
+    tasks = json.loads(tasks_file.read_text(encoding="utf-8"))
+    if len(tasks) != num_images:
+        raise RuntimeError(
+            f"label-studio-converter 只生成了 {len(tasks)} 个任务，应为 {num_images}；"
+            f"检查 {out_dir / 'images'}"
+        )
+    return {
+        "tasks": len(tasks),
+        "document_root": str(out_dir.resolve()),
+        "tasks_file": str(tasks_file.resolve()),
+        "label_config": str((out_dir / "tasks.label_config.xml").resolve()),
+    }
+
+
+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:
+    global logger
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--source", required=True, help="generate_boxes.py 的产出目录")
+    parser.add_argument("--out", default=None, help="输出目录；缺省 data/<source 目录名>")
+    parser.add_argument("--data-root", default="data", help="自动推导时的根目录，默认 data")
+    parser.add_argument(
+        "--empty-policy", nargs="*", default=[], metavar="POLICY|类别=POLICY",
+        help=(
+            "怎么处置没有框的图（标签为空）。可写一个全局默认，再按类别覆盖，如 "
+            "`--empty-policy drop normal=keep black=full`。"
+            "keep=当负样本保留（默认）、drop=丢弃、full=退化成整图框。"
+            "注意空标签在 YOLO 里的意思是『此图无缺陷』，缺陷类的空标签照搬会教模型抑制真实缺陷"
+        ),
+    )
+    parser.add_argument(
+        "--keep-names", action="store_true",
+        help="保留源文件名，不按 <类别>_<序号> 重新编号（源文件名需去掉后缀后全局唯一）",
+    )
+    parser.add_argument(
+        "--label-studio", action="store_true",
+        help="额外在数据集根目录生成 tasks.json 与 labeling config，可直接导入 Label Studio 复核",
+    )
+    parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
+    parser.add_argument("--dry-run", action="store_true", help="只打印计划")
+    args = parser.parse_args()
+
+    source = Path(args.source)
+    out_dir = Path(args.out) if args.out else Path(args.data_root) / source.name
+    logger = setup_logging()
+
+    try:
+        samples = collect(source)
+    except (FileNotFoundError, ValueError) as exc:
+        logger.error("%s", exc)
+        return 1
+
+    meta_path = source / "data.yaml"
+    if not meta_path.exists():
+        logger.error("找不到 %s", meta_path)
+        return 1
+    meta = yaml.safe_load(meta_path.read_text(encoding="utf-8"))
+    names = [meta["names"][i] for i in sorted(meta["names"])]
+
+    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),
+            )
+
+    counts = 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 samples
+    )
+    logger.info("%-14s %8s", "类别", "数量")
+    for name in names + ["(无框)"]:
+        if counts[name]:
+            logger.info("%-14s %8d", name, counts[name])
+
+    to_convert = sum(not is_png(sample.image) for sample in samples)
+    if to_convert:
+        logger.info(
+            "其中 %d 张不是 PNG，会无损转成 PNG（无法硬链接，占额外磁盘）；其余 %d 张硬链接",
+            to_convert, len(samples) - to_convert,
+        )
+
+    if args.dry_run:
+        if args.label_studio:
+            logger.info("--label-studio：会在 %s 下生成 tasks.json", out_dir)
+        logger.info("--dry-run：未写入任何文件")
+        return 0
+
+    if out_dir.exists() and any(out_dir.iterdir()):
+        if not args.overwrite:
+            logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
+            return 1
+        shutil.rmtree(out_dir)
+
+    renamed = None if args.keep_names else assign_names(samples, names)
+    modes: Counter[str] = Counter()
+    seen: dict[str, Path] = {}
+    name_rows = []
+    for sample in samples:
+        image = sample.image
+        stem = image.stem if renamed is None else renamed[image]
+        output_name = f"{stem}.png"
+        # 保留原名时，拍平并统一后缀后不同目录下的 a.png 与 a.bmp 会撞成同一个 a.png；
+        # 重新编号时只有类别名仅大小写不同才会撞。按小写比较：macOS 默认文件系统
+        # 不区分大小写，A.png 与 a.png 也会互相覆盖。
+        key = output_name.lower()
+        if key in seen:
+            hint = (
+                "源数据集的文件名（去掉后缀后）需要全局唯一，"
+                "去掉 --keep-names 按类别重新编号，"
+                "或用 scripts/export_clusters.py --prefix-parent 重新导出"
+                if renamed is None
+                else "data.yaml 里有只差大小写的类别名，请先改成互不相同"
+            )
+            logger.error("输出文件名冲突：%s 与 %s 都会落成 %s。%s", seen[key], image, output_name, hint)
+            return 1
+        seen[key] = image
+        modes[materialize_png(image, out_dir / "images" / output_name)] += 1
+        name_rows.append((output_name, class_of(sample, names), str(image.resolve())))
+        destination = out_dir / "labels" / f"{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")
+
+    # 不划分时 train 与 val 只能指向同一目录：ultralytics 能读，但验证指标算的是训练集本身
+    write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(), train="images", val="images")
+    write_classes_txt(out_dir / "classes.txt", names)
+
+    label_studio = None
+    if args.label_studio:
+        try:
+            label_studio = build_label_studio_tasks(out_dir, len(samples))
+        except RuntimeError as exc:
+            logger.error("%s", exc)
+            return 1
+    # utf-8-sig：带 BOM，Excel / Numbers 直接打开中文源路径不乱码
+    with (out_dir / "name_map.csv").open("w", encoding="utf-8-sig", newline="") as handle:
+        writer = csv.writer(handle)
+        writer.writerow(["name", "class", "source"])
+        writer.writerows(sorted(name_rows, key=lambda row: (row[1], row[0])))
+    (out_dir / "summary.json").write_text(
+        json.dumps(
+            {
+                "source": str(source.resolve()),
+                "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},
+                "renamed": renamed is not None,
+                "label_studio": label_studio,
+                "split": False,
+                "per_class": dict(counts),
+                "link_modes": dict(modes),
+            },
+            ensure_ascii=False,
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+    logger.info("图片落盘方式: %s（hard=硬链接, copy=跨盘复制, convert=转成 PNG）", dict(modes))
+    logger.info("完成 -> %s", out_dir)
+    logger.warning(
+        "数据集未划分 train/val，data.yaml 的 train 与 val 都指向 images/。"
+        "直接拿去训练的话验证指标算的是训练集本身，会虚高；训练前需要另行划分。"
+    )
+    if label_studio:
+        # label-studio-converter 自带的提示漏了环境变量（它在循环里改写了 image_root_url，
+        # 与默认值比较时就不相等了），这里给出带实际路径的完整步骤
+        logger.info(
+            "Label Studio 任务 %d 个，导入步骤：\n"
+            "  1) export LABEL_STUDIO_LOCAL_FILES_SERVING_ENABLED=true\n"
+            "     export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=%s\n"
+            "     uv run label-studio start\n"
+            "  2) 新建项目 -> Labeling Setup -> Code，粘贴 %s\n"
+            "  3) Settings -> Cloud Storage -> Add Source Storage -> Local files，\n"
+            "     绝对路径填 %s/images。只保存，不要点 Sync（会给每张图再建一个空任务）\n"
+            "  4) 导入 %s",
+            label_studio["tasks"], label_studio["document_root"], label_studio["label_config"],
+            label_studio["document_root"], label_studio["tasks_file"],
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())