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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
1a43267b55
commit
af0897f8c7
+78
-110
@@ -1,5 +1,5 @@
|
||||
#!/usr/bin/env python
|
||||
"""把 generate_boxes.py 的产出整理成 YOLO 训练所需的数据集，放到 data/ 下。
|
||||
"""把 generate_boxes.py 的产出整理成扁平的 YOLO 数据集，放到 data/ 下。
|
||||
|
||||
uv run python scripts/prepare_yolo_dataset.py --source runs/boxes_gradcam
|
||||
|
||||
@@ -11,33 +11,33 @@ runs/boxes_gradcam/
|
||||
└── labels/<类别>/<图名>.txt
|
||||
```
|
||||
|
||||
而 YOLO 检测训练要的是 **train/val 划分**——类别信息在标签文件的第一列里，
|
||||
不靠目录名。这一步做三件事：
|
||||
而 YOLO 的类别信息在标签文件的第一列里，不靠目录名。这一步做三件事：
|
||||
|
||||
1. 把按类别分的目录拍平，按 train/val 重新组织；
|
||||
2. **按类别分层**切分，小类不至于在验证集里挂零；
|
||||
3. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式
|
||||
1. 把按类别分的目录**拍平**，不划分 train/val——这份数据集先导入 Label Studio 复核，
|
||||
划分留到复核之后；
|
||||
2. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式
|
||||
（pearlescent / blue / uv 原始数据都是 BMP）无损转成 PNG。不缩放像素，
|
||||
归一化的框坐标原样有效。
|
||||
|
||||
4. 文件**按类别重新编号**：`<类别>_<序号>.png`，序号在类内按源路径排序，重跑结果一致。
|
||||
3. 文件**按类别重新编号**：`<类别>_<序号>.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
|
||||
├── images/<类别>_<序号>.png
|
||||
├── labels/<类别>_<序号>.txt
|
||||
├── classes.txt # 每行一个类别名，行号即 class_id
|
||||
├── name_map.csv # 新名 ↔ 源文件对照
|
||||
├── data.yaml # train 与 val 都指向 images/，不能直接用于评估
|
||||
├── summary.json
|
||||
└── label_studio/ # 仅 --label-studio：导入 Label Studio 人工复核用
|
||||
├── images/*.png labels/*.txt classes.txt # 两个 split 拍平，硬链接
|
||||
├── tasks.json # 框作为 predictions
|
||||
└── tasks.label_config.xml
|
||||
├── 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` 按来源类别分别处置：
|
||||
@@ -56,7 +56,6 @@ import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import sys
|
||||
from collections import Counter, defaultdict
|
||||
@@ -196,8 +195,8 @@ def class_of(sample: Sample, names: Sequence[str]) -> str:
|
||||
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:
|
||||
@@ -210,16 +209,12 @@ def assign_names(samples: Sequence[Sample], names: Sequence[str]) -> dict[Path,
|
||||
return mapping
|
||||
|
||||
|
||||
LABEL_STUDIO_DIR = "label_studio"
|
||||
def build_label_studio_tasks(out_dir: Path, num_images: int) -> dict:
|
||||
"""在数据集根目录上调用 label-studio-converter，生成 tasks.json 与 labeling config。
|
||||
|
||||
|
||||
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 之下，软链接解析后可能落到外面。
|
||||
converter 的 YOLO 导入只认 `<input>/{images,labels,classes.txt}`，且 os.listdir 不递归——
|
||||
这正是本脚本产出的扁平结构。图片是硬链接而不是软链接也有讲究：Label Studio 本地存储
|
||||
会检查文件是否在 DOCUMENT_ROOT 之下，软链接解析后可能落到外面。
|
||||
|
||||
框以 predictions 导入：它们是 SHAP / Grad-CAM 生成的，不是人工标注。审核员修改提交
|
||||
之后才成为 annotation，Label Studio 导出时也只导出 annotation，没审过的框不会混进去。
|
||||
@@ -229,14 +224,7 @@ def build_label_studio_view(out_dir: Path, num_images: int) -> dict:
|
||||
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"
|
||||
tasks_file = out_dir / "tasks.json"
|
||||
# converter 会往 stdout 打一份缺了环境变量的导入说明，还会刷一串 root logger 的 INFO；
|
||||
# 静音它们，由 main() 打印带实际路径的完整步骤
|
||||
root_logger = logging.getLogger()
|
||||
@@ -246,7 +234,7 @@ def build_label_studio_view(out_dir: Path, num_images: int) -> dict:
|
||||
with contextlib.redirect_stdout(io.StringIO()):
|
||||
# 不传 image_dims：图片尺寸不统一，传了会把同一个宽高写进每一条标注
|
||||
convert_yolo_to_ls(
|
||||
input_dir=str(view),
|
||||
input_dir=str(out_dir),
|
||||
out_file=str(tasks_file),
|
||||
out_type="predictions",
|
||||
image_ext=".png",
|
||||
@@ -257,13 +245,13 @@ def build_label_studio_view(out_dir: Path, num_images: int) -> dict:
|
||||
if len(tasks) != num_images:
|
||||
raise RuntimeError(
|
||||
f"label-studio-converter 只生成了 {len(tasks)} 个任务，应为 {num_images}；"
|
||||
f"检查 {view / 'images'}"
|
||||
f"检查 {out_dir / 'images'}"
|
||||
)
|
||||
return {
|
||||
"tasks": len(tasks),
|
||||
"document_root": str(view.resolve()),
|
||||
"document_root": str(out_dir.resolve()),
|
||||
"tasks_file": str(tasks_file.resolve()),
|
||||
"label_config": str((view / "tasks.label_config.xml").resolve()),
|
||||
"label_config": str((out_dir / "tasks.label_config.xml").resolve()),
|
||||
}
|
||||
|
||||
|
||||
@@ -337,8 +325,6 @@ def main() -> int:
|
||||
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=(
|
||||
@@ -354,7 +340,7 @@ def main() -> int:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--label-studio", action="store_true",
|
||||
help="额外生成 <out>/label_studio/：拍平的硬链接视图 + tasks.json，可直接导入 Label Studio 复核",
|
||||
help="额外在数据集根目录生成 tasks.json 与 labeling config，可直接导入 Label Studio 复核",
|
||||
)
|
||||
parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
|
||||
parser.add_argument("--dry-run", action="store_true", help="只打印计划")
|
||||
@@ -403,35 +389,16 @@ def main() -> int:
|
||||
"/".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")
|
||||
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["train"][name] or counts["val"][name]:
|
||||
logger.info("%-14s %8d %8d", name, counts["train"][name], counts["val"][name])
|
||||
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:
|
||||
@@ -442,7 +409,7 @@ def main() -> int:
|
||||
|
||||
if args.dry_run:
|
||||
if args.label_studio:
|
||||
logger.info("--label-studio：会在 %s 下生成导入视图", out_dir / LABEL_STUDIO_DIR)
|
||||
logger.info("--label-studio：会在 %s 下生成 tasks.json", out_dir)
|
||||
logger.info("--dry-run：未写入任何文件")
|
||||
return 0
|
||||
|
||||
@@ -456,52 +423,51 @@ def main() -> int:
|
||||
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")
|
||||
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")
|
||||
|
||||
write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(),
|
||||
train="images/train", val="images/val")
|
||||
# 不划分时 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_view(out_dir, len(samples))
|
||||
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(["split", "name", "class", "source"])
|
||||
writer.writerows(sorted(name_rows, key=lambda row: (row[2], row[1])))
|
||||
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(
|
||||
{
|
||||
@@ -514,10 +480,8 @@ def main() -> int:
|
||||
"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()},
|
||||
"split": False,
|
||||
"per_class": dict(counts),
|
||||
"link_modes": dict(modes),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
@@ -527,6 +491,10 @@ def main() -> int:
|
||||
)
|
||||
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，
|
||||
# 与默认值比较时就不相等了），这里给出带实际路径的完整步骤
|
||||
|
||||
Reference in New Issue
Block a user