feat(detect): 把 Label Studio 的 YOLO 导出装配成带 train/val 划分的数据集
scripts/build_yolo_dataset.py:导出(zip 或目录)+ 复核材料里的 images/ -> data/ 下 可训练的数据集。Label Studio 的 YOLO 导出只有标签,本地存储的图片不会打包进去 (get_local_path 对 local-files 直接返回磁盘路径,YOLO_WITH_IMAGES 也一样), 所以图片从准备复核材料的目录按文件名取,硬链接过去。 - 类别顺序照抄导出的 classes.txt:转换器把标签收进 set 再排序编号,与 labeling config 的顺序无关,标签里的 class_id 是按它写的。 - 按类别分层切分,分层用"框最多的那个类别"——人工复核后一张图可能有多个类别的框。 - 标签逐行校验:字段数、类别下标、坐标范围。坏标签训练时不会崩,只会悄悄拉低指标。 - 空标签默认保留为负样本;标签找不到图直接报错,图多于标签则告警。 顺带把 link_or_copy 提到 utils/fs.py 共用(此前 generate_boxes 与 prepare_label_studio 各有一份)。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
65bc6b2045
commit
998e6601d4
@@ -0,0 +1,295 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
"""把 Label Studio 的 YOLO 导出装配成可训练的数据集,划分 train/val 放到 data/ 下。
|
||||||
|
|
||||||
|
uv run python scripts/build_yolo_dataset.py \\
|
||||||
|
--export data/film_yolo_shap.zip \\
|
||||||
|
--images runs/boxes_shap_label_studio/images \\
|
||||||
|
--out data/film_yolo
|
||||||
|
|
||||||
|
Label Studio 的 YOLO 导出**只有标签**:本地存储的图片它不会打包进去(`get_local_path` 对
|
||||||
|
local-files 直接返回磁盘路径,不复制,`YOLO_WITH_IMAGES` 也一样)。所以图片从我们自己准备
|
||||||
|
复核材料的那个目录取——文件名是一一对应的。
|
||||||
|
|
||||||
|
```
|
||||||
|
<export>.zip runs/<...>_label_studio/images/
|
||||||
|
├── classes.txt └── <类别>_<序号>.png
|
||||||
|
├── labels/<类别>_<序号>.txt
|
||||||
|
└── images/ ← 空的
|
||||||
|
```
|
||||||
|
|
||||||
|
装配结果:
|
||||||
|
|
||||||
|
```
|
||||||
|
data/<name>/
|
||||||
|
├── images/{train,val}/<类别>_<序号>.png # 硬链接,不占额外磁盘
|
||||||
|
├── labels/{train,val}/<类别>_<序号>.txt
|
||||||
|
├── classes.txt # 原样沿用导出的顺序
|
||||||
|
├── data.yaml
|
||||||
|
└── summary.json
|
||||||
|
```
|
||||||
|
|
||||||
|
**类别顺序以导出的 classes.txt 为准。** Label Studio 的转换器把标签收进 set 再排序编号,
|
||||||
|
与 labeling config 里写的顺序无关,也与我们 `generate_boxes.py` 时的顺序无关。标签里的
|
||||||
|
class_id 是按导出的顺序写的,所以 data.yaml 必须照抄,不能另排。
|
||||||
|
|
||||||
|
**空标签保留为负样本。** 导出后仍为空的图(本例是 normal,预测已在 Label Studio 里批量删掉)
|
||||||
|
表示"这张图没有缺陷",在检测训练里是压低误检的负样本。`--drop-empty` 可以丢掉它们。
|
||||||
|
注意:没复核、也没把预测转成标注的任务同样会导成空标签,那种"空"是信息缺失而不是负样本,
|
||||||
|
装配前应先在 Label Studio 里用 Create Annotations From Predictions 补上。
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import random
|
||||||
|
import shutil
|
||||||
|
import sys
|
||||||
|
import zipfile
|
||||||
|
from collections import Counter, defaultdict
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import NamedTuple
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt, write_data_yaml # noqa: E402
|
||||||
|
from utils import link_or_copy, setup_logging # noqa: E402
|
||||||
|
|
||||||
|
logger = None
|
||||||
|
|
||||||
|
EMPTY = "(无框)"
|
||||||
|
|
||||||
|
|
||||||
|
class Sample(NamedTuple):
|
||||||
|
"""一张图、它的标签文本,以及分层划分用的类别。"""
|
||||||
|
|
||||||
|
image: Path
|
||||||
|
label_text: str
|
||||||
|
stratum: str
|
||||||
|
|
||||||
|
|
||||||
|
class ExportError(RuntimeError):
|
||||||
|
"""导出内容不自洽,装配无法继续。"""
|
||||||
|
|
||||||
|
|
||||||
|
def read_export(path: Path) -> tuple[list[str], dict[str, str]]:
|
||||||
|
"""读 Label Studio 的 YOLO 导出(zip 或已解压目录)。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(类别名列表, {文件名主干: 标签文本})
|
||||||
|
"""
|
||||||
|
if path.is_dir():
|
||||||
|
classes_file = path / "classes.txt"
|
||||||
|
if not classes_file.exists():
|
||||||
|
raise ExportError(f"{path} 下没有 classes.txt,这不像 Label Studio 的 YOLO 导出")
|
||||||
|
classes = classes_file.read_text(encoding="utf-8").split()
|
||||||
|
labels = {
|
||||||
|
p.stem: p.read_text(encoding="utf-8")
|
||||||
|
for p in sorted((path / "labels").glob("*.txt"))
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
with zipfile.ZipFile(path) as archive:
|
||||||
|
names = archive.namelist()
|
||||||
|
if "classes.txt" not in names:
|
||||||
|
raise ExportError(f"{path} 里没有 classes.txt,这不像 Label Studio 的 YOLO 导出")
|
||||||
|
classes = archive.read("classes.txt").decode("utf-8").split()
|
||||||
|
labels = {
|
||||||
|
Path(name).stem: archive.read(name).decode("utf-8")
|
||||||
|
for name in sorted(names)
|
||||||
|
if name.startswith("labels/") and name.endswith(".txt")
|
||||||
|
}
|
||||||
|
if not classes:
|
||||||
|
raise ExportError("classes.txt 是空的")
|
||||||
|
if not labels:
|
||||||
|
raise ExportError("导出里没有任何标签文件")
|
||||||
|
return classes, labels
|
||||||
|
|
||||||
|
|
||||||
|
def parse_boxes(text: str, stem: str, num_classes: int) -> list[int]:
|
||||||
|
"""校验一个标签文件,返回其中所有框的类别下标。
|
||||||
|
|
||||||
|
坐标越界或类别越界都直接报错:这类问题在训练时不会崩,只会悄悄拉低指标。
|
||||||
|
"""
|
||||||
|
class_ids = []
|
||||||
|
for number, line in enumerate(text.splitlines(), start=1):
|
||||||
|
line = line.strip()
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
parts = line.split()
|
||||||
|
if len(parts) != 5:
|
||||||
|
raise ExportError(f"{stem}.txt 第 {number} 行有 {len(parts)} 个字段,YOLO 需要 5 个:{line!r}")
|
||||||
|
try:
|
||||||
|
class_id = int(parts[0])
|
||||||
|
coordinates = [float(v) for v in parts[1:]]
|
||||||
|
except ValueError as exc:
|
||||||
|
raise ExportError(f"{stem}.txt 第 {number} 行无法解析:{line!r}") from exc
|
||||||
|
if not 0 <= class_id < num_classes:
|
||||||
|
raise ExportError(f"{stem}.txt 第 {number} 行的类别下标 {class_id} 超出 0~{num_classes - 1}")
|
||||||
|
if not all(0.0 <= v <= 1.0 for v in coordinates):
|
||||||
|
raise ExportError(f"{stem}.txt 第 {number} 行的坐标不在 [0,1] 内:{coordinates}")
|
||||||
|
class_ids.append(class_id)
|
||||||
|
return class_ids
|
||||||
|
|
||||||
|
|
||||||
|
def collect(
|
||||||
|
classes: list[str], labels: dict[str, str], images_dir: Path
|
||||||
|
) -> tuple[list[Sample], list[str]]:
|
||||||
|
"""把标签与图片配对,返回 (样本列表, 多余的图片名)。"""
|
||||||
|
samples = []
|
||||||
|
missing = []
|
||||||
|
for stem, text in labels.items():
|
||||||
|
candidates = [
|
||||||
|
images_dir / f"{stem}{suffix}" for suffix in (".png", *IMAGE_SUFFIXES)
|
||||||
|
]
|
||||||
|
image = next((p for p in candidates if p.exists()), None)
|
||||||
|
if image is None:
|
||||||
|
missing.append(stem)
|
||||||
|
continue
|
||||||
|
class_ids = parse_boxes(text, stem, len(classes))
|
||||||
|
# 分层用"框最多的那个类别":一张图可能有多个类别的框,人工复核后尤其常见
|
||||||
|
stratum = classes[Counter(class_ids).most_common(1)[0][0]] if class_ids else EMPTY
|
||||||
|
samples.append(Sample(image, text, stratum))
|
||||||
|
|
||||||
|
if missing:
|
||||||
|
raise ExportError(
|
||||||
|
f"有 {len(missing)} 个标签找不到对应的图片,例如 {missing[0]};"
|
||||||
|
f"--images 是否指向了准备复核材料时的 images/ 目录(当前 {images_dir})"
|
||||||
|
)
|
||||||
|
known = {s.image.stem for s in samples}
|
||||||
|
extra = sorted(
|
||||||
|
p.stem for p in images_dir.iterdir()
|
||||||
|
if p.suffix.lower() in IMAGE_SUFFIXES and p.stem not in known
|
||||||
|
)
|
||||||
|
return samples, extra
|
||||||
|
|
||||||
|
|
||||||
|
def split(samples: list[Sample], val_ratio: float, seed: int) -> dict[str, list[Sample]]:
|
||||||
|
"""按类别分层切分。小类只有几十张,随机整体切很容易让它在验证集里挂零。"""
|
||||||
|
grouped: dict[str, list[Sample]] = defaultdict(list)
|
||||||
|
for sample in samples:
|
||||||
|
grouped[sample.stratum].append(sample)
|
||||||
|
|
||||||
|
rng = random.Random(seed)
|
||||||
|
splits: dict[str, list[Sample]] = {"train": [], "val": []}
|
||||||
|
for _stratum, group in sorted(grouped.items()):
|
||||||
|
order = sorted(group, key=lambda s: s.image.name) # 先定序,结果才可复现
|
||||||
|
rng.shuffle(order)
|
||||||
|
n_val = max(1, round(len(order) * val_ratio)) if len(order) > 1 else 0
|
||||||
|
splits["val"].extend(order[:n_val])
|
||||||
|
splits["train"].extend(order[n_val:])
|
||||||
|
return splits
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
global logger
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||||
|
)
|
||||||
|
parser.add_argument("--export", required=True, help="Label Studio 的 YOLO 导出(zip 或目录)")
|
||||||
|
parser.add_argument(
|
||||||
|
"--images", required=True, help="图片目录,即准备复核材料时的 <...>_label_studio/images",
|
||||||
|
)
|
||||||
|
parser.add_argument("--out", required=True, help="输出目录,建议放在 data/ 下")
|
||||||
|
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="丢掉空标签的图。默认保留——它们是压低误检的负样本",
|
||||||
|
)
|
||||||
|
parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
|
||||||
|
parser.add_argument("--dry-run", action="store_true", help="只打印计划")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
out_dir = Path(args.out)
|
||||||
|
logger = setup_logging()
|
||||||
|
|
||||||
|
try:
|
||||||
|
classes, labels = read_export(Path(args.export))
|
||||||
|
samples, extra = collect(classes, labels, Path(args.images))
|
||||||
|
except (ExportError, FileNotFoundError) as exc:
|
||||||
|
logger.error("%s", exc)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
logger.info("导出: %s(%d 个标签,%d 类)", args.export, len(labels), len(classes))
|
||||||
|
if extra:
|
||||||
|
logger.warning(
|
||||||
|
"%s 里有 %d 张图片没有对应的标签,不会进入数据集(例如 %s)",
|
||||||
|
args.images, len(extra), extra[0],
|
||||||
|
)
|
||||||
|
|
||||||
|
empty = [s for s in samples if s.stratum == EMPTY]
|
||||||
|
if args.drop_empty:
|
||||||
|
samples = [s for s in samples if s.stratum != EMPTY]
|
||||||
|
logger.info(
|
||||||
|
"空标签 %d 张(%.1f%%):%s",
|
||||||
|
len(empty), 100 * len(empty) / max(len(labels), 1),
|
||||||
|
"已丢弃" if args.drop_empty else "保留为负样本",
|
||||||
|
)
|
||||||
|
|
||||||
|
splits = split(samples, args.val_ratio, args.seed)
|
||||||
|
counts = {name: Counter(s.stratum for s in items) for name, items in splits.items()}
|
||||||
|
boxes = Counter()
|
||||||
|
for sample in samples:
|
||||||
|
for line in sample.label_text.splitlines():
|
||||||
|
if line.strip():
|
||||||
|
boxes[classes[int(line.split()[0])]] += 1
|
||||||
|
|
||||||
|
logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"]))
|
||||||
|
logger.info("%-12s %8s %8s %8s", "类别", "train", "val", "框数")
|
||||||
|
for name in [*classes, EMPTY]:
|
||||||
|
if counts["train"][name] or counts["val"][name]:
|
||||||
|
logger.info(
|
||||||
|
"%-12s %8d %8d %8d",
|
||||||
|
name, counts["train"][name], counts["val"][name], boxes[name],
|
||||||
|
)
|
||||||
|
|
||||||
|
if args.dry_run:
|
||||||
|
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)
|
||||||
|
|
||||||
|
modes: Counter[str] = Counter()
|
||||||
|
for name, items in splits.items():
|
||||||
|
for sample in items:
|
||||||
|
modes[link_or_copy(sample.image, out_dir / "images" / name / sample.image.name)] += 1
|
||||||
|
label_path = out_dir / "labels" / name / f"{sample.image.stem}.txt"
|
||||||
|
label_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
label_path.write_text(sample.label_text, encoding="utf-8")
|
||||||
|
|
||||||
|
write_data_yaml(out_dir / "data.yaml", classes, root=out_dir.resolve(),
|
||||||
|
train="images/train", val="images/val")
|
||||||
|
write_classes_txt(out_dir / "classes.txt", classes)
|
||||||
|
(out_dir / "summary.json").write_text(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"export": str(Path(args.export).resolve()),
|
||||||
|
"images": str(Path(args.images).resolve()),
|
||||||
|
"num_images": len(samples),
|
||||||
|
"num_boxes": sum(boxes.values()),
|
||||||
|
"empty_labels": len(empty),
|
||||||
|
"dropped_empty": args.drop_empty,
|
||||||
|
"images_without_label": len(extra),
|
||||||
|
"val_ratio": args.val_ratio,
|
||||||
|
"seed": args.seed,
|
||||||
|
"split_sizes": {k: len(v) for k, v in splits.items()},
|
||||||
|
"per_class_images": {k: dict(v) for k, v in counts.items()},
|
||||||
|
"per_class_boxes": dict(boxes),
|
||||||
|
"link_modes": dict(modes),
|
||||||
|
},
|
||||||
|
ensure_ascii=False,
|
||||||
|
indent=2,
|
||||||
|
),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
logger.info("图片落盘方式: %s", dict(modes))
|
||||||
|
logger.info("完成 -> %s(%d 个框)", out_dir, sum(boxes.values()))
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
@@ -55,7 +55,6 @@ from __future__ import annotations
|
|||||||
import argparse
|
import argparse
|
||||||
import contextlib
|
import contextlib
|
||||||
import csv
|
import csv
|
||||||
import errno
|
|
||||||
import io
|
import io
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
@@ -73,7 +72,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|||||||
import yaml # noqa: E402
|
import yaml # noqa: E402
|
||||||
|
|
||||||
from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt # noqa: E402
|
from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt # noqa: E402
|
||||||
from utils import setup_logging # noqa: E402
|
from utils import link_or_copy, setup_logging # noqa: E402
|
||||||
|
|
||||||
logger = None
|
logger = None
|
||||||
|
|
||||||
@@ -116,21 +115,6 @@ def materialize_png(source: Path, destination: Path) -> str:
|
|||||||
return "convert"
|
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):
|
class Sample(NamedTuple):
|
||||||
"""一张图及其标签。
|
"""一张图及其标签。
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,229 @@
|
|||||||
|
"""把 Label Studio 的 YOLO 导出装配成带 train/val 划分的数据集。"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import zipfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
pytest.importorskip("PIL")
|
||||||
|
|
||||||
|
import numpy as np # noqa: E402
|
||||||
|
import yaml # noqa: E402
|
||||||
|
from PIL import Image # noqa: E402
|
||||||
|
|
||||||
|
sys.path.insert(0, "scripts")
|
||||||
|
import build_yolo_dataset as byd # noqa: E402
|
||||||
|
|
||||||
|
#: 与 Label Studio 导出一致:类别按字母序,标签里的 class_id 就是这个顺序的下标
|
||||||
|
CLASSES = ["blackdot", "grease", "normal", "wrinkle"]
|
||||||
|
|
||||||
|
#: {文件名主干: 标签内容}。normal 是空标签(负样本),wrinkle_0001 有两个不同类别的框
|
||||||
|
LABELS = {
|
||||||
|
"blackdot_0001": "0 0.5 0.5 0.1 0.1\n",
|
||||||
|
"blackdot_0002": "0 0.2 0.3 0.1 0.1\n",
|
||||||
|
"blackdot_0003": "0 0.6 0.6 0.2 0.2\n",
|
||||||
|
"grease_0001": "1 0.4 0.4 0.3 0.3\n",
|
||||||
|
"grease_0002": "1 0.5 0.5 0.2 0.2\n",
|
||||||
|
"wrinkle_0001": "3 0.5 0.5 0.4 0.4\n0 0.1 0.1 0.05 0.05\n3 0.8 0.8 0.1 0.1\n",
|
||||||
|
"wrinkle_0002": "3 0.3 0.3 0.3 0.3\n",
|
||||||
|
"normal_0001": "",
|
||||||
|
"normal_0002": "",
|
||||||
|
"normal_0003": "",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def export(tmp_path):
|
||||||
|
"""一个 zip 形式的导出 + 与之配套的图片目录。"""
|
||||||
|
images = tmp_path / "label_studio" / "images"
|
||||||
|
images.mkdir(parents=True)
|
||||||
|
rng = np.random.default_rng(0)
|
||||||
|
for stem in LABELS:
|
||||||
|
Image.fromarray(rng.integers(0, 255, (32, 40, 3), dtype=np.uint8)).save(
|
||||||
|
images / f"{stem}.png"
|
||||||
|
)
|
||||||
|
|
||||||
|
archive = tmp_path / "export.zip"
|
||||||
|
with zipfile.ZipFile(archive, "w") as zipped:
|
||||||
|
zipped.writestr("classes.txt", "".join(f"{c}\n" for c in CLASSES))
|
||||||
|
zipped.writestr("notes.json", "{}")
|
||||||
|
for stem, text in LABELS.items():
|
||||||
|
zipped.writestr(f"labels/{stem}.txt", text)
|
||||||
|
return archive, images
|
||||||
|
|
||||||
|
|
||||||
|
def run(export, out, *extra):
|
||||||
|
archive, images = export
|
||||||
|
argv = ["build_yolo_dataset.py", "--export", str(archive), "--images", str(images),
|
||||||
|
"--out", str(out), *extra]
|
||||||
|
original, sys.argv = sys.argv, argv
|
||||||
|
try:
|
||||||
|
return byd.main()
|
||||||
|
finally:
|
||||||
|
sys.argv = original
|
||||||
|
|
||||||
|
|
||||||
|
def outputs(out, kind):
|
||||||
|
return {p.stem: p for split in ("train", "val") for p in (out / kind / split).iterdir()}
|
||||||
|
|
||||||
|
|
||||||
|
def test_builds_a_split_dataset(export, tmp_path):
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
|
||||||
|
images, labels = outputs(out, "images"), outputs(out, "labels")
|
||||||
|
assert set(images) == set(LABELS) == set(labels)
|
||||||
|
for stem, path in labels.items():
|
||||||
|
assert path.read_text(encoding="utf-8") == LABELS[stem]
|
||||||
|
# 图与标签必须落在同一个 split,否则验证集会混进训练图
|
||||||
|
assert path.parent.name == images[stem].parent.name
|
||||||
|
|
||||||
|
|
||||||
|
def test_class_order_comes_from_the_export(export, tmp_path):
|
||||||
|
"""类别顺序必须照抄导出的 classes.txt:标签里的 class_id 是按它写的。"""
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
|
||||||
|
meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
|
||||||
|
assert [meta["names"][i] for i in range(meta["nc"])] == CLASSES
|
||||||
|
assert (out / "classes.txt").read_text(encoding="utf-8").split() == CLASSES
|
||||||
|
assert Path(meta["path"]).is_absolute()
|
||||||
|
assert meta["train"] == "images/train" and meta["val"] == "images/val"
|
||||||
|
|
||||||
|
|
||||||
|
def test_images_are_hard_linked(export, tmp_path):
|
||||||
|
_, images_dir = export
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
|
||||||
|
linked = outputs(out, "images")["grease_0001"]
|
||||||
|
assert linked.stat().st_ino == (images_dir / "grease_0001.png").stat().st_ino
|
||||||
|
assert json.loads((out / "summary.json").read_text(encoding="utf-8"))["link_modes"] == {
|
||||||
|
"hard": len(LABELS)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_small_classes_appear_in_both_splits(export, tmp_path):
|
||||||
|
"""分层切分:三张图的类别也要在验证集里留下一张。"""
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
per_class = json.loads((out / "summary.json").read_text(encoding="utf-8"))["per_class_images"]
|
||||||
|
for name in ("blackdot", "(无框)"):
|
||||||
|
assert per_class["train"][name] >= 1 and per_class["val"][name] == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_multi_class_image_is_stratified_by_its_dominant_class(export, tmp_path):
|
||||||
|
"""一张图可能有多个类别的框(人工复核后常见),按框最多的类别分层。"""
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||||||
|
# wrinkle_0001 有 2 个 wrinkle 框 + 1 个 blackdot 框 -> 计入 wrinkle
|
||||||
|
assert summary["per_class_images"]["train"].get("wrinkle", 0) + summary[
|
||||||
|
"per_class_images"
|
||||||
|
]["val"].get("wrinkle", 0) == 2
|
||||||
|
# 框数按每个框各自的类别统计
|
||||||
|
assert summary["per_class_boxes"] == {"blackdot": 4, "grease": 2, "wrinkle": 3}
|
||||||
|
assert summary["num_boxes"] == 9
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_labels_kept_as_negatives_by_default(export, tmp_path):
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out) == 0
|
||||||
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||||||
|
assert summary["empty_labels"] == 3
|
||||||
|
assert summary["num_images"] == 10
|
||||||
|
assert outputs(out, "labels")["normal_0001"].read_text(encoding="utf-8") == ""
|
||||||
|
|
||||||
|
dropped = tmp_path / "data" / "dropped"
|
||||||
|
assert run(export, dropped, "--drop-empty") == 0
|
||||||
|
assert json.loads((dropped / "summary.json").read_text(encoding="utf-8"))["num_images"] == 7
|
||||||
|
assert "normal_0001" not in outputs(dropped, "images")
|
||||||
|
|
||||||
|
|
||||||
|
def test_reads_an_extracted_directory_too(export, tmp_path):
|
||||||
|
"""zip 与解压后的目录应当等价。"""
|
||||||
|
archive, images = export
|
||||||
|
extracted = tmp_path / "extracted"
|
||||||
|
with zipfile.ZipFile(archive) as zipped:
|
||||||
|
zipped.extractall(extracted)
|
||||||
|
|
||||||
|
from_zip, from_dir = tmp_path / "a", tmp_path / "b"
|
||||||
|
assert run(export, from_zip) == 0
|
||||||
|
assert run((extracted, images), from_dir) == 0
|
||||||
|
assert sorted(outputs(from_zip, "images")) == sorted(outputs(from_dir, "images"))
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_image_is_an_error(export, tmp_path, capsys):
|
||||||
|
"""标签找不到图就报错,而不是悄悄少装配一张。"""
|
||||||
|
_, images = export
|
||||||
|
(images / "grease_0001.png").unlink()
|
||||||
|
assert run(export, tmp_path / "out") != 0
|
||||||
|
assert "找不到对应的图片" in capsys.readouterr().out
|
||||||
|
|
||||||
|
|
||||||
|
def test_extra_images_are_reported(export, tmp_path, capsys):
|
||||||
|
"""图片多于标签时告警——多半是 --images 指错了目录或导出不完整。"""
|
||||||
|
_, images = export
|
||||||
|
Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)).save(images / "stray_0001.png")
|
||||||
|
assert run(export, tmp_path / "out") == 0
|
||||||
|
assert "没有对应的标签" in capsys.readouterr().out
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"text,expected",
|
||||||
|
[
|
||||||
|
("0 0.5 0.5 0.1\n", "需要 5 个"), # 字段数不对
|
||||||
|
("9 0.5 0.5 0.1 0.1\n", "超出"), # 类别下标越界
|
||||||
|
("0 0.5 1.5 0.1 0.1\n", "不在 [0,1]"), # 坐标越界
|
||||||
|
("x 0.5 0.5 0.1 0.1\n", "无法解析"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_malformed_labels_are_rejected(export, tmp_path, capsys, text, expected):
|
||||||
|
"""坏标签在训练时不会崩,只会悄悄拉低指标,所以这里必须拦住。"""
|
||||||
|
archive, images = export
|
||||||
|
broken = tmp_path / "broken.zip"
|
||||||
|
with zipfile.ZipFile(archive) as source, zipfile.ZipFile(broken, "w") as target:
|
||||||
|
for item in source.namelist():
|
||||||
|
data = source.read(item)
|
||||||
|
if item == "labels/grease_0001.txt":
|
||||||
|
data = text.encode("utf-8")
|
||||||
|
target.writestr(item, data)
|
||||||
|
|
||||||
|
assert run((broken, images), tmp_path / "out") != 0
|
||||||
|
assert expected in capsys.readouterr().out
|
||||||
|
|
||||||
|
|
||||||
|
def test_not_a_label_studio_export(tmp_path, capsys):
|
||||||
|
empty = tmp_path / "empty"
|
||||||
|
(empty / "labels").mkdir(parents=True)
|
||||||
|
assert run((empty, tmp_path), tmp_path / "out") != 0
|
||||||
|
assert "classes.txt" in capsys.readouterr().out
|
||||||
|
|
||||||
|
|
||||||
|
def test_split_is_reproducible_and_ratio_respected(export, tmp_path):
|
||||||
|
first, second = tmp_path / "a", tmp_path / "b"
|
||||||
|
assert run(export, first, "--seed", "7") == 0
|
||||||
|
assert run(export, second, "--seed", "7") == 0
|
||||||
|
assert sorted(outputs(first, "images")) == sorted(outputs(second, "images"))
|
||||||
|
for split in ("train", "val"):
|
||||||
|
assert sorted(p.name for p in (first / "images" / split).iterdir()) == sorted(
|
||||||
|
p.name for p in (second / "images" / split).iterdir()
|
||||||
|
)
|
||||||
|
|
||||||
|
# 比例是按每个分层分别取整的:blackdot 3 张 -> 2 张进验证集,所以合计不是对半
|
||||||
|
other = tmp_path / "c"
|
||||||
|
assert run(export, other, "--val-ratio", "0.5") == 0
|
||||||
|
sizes = json.loads((other / "summary.json").read_text(encoding="utf-8"))["split_sizes"]
|
||||||
|
assert sizes["val"] == 6 and sizes["train"] == 4
|
||||||
|
|
||||||
|
|
||||||
|
def test_refuses_to_clobber_and_dry_run_writes_nothing(export, tmp_path):
|
||||||
|
out = tmp_path / "data" / "film"
|
||||||
|
assert run(export, out, "--dry-run") == 0
|
||||||
|
assert not out.exists()
|
||||||
|
|
||||||
|
assert run(export, out) == 0
|
||||||
|
assert run(export, out) != 0 # 已存在且非空
|
||||||
|
assert run(export, out, "--overwrite") == 0
|
||||||
+3
-1
@@ -1,4 +1,4 @@
|
|||||||
"""通用工具:配置、注册表、日志、随机种子、设备检查。"""
|
"""通用工具:配置、注册表、日志、随机种子、设备检查、文件落盘。"""
|
||||||
|
|
||||||
from utils.config import Config, load_config
|
from utils.config import Config, load_config
|
||||||
from utils.device import (
|
from utils.device import (
|
||||||
@@ -10,6 +10,7 @@ from utils.device import (
|
|||||||
resolve_amp,
|
resolve_amp,
|
||||||
resolve_device,
|
resolve_device,
|
||||||
)
|
)
|
||||||
|
from utils.fs import link_or_copy
|
||||||
from utils.logging import get_logger, setup_logging
|
from utils.logging import get_logger, setup_logging
|
||||||
from utils.registry import Registry
|
from utils.registry import Registry
|
||||||
from utils.seed import seed_everything
|
from utils.seed import seed_everything
|
||||||
@@ -28,4 +29,5 @@ __all__ = [
|
|||||||
"setup_logging",
|
"setup_logging",
|
||||||
"Registry",
|
"Registry",
|
||||||
"seed_everything",
|
"seed_everything",
|
||||||
|
"link_or_copy",
|
||||||
]
|
]
|
||||||
|
|||||||
+28
@@ -0,0 +1,28 @@
|
|||||||
|
"""文件落盘的小工具。"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import errno
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def link_or_copy(source: str | Path, destination: str | Path) -> str:
|
||||||
|
"""给文件建硬链接;跨文件系统时退回复制。返回 "hard" 或 "copy"。
|
||||||
|
|
||||||
|
硬链接不占额外磁盘,却表现得和真实文件一样——产出目录可以整个拷走。
|
||||||
|
数据集动辄上万张图,链接与复制的差别是几百 MB 起步。
|
||||||
|
"""
|
||||||
|
destination = Path(destination)
|
||||||
|
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"
|
||||||
Reference in New Issue
Block a user