diff --git a/scripts/build_yolo_dataset.py b/scripts/build_yolo_dataset.py new file mode 100644 index 0000000..55253d1 --- /dev/null +++ b/scripts/build_yolo_dataset.py @@ -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` 也一样)。所以图片从我们自己准备 +复核材料的那个目录取——文件名是一一对应的。 + +``` +.zip runs/<...>_label_studio/images/ +├── classes.txt └── <类别>_<序号>.png +├── labels/<类别>_<序号>.txt +└── images/ ← 空的 +``` + +装配结果: + +``` +data// +├── 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()) diff --git a/scripts/prepare_label_studio.py b/scripts/prepare_label_studio.py index f149897..fb8e156 100644 --- a/scripts/prepare_label_studio.py +++ b/scripts/prepare_label_studio.py @@ -55,7 +55,6 @@ from __future__ import annotations import argparse import contextlib import csv -import errno import io import json import logging @@ -73,7 +72,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import yaml # 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 @@ -116,21 +115,6 @@ def materialize_png(source: Path, destination: Path) -> str: 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): """一张图及其标签。 diff --git a/tests/test_build_yolo_dataset.py b/tests/test_build_yolo_dataset.py new file mode 100644 index 0000000..f3a9eda --- /dev/null +++ b/tests/test_build_yolo_dataset.py @@ -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 diff --git a/utils/__init__.py b/utils/__init__.py index d7cad96..8c11c35 100644 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -1,4 +1,4 @@ -"""通用工具:配置、注册表、日志、随机种子、设备检查。""" +"""通用工具:配置、注册表、日志、随机种子、设备检查、文件落盘。""" from utils.config import Config, load_config from utils.device import ( @@ -10,6 +10,7 @@ from utils.device import ( resolve_amp, resolve_device, ) +from utils.fs import link_or_copy from utils.logging import get_logger, setup_logging from utils.registry import Registry from utils.seed import seed_everything @@ -28,4 +29,5 @@ __all__ = [ "setup_logging", "Registry", "seed_everything", + "link_or_copy", ] diff --git a/utils/fs.py b/utils/fs.py new file mode 100644 index 0000000..002aee9 --- /dev/null +++ b/utils/fs.py @@ -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"