feat(detect): YOLO 数据集统一为 PNG、按类别重命名并附 classes.txt
- 图片统一成 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 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
cc053a6e9c
commit
7adc5b6332
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user