Files
ModelTrainEval/scripts/prepare_yolo_dataset.py
T
Shaun ZhouandClaude Opus 5 1a43267b55 feat(detect): --label-studio 一步生成可导入 Label Studio 的复核视图
label-studio-converter import yolo 只认 <input>/{images,labels,classes.txt},且
os.listdir 不递归,看不见 train/val 子目录,直接喂数据集会 "No labels converted"。
按类别重命名后文件名全局唯一,于是在 <out>/label_studio/ 下把两个 split 拍平
(硬链接:Label Studio 本地存储会校验路径在 DOCUMENT_ROOT 之下,软链接可能越界),
再调用 converter 生成 tasks.json 与 labeling config。

- 框以 predictions 导入:它们是 SHAP/Grad-CAM 生成的,审核员提交后才成为
  annotation,导出时没审过的框不会混进去。
- 不传 image_dims:真实数据有 3223 种尺寸,传了会写错每条标注的宽高。
- 任务数与图片数不符时报错,而不是静默少导。
- converter 自带的导入提示漏掉了 LABEL_STUDIO_LOCAL_FILES_* 环境变量(它在循环里
  改写了 image_root_url,之后与默认值比较失败),静音后由脚本打印带实际路径的步骤。

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

550 lines
23 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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 检测训练要的是 **train/val 划分**——类别信息在标签文件的第一列里,
不靠目录名。这一步做三件事:
1. 把按类别分的目录拍平,按 train/val 重新组织;
2. **按类别分层**切分,小类不至于在验证集里挂零;
3. 图片**统一成 PNG**:本来就是 PNG 的硬链接过去,不占额外磁盘;BMP 等其他格式
(pearlescent / blue / uv 原始数据都是 BMP)无损转成 PNG。不缩放像素,
归一化的框坐标原样有效。
4. 文件**按类别重新编号**:`<类别>_<序号>.png`,序号在类内按源路径排序,重跑结果一致。
原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里,不会丢。
想保留原名加 `--keep-names`。
```
data/<name>/
├── images/{train,val}/<类别>_<序号>.png
├── labels/{train,val}/<类别>_<序号>.txt
├── classes.txt # 每行一个类别名,行号即 class_id,给标注工具复核用
├── name_map.csv # 新名 ↔ 源文件对照
├── data.yaml
├── summary.json
└── label_studio/ # 仅 --label-studio:导入 Label Studio 人工复核用
├── images/*.png labels/*.txt classes.txt # 两个 split 拍平,硬链接
├── tasks.json # 框作为 predictions
└── tasks.label_config.xml
```
没有框的图(标签为空)默认保留: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 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]))
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`。
序号在类内按源路径排序、与 train/val 划分无关:换 --val-ratio 或 --seed 重跑,
同一张图拿到的名字不变,已经做过的人工复核记录还能对得上。
"""
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
LABEL_STUDIO_DIR = "label_studio"
def build_label_studio_view(out_dir: Path, num_images: int) -> dict:
"""在 <out>/label_studio/ 下搭一份 label-studio-converter 能直接吃的扁平视图并转换。
label-studio-converter 的 YOLO 导入只认 `<input>/{images,labels,classes.txt}`,
且 os.listdir 不递归——train/val 子目录它看不见。重命名后文件名全局唯一,
两个 split 可以直接拍平;用硬链接而不是软链接,因为 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
view = out_dir / LABEL_STUDIO_DIR
for kind in ("images", "labels"):
for split in ("train", "val"):
for path in sorted((out_dir / kind / split).iterdir()):
link_or_copy(path, view / kind / path.name)
link_or_copy(out_dir / "classes.txt", view / "classes.txt")
tasks_file = view / "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(view),
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"检查 {view / 'images'}"
)
return {
"tasks": len(tasks),
"document_root": str(view.resolve()),
"tasks_file": str(tasks_file.resolve()),
"label_config": str((view / "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("--val-ratio", type=float, default=0.2)
parser.add_argument("--seed", type=int, default=42)
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="额外生成 <out>/label_studio/:拍平的硬链接视图 + tasks.json,可直接导入 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),
)
# 按类别分层切分;无框的图单独作为一组分配,避免它们全落进同一侧
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": []}
for _class_id, group in sorted(grouped.items(), key=lambda kv: (kv[0] is None, kv[0])):
order = group[:]
rng.shuffle(order)
n_val = max(1, round(len(order) * args.val_ratio)) if len(order) > 1 else 0
splits["val"].extend(order[:n_val])
splits["train"].extend(order[n_val:])
counts = {
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"]))
logger.info("%-14s %8s %8s", "类别", "train", "val")
for name in names + ["(无框)"]:
if counts["train"][name] or counts["val"][name]:
logger.info("%-14s %8d %8d", name, counts["train"][name], counts["val"][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 下生成导入视图", out_dir / LABEL_STUDIO_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 split, items in splits.items():
for sample in items:
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" / split / output_name)] += 1
name_rows.append((split, output_name, class_of(sample, names), str(image.resolve())))
destination = out_dir / "labels" / split / 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")
write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(),
train="images/train", val="images/val")
write_classes_txt(out_dir / "classes.txt", names)
label_studio = None
if args.label_studio:
try:
label_studio = build_label_studio_view(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(["split", "name", "class", "source"])
writer.writerows(sorted(name_rows, key=lambda row: (row[2], row[1])))
(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,
"val_ratio": args.val_ratio,
"seed": args.seed,
"split_sizes": {k: len(v) for k, v in splits.items()},
"per_class": {k: dict(v) for k, v in counts.items()},
"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)
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())