Files
ModelTrainEval/scripts/prepare_yolo_dataset.py
T
Shaun ZhouandClaude Opus 5 af0897f8c7 refactor(detect): 数据集改为扁平结构、不划分 train/val,Label Studio 文件放到根目录
prepare_yolo_dataset.py 的产出先导入 Label Studio 复核,划分留到复核之后,于是:

- 去掉按类别分层的 train/val 切分以及 --val-ratio / --seed,images/ 与 labels/
  下直接是文件;name_map.csv 去掉 split 列,summary.json 的逐类数量改为单列。
- label_studio/ 子目录的内容上移一层:扁平的数据集本身就是 label-studio-converter
  要求的 <input>/{images,labels,classes.txt},tasks.json 与 labeling config 直接
  生成在根目录,不再需要额外的硬链接视图,DOCUMENT_ROOT 即数据集根目录。
- data.yaml 的 train 与 val 只能指向同一个 images/。ultralytics 能读,但验证指标
  算的是训练集本身,mAP 虚高、早停与 best.pt 的选择失真。脚本结束时告警,
  detect.py 检测到 train == val 时也告警。

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

518 lines
22 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 的类别信息在标签文件的第一列里,不靠目录名。这一步做三件事:
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())