From 7adc5b6332e0fecfe91e6aef087b085b13f8cd66 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Tue, 15 Sep 2026 16:45:58 +0800 Subject: [PATCH] =?UTF-8?q?feat(detect):=20YOLO=20=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E9=9B=86=E7=BB=9F=E4=B8=80=E4=B8=BA=20PNG=E3=80=81=E6=8C=89?= =?UTF-8?q?=E7=B1=BB=E5=88=AB=E9=87=8D=E5=91=BD=E5=90=8D=E5=B9=B6=E9=99=84?= =?UTF-8?q?=20classes.txt?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 图片统一成 PNG。pearlescent / blue / uv 的原始数据是 BMP，generate_boxes.py 按原后缀链接，于是 runs/boxes_shap 里 3685 张 BMP 与 7526 张 PNG 混在一起， ripple 与 uneven 两类全是 BMP。按文件头（而非后缀）判断：真 PNG 照旧硬链接， 其余无损转换，不缩放像素，归一化框坐标原样有效；CMYK 等 PNG 存不下的模式转 RGB。 在真实数据上核对过 3685 张转换结果逐像素一致。 - 文件按 <类别>_<序号> 重新编号。类别取标签第一列，文件名不会与标签矛盾； 序号在类内按源路径排序、与划分无关，换种子重跑名字不变。新旧名对照写入 name_map.csv，源文件名里的正反面、位置、时间戳不丢。--keep-names 保留原名。 - 输出 classes.txt（行号即 class_id），供 LabelImg / X-AnyLabeling 复核框时使用。 - 撞名检查改为按输出文件名的小写比较：统一后缀后 a.png 与 a.bmp 会撞， macOS 默认文件系统也不区分大小写。 Co-Authored-By: Claude Opus 5 --- README.md | 23 +++- datasets/yolo.py | 15 +++ scripts/prepare_yolo_dataset.py | 129 +++++++++++++++--- tests/test_prepare_yolo_dataset.py | 208 +++++++++++++++++++++++++++-- 4 files changed, 346 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index ea68a9c..0322d01 100644 --- a/README.md +++ b/README.md @@ -610,12 +610,29 @@ uv run python detect.py --config configs/detect/film.yaml --resume ``` data/boxes_gradcam/ -├── images/{train,val}/<图名>.png # 硬链接到 runs/，再到 data/raw/ 的原图 -├── labels/{train,val}/<图名>.txt +├── images/{train,val}/<类别>_<序号>.png # 一律 PNG：PNG 原图硬链接，BMP 等无损转换 +├── labels/{train,val}/<类别>_<序号>.txt +├── classes.txt # 每行一个类别名，行号即 class_id +├── name_map.csv # 新文件名 ↔ 源文件路径 ├── data.yaml # path 写绝对路径，ultralytics 从哪跑都能读到 -└── summary.json # 划分明细、逐类数量、链接方式 +└── summary.json # 划分明细、逐类数量、落盘方式 ``` +**图片统一成 PNG。** pearlescent / blue / uv 的原始数据是 BMP（共 10,814 张）， +`generate_boxes.py` 按原后缀链接，于是 YOLO 数据集里会混着两种格式。这一步按**文件头** +（不是后缀）判断：真 PNG 照旧硬链接不占磁盘，其余无损转 PNG——不缩放像素，归一化的框坐标 +原样有效；CMYK 这类 PNG 存不下的模式转 RGB。转换无法硬链接，会占额外磁盘，`--dry-run` +会先报出要转多少张。统一后缀后 `a.png` 与另一目录的 `a.bmp` 会撞名，脚本直接报错。 + +**按类别重新编号。** 源文件名五花八门（`反面纵黑点_1_550mmX1561.71m_20260416160441_0.png`、 +`pearlescent_404.bmp`……），输出统一改成 `<类别>_<序号>.png`，如 `blackdot_0042.png`。 +类别以标签第一列为准，文件名永远不会与标签矛盾；序号在类内按源路径排序、与 train/val 划分 +无关，换 `--seed` 或 `--val-ratio` 重跑，同一张图拿到的名字不变。原文件名里的正反面、位置、 +时间戳等信息都在 `name_map.csv` 里（带 BOM，Excel 直接打开中文不乱码）。想保留原名加 `--keep-names`。 + +**classes.txt** 是给 LabelImg / X-AnyLabeling 这类标注工具用的，人工复核、修改框时 +靠它把 class_id 映射回类别名；ultralytics 训练本身只读 `data.yaml`，两者内容严格一致。 + ### 划分 按类别分层（`--val-ratio` 默认 0.2，`--seed` 默认 42，两次运行结果一致）。分层是必需的： diff --git a/datasets/yolo.py b/datasets/yolo.py index 6a06e95..33de220 100644 --- a/datasets/yolo.py +++ b/datasets/yolo.py @@ -18,6 +18,21 @@ from pathlib import Path IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp") +def write_classes_txt(path: str | Path, classes: Sequence[str]) -> Path: + """写 classes.txt：每行一个类别名，行号（从 0 起）即标签里的 class_id。 + + ultralytics 训练只读 data.yaml；这份是给 LabelImg / X-AnyLabeling 之类 + 标注工具用的——人工复核、修改框时它们靠它把 class_id 映射回类别名。 + """ + bad = [name for name in classes if not name.strip() or "\n" in name] + if bad: + raise ValueError(f"类别名不能为空或包含换行，否则行号与 class_id 对不上: {bad!r}") + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("".join(f"{name}\n" for name in classes), encoding="utf-8") + return path + + def write_data_yaml( path: str | Path, classes: Sequence[str], diff --git a/scripts/prepare_yolo_dataset.py b/scripts/prepare_yolo_dataset.py index 7faa536..84c2617 100644 --- a/scripts/prepare_yolo_dataset.py +++ b/scripts/prepare_yolo_dataset.py @@ -16,13 +16,22 @@ runs/boxes_gradcam/ 1. 把按类别分的目录拍平，按 train/val 重新组织； 2. **按类别分层**切分，小类不至于在验证集里挂零； -3. 全部用**硬链接**，不占额外磁盘。 +3. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式 + （pearlescent / blue / uv 原始数据都是 BMP）无损转成 PNG。不缩放像素， + 归一化的框坐标原样有效。 + +4. 文件**按类别重新编号**：`<类别>_<序号>.png`，序号在类内按源路径排序，重跑结果一致。 + 原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里，不会丢。 + 想保留原名加 `--keep-names`。 ``` data// -├── images/{train,val}/<图名>.png -├── labels/{train,val}/<图名>.txt -└── data.yaml +├── images/{train,val}/<类别>_<序号>.png +├── labels/{train,val}/<类别>_<序号>.txt +├── classes.txt # 每行一个类别名，行号即 class_id，给标注工具复核用 +├── name_map.csv # 新名 ↔ 源文件对照 +├── data.yaml +└── summary.json ``` 没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。但**空标签的 @@ -36,6 +45,7 @@ data// from __future__ import annotations import argparse +import csv import errno import json import os @@ -51,7 +61,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1])) import yaml # noqa: E402 -from datasets.yolo import IMAGE_SUFFIXES, write_data_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 @@ -64,6 +74,37 @@ EMPTY_POLICY_DESC = { } +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(): @@ -133,6 +174,35 @@ def collect(source: Path) -> list[Sample]: 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 + + def parse_empty_policy(tokens: Sequence[str]) -> tuple[str, dict[str, str]]: """解析 --empty-policy：一个可选的全局默认，加若干 `类别=策略`。 @@ -214,6 +284,10 @@ def main() -> int: "注意空标签在 YOLO 里的意思是『此图无缺陷』，缺陷类的空标签照搬会教模型抑制真实缺陷" ), ) + parser.add_argument( + "--keep-names", 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() @@ -291,6 +365,13 @@ def main() -> int: 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: logger.info("--dry-run：未写入任何文件") return 0 @@ -301,22 +382,33 @@ def main() -> int: 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 - # 拍平之后不同类别目录下的同名文件会撞车，显式挡住 - if image.name in seen: - logger.error( - "拍平后文件名冲突：%s 与 %s 同名。源数据集的文件名需要全局唯一，" - "用 scripts/export_clusters.py --prefix-parent 重新导出即可", - seen[image.name], 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[image.name] = image - modes[link_or_copy(image, out_dir / "images" / split / image.name)] += 1 - destination = out_dir / "labels" / split / f"{image.stem}.txt" + 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: @@ -326,6 +418,12 @@ def main() -> int: 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) + # 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( { @@ -336,6 +434,7 @@ def main() -> int: for group, (policy, count) in sorted(empty_stats.items()) }, "empty_policy": {"default": default_policy, **per_class_policy}, + "renamed": renamed is not None, "val_ratio": args.val_ratio, "seed": args.seed, "split_sizes": {k: len(v) for k, v in splits.items()}, @@ -347,7 +446,7 @@ def main() -> int: ), encoding="utf-8", ) - logger.info("链接方式: %s", dict(modes)) + logger.info("图片落盘方式: %s（hard=硬链接, copy=跨盘复制, convert=转成 PNG）", dict(modes)) logger.info("完成 -> %s", out_dir) return 0 diff --git a/tests/test_prepare_yolo_dataset.py b/tests/test_prepare_yolo_dataset.py index 0a0ac39..f453e67 100644 --- a/tests/test_prepare_yolo_dataset.py +++ b/tests/test_prepare_yolo_dataset.py @@ -1,5 +1,6 @@ """把 generate_boxes.py 的产出整理成 YOLO 训练集。""" +import csv import json import sys from pathlib import Path @@ -44,6 +45,15 @@ def source(tmp_path): return root +def output_of(out, source_stem): + """通过 name_map.csv 找源文件对应的输出图片——重命名后只能这么找。""" + with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle: + for row in csv.DictReader(handle): + if Path(row["source"]).stem == source_stem: + return out / "images" / row["split"] / row["name"] + return None + + def run(source, out, *extra): argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra] original, sys.argv = sys.argv, argv @@ -72,8 +82,8 @@ def test_hard_links_share_inode(source, tmp_path): """用硬链接而不是复制，不占额外磁盘。""" out = tmp_path / "data" / "boxes" assert run(source, out) == 0 - linked = next((out / "images" / "train").iterdir()) - original = next(p for p in (source / "images").rglob("*.png") if p.name == linked.name) + original = source / "images" / "grease" / "grease_0.png" + linked = output_of(out, "grease_0") assert linked.stat().st_ino == original.stat().st_ino @@ -101,8 +111,7 @@ def test_empty_policy_drop_removes_the_image(source, tmp_path): summary = json.loads((out / "summary.json").read_text(encoding="utf-8")) assert summary["num_images"] == 12 assert summary["empty_labels"]["normal"]["policy"] == "drop" - assert not (out / "labels" / "train" / "normal_0.txt").exists() - assert not (out / "images" / "train" / "normal_0.png").exists() + assert output_of(out, "normal_0") is None def test_empty_policy_is_per_class(source, tmp_path): @@ -119,11 +128,8 @@ def test_empty_policy_full_writes_a_whole_image_box(source, tmp_path): out = tmp_path / "data" / "boxes" assert run(source, out, "--empty-policy", "normal=full") == 0 - label = next( - p for split in ("train", "val") - for p in (out / "labels" / split).iterdir() - if p.stem == "normal_0" - ) + image = output_of(out, "normal_0") + label = out / "labels" / image.parent.name / f"{image.stem}.txt" assert label.read_text(encoding="utf-8").split() == ["1", "0.5", "0.5", "1.0", "1.0"] # 重写而不是硬链接，源标签必须还是空的 assert (source / "labels" / "normal" / "normal_0.txt").read_text(encoding="utf-8") == "" @@ -171,12 +177,12 @@ def test_reproducible_split(source, tmp_path): def test_name_collision_after_flattening_is_reported(source, tmp_path, capsys): - """不同类别目录下的同名文件，拍平后会互相覆盖——必须报错而不是静默丢数据。""" + """保留原名时，不同类别目录下的同名文件拍平后会互相覆盖——必须报错而不是静默丢数据。""" duplicate = source / "images" / "normal" / "grease_0.png" Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(duplicate) (source / "labels" / "normal" / "grease_0.txt").write_text("1 0.5 0.5 0.2 0.2\n", encoding="utf-8") - assert run(source, tmp_path / "data" / "boxes") != 0 + assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0 assert "冲突" in capsys.readouterr().out @@ -201,3 +207,183 @@ def test_dry_run_writes_nothing(source, tmp_path): out = tmp_path / "data" / "boxes" assert run(source, out, "--dry-run") == 0 assert not out.exists() + + +def _all_outputs(out, kind): + return [p for split in ("train", "val") for p in (out / kind / split).iterdir()] + + +def _add_sample(source, group, stem, suffix, class_id, *, fmt=None, mode="RGB"): + """往来源目录里加一张指定格式的图 + 一个非空标签。""" + size = (32, 24) + array = np.random.default_rng(1).integers(0, 255, (size[1], size[0], 3), dtype=np.uint8) + image = Image.fromarray(array) + if mode != "RGB": + image = image.convert(mode) + path = source / "images" / group / f"{stem}{suffix}" + image.save(path, format=fmt) + (source / "labels" / group / f"{stem}.txt").write_text( + f"{class_id} 0.5 0.5 0.2 0.2\n", encoding="utf-8" + ) + return path, image + + +def test_all_output_images_are_png(source, tmp_path): + """BMP / JPEG 来源都要落成真正的 PNG——后缀和文件头都得是。""" + _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP") + _add_sample(source, "normal", "jpg_img", ".jpg", 1, fmt="JPEG") + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + + images = _all_outputs(out, "images") + assert len(images) == 15 + assert {p.suffix for p in images} == {".png"} + assert all(pyd.is_png(p) for p in images) + # 标签跟着改名后的图片走，一一对应 + assert sorted(p.stem for p in images) == sorted(p.stem for p in _all_outputs(out, "labels")) + + +def test_bmp_conversion_is_lossless(source, tmp_path): + """转换不能缩放或有损压缩，否则归一化的框就不再对应原来的像素。""" + _, original = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP") + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + + converted = output_of(out, "bmp_img") + with Image.open(converted) as image: + assert image.size == original.size + assert np.array_equal(np.asarray(image), np.asarray(original)) + + +def test_png_sources_are_still_hard_linked(source, tmp_path): + """只有非 PNG 才转换；PNG 照旧硬链接，不占额外磁盘。""" + bmp, _ = _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP") + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + + summary = json.loads((out / "summary.json").read_text(encoding="utf-8")) + assert summary["link_modes"] == {"hard": 13, "convert": 1} + converted = output_of(out, "bmp_img") + assert converted.stat().st_ino != bmp.stat().st_ino + assert bmp.read_bytes()[:2] == b"BM" # 源文件没被动过 + + +def test_png_suffix_with_non_png_content_is_converted(source, tmp_path): + """按文件头判断而不是后缀：写着 .png 其实是 JPEG 的也要转。""" + fake, _ = _add_sample(source, "grease", "fake", ".png", 0, fmt="JPEG") + assert not pyd.is_png(fake) + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + output = output_of(out, "fake") + assert pyd.is_png(output) + + +@pytest.mark.parametrize("mode", ["L", "RGBA", "CMYK"]) +def test_conversion_handles_pixel_modes(source, tmp_path, mode): + """灰度 / 带透明度原样保留；PNG 存不下的 CMYK 转成 RGB。""" + _add_sample(source, "grease", "moded", ".tif", 0, fmt="TIFF", mode=mode) + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + output = output_of(out, "moded") + with Image.open(output) as image: + assert image.format == "PNG" + assert image.mode == ("RGB" if mode == "CMYK" else mode) + + +def test_same_stem_different_suffix_collides(source, tmp_path, capsys): + """grease_0.png 与另一目录的 grease_0.bmp 统一后缀后会撞成同一个文件，必须报错。""" + _add_sample(source, "normal", "grease_0", ".bmp", 1, fmt="BMP") + assert run(source, tmp_path / "data" / "boxes", "--keep-names") != 0 + assert "冲突" in capsys.readouterr().out + + +def test_dry_run_reports_conversion_count(source, tmp_path, capsys): + _add_sample(source, "grease", "bmp_img", ".bmp", 0, fmt="BMP") + assert run(source, tmp_path / "out", "--dry-run") == 0 + assert "1 张不是 PNG" in capsys.readouterr().out + + +def test_classes_txt_matches_data_yaml(source, tmp_path): + """classes.txt 的行号就是 class_id，必须与 data.yaml 的 names 严格一致。""" + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + + lines = (out / "classes.txt").read_text(encoding="utf-8").splitlines() + meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8")) + assert lines == CLASSES + assert lines == [meta["names"][i] for i in range(meta["nc"])] + assert (out / "classes.txt").read_text(encoding="utf-8").endswith("\n") + + +def test_classes_txt_rejects_names_that_break_line_numbering(tmp_path): + from datasets.yolo import write_classes_txt + + with pytest.raises(ValueError, match="换行"): + write_classes_txt(tmp_path / "classes.txt", ["ok", "bad\nname"]) + with pytest.raises(ValueError, match="为空"): + write_classes_txt(tmp_path / "classes.txt", ["ok", " "]) + + +def _stems(out): + return sorted(p.stem for p in _all_outputs(out, "images")) + + +def test_files_renamed_by_class(source, tmp_path): + """默认按 <类别>_<序号> 重新编号，序号在类内从 1 连续编到该类总数。""" + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + assert _stems(out) == ( + [f"grease_{i:04d}" for i in range(1, 7)] + + [f"normal_{i:04d}" for i in range(1, 6)] + + [f"whitedot_{i:04d}" for i in range(1, 3)] + ) + assert _stems(out) == sorted(p.stem for p in _all_outputs(out, "labels")) + + +def test_renamed_file_class_matches_its_label(source, tmp_path): + """文件名里的类别必须与标签第一列一致，否则改名反倒制造了误导。""" + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + for label in _all_outputs(out, "labels"): + text = label.read_text(encoding="utf-8").split() + if text: + assert label.stem.rsplit("_", 1)[0] == CLASSES[int(text[0])] + + +def test_name_map_traces_back_to_source(source, tmp_path): + """name_map.csv 覆盖每一张输出，且源路径真实存在、内容一致。""" + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + with (out / "name_map.csv").open(encoding="utf-8-sig", newline="") as handle: + rows = list(csv.DictReader(handle)) + assert len(rows) == 13 + for row in rows: + output = out / "images" / row["split"] / row["name"] + assert output.read_bytes() == Path(row["source"]).read_bytes() + assert row["name"].startswith(row["class"] + "_") + + +def test_names_stable_across_different_splits(source, tmp_path): + """序号与划分无关：换种子重跑，同一张源图拿到同一个名字。""" + first, second = tmp_path / "a", tmp_path / "b" + assert run(source, first, "--seed", "1") == 0 + assert run(source, second, "--seed", "2") == 0 + for stem in ("grease_0", "grease_5", "normal_3", "whitedot_1"): + assert output_of(first, stem).name == output_of(second, stem).name + + +def test_keep_names_preserves_source_stems(source, tmp_path): + """--keep-names 保留原名，只是后缀统一成 .png。""" + _add_sample(source, "grease", "原名_中文", ".bmp", 0, fmt="BMP") + out = tmp_path / "data" / "boxes" + assert run(source, out, "--keep-names") == 0 + assert "原名_中文" in _stems(out) + assert output_of(out, "原名_中文").name == "原名_中文.png" + + +def test_kept_empty_labels_named_after_source_folder(source, tmp_path): + """空标签没有 class_id，按来源目录名编号；与同类有框的图共用序号，不会撞名。""" + out = tmp_path / "data" / "boxes" + assert run(source, out) == 0 + assert output_of(out, "normal_0").stem.startswith("normal_") + assert len([s for s in _stems(out) if s.startswith("normal_")]) == 5