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:
co-authored by
Claude Opus 5
parent
cc053a6e9c
commit
7adc5b6332
@@ -610,12 +610,29 @@ uv run python detect.py --config configs/detect/film.yaml --resume
|
|||||||
|
|
||||||
```
|
```
|
||||||
data/boxes_gradcam/
|
data/boxes_gradcam/
|
||||||
├── images/{train,val}/<图名>.png # 硬链接到 runs/,再到 data/raw/ 的原图
|
├── images/{train,val}/<类别>_<序号>.png # 一律 PNG:PNG 原图硬链接,BMP 等无损转换
|
||||||
├── labels/{train,val}/<图名>.txt
|
├── labels/{train,val}/<类别>_<序号>.txt
|
||||||
|
├── classes.txt # 每行一个类别名,行号即 class_id
|
||||||
|
├── name_map.csv # 新文件名 ↔ 源文件路径
|
||||||
├── data.yaml # path 写绝对路径,ultralytics 从哪跑都能读到
|
├── 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,两次运行结果一致)。分层是必需的:
|
按类别分层(`--val-ratio` 默认 0.2,`--seed` 默认 42,两次运行结果一致)。分层是必需的:
|
||||||
|
|||||||
@@ -18,6 +18,21 @@ from pathlib import Path
|
|||||||
IMAGE_SUFFIXES = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
|
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(
|
def write_data_yaml(
|
||||||
path: str | Path,
|
path: str | Path,
|
||||||
classes: Sequence[str],
|
classes: Sequence[str],
|
||||||
|
|||||||
+114
-15
@@ -16,13 +16,22 @@ runs/boxes_gradcam/
|
|||||||
|
|
||||||
1. 把按类别分的目录拍平,按 train/val 重新组织;
|
1. 把按类别分的目录拍平,按 train/val 重新组织;
|
||||||
2. **按类别分层**切分,小类不至于在验证集里挂零;
|
2. **按类别分层**切分,小类不至于在验证集里挂零;
|
||||||
3. 全部用**硬链接**,不占额外磁盘。
|
3. 图片**统一成 PNG**:本来就是 PNG 的硬链接过去,不占额外磁盘;BMP 等其他格式
|
||||||
|
(pearlescent / blue / uv 原始数据都是 BMP)无损转成 PNG。不缩放像素,
|
||||||
|
归一化的框坐标原样有效。
|
||||||
|
|
||||||
|
4. 文件**按类别重新编号**:`<类别>_<序号>.png`,序号在类内按源路径排序,重跑结果一致。
|
||||||
|
原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里,不会丢。
|
||||||
|
想保留原名加 `--keep-names`。
|
||||||
|
|
||||||
```
|
```
|
||||||
data/<name>/
|
data/<name>/
|
||||||
├── images/{train,val}/<图名>.png
|
├── images/{train,val}/<类别>_<序号>.png
|
||||||
├── labels/{train,val}/<图名>.txt
|
├── labels/{train,val}/<类别>_<序号>.txt
|
||||||
└── data.yaml
|
├── classes.txt # 每行一个类别名,行号即 class_id,给标注工具复核用
|
||||||
|
├── name_map.csv # 新名 ↔ 源文件对照
|
||||||
|
├── data.yaml
|
||||||
|
└── summary.json
|
||||||
```
|
```
|
||||||
|
|
||||||
没有框的图(标签为空)默认保留:YOLO 训练里它们是负样本,有助于压低误检。但**空标签的
|
没有框的图(标签为空)默认保留:YOLO 训练里它们是负样本,有助于压低误检。但**空标签的
|
||||||
@@ -36,6 +45,7 @@ data/<name>/
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
|
import csv
|
||||||
import errno
|
import errno
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
@@ -51,7 +61,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_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
|
from utils import setup_logging # noqa: E402
|
||||||
|
|
||||||
logger = None
|
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:
|
def link_or_copy(source: Path, destination: Path) -> str:
|
||||||
"""硬链接;跨文件系统时退回复制。"""
|
"""硬链接;跨文件系统时退回复制。"""
|
||||||
if destination.exists() or destination.is_symlink():
|
if destination.exists() or destination.is_symlink():
|
||||||
@@ -133,6 +174,35 @@ def collect(source: Path) -> list[Sample]:
|
|||||||
return samples
|
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]]:
|
def parse_empty_policy(tokens: Sequence[str]) -> tuple[str, dict[str, str]]:
|
||||||
"""解析 --empty-policy:一个可选的全局默认,加若干 `类别=策略`。
|
"""解析 --empty-policy:一个可选的全局默认,加若干 `类别=策略`。
|
||||||
|
|
||||||
@@ -214,6 +284,10 @@ def main() -> int:
|
|||||||
"注意空标签在 YOLO 里的意思是『此图无缺陷』,缺陷类的空标签照搬会教模型抑制真实缺陷"
|
"注意空标签在 YOLO 里的意思是『此图无缺陷』,缺陷类的空标签照搬会教模型抑制真实缺陷"
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--keep-names", action="store_true",
|
||||||
|
help="保留源文件名,不按 <类别>_<序号> 重新编号(源文件名需去掉后缀后全局唯一)",
|
||||||
|
)
|
||||||
parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
|
parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
|
||||||
parser.add_argument("--dry-run", action="store_true", help="只打印计划")
|
parser.add_argument("--dry-run", action="store_true", help="只打印计划")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
@@ -291,6 +365,13 @@ def main() -> int:
|
|||||||
if counts["train"][name] or counts["val"][name]:
|
if counts["train"][name] or counts["val"][name]:
|
||||||
logger.info("%-14s %8d %8d", name, counts["train"][name], 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:
|
if args.dry_run:
|
||||||
logger.info("--dry-run:未写入任何文件")
|
logger.info("--dry-run:未写入任何文件")
|
||||||
return 0
|
return 0
|
||||||
@@ -301,22 +382,33 @@ def main() -> int:
|
|||||||
return 1
|
return 1
|
||||||
shutil.rmtree(out_dir)
|
shutil.rmtree(out_dir)
|
||||||
|
|
||||||
|
renamed = None if args.keep_names else assign_names(samples, names)
|
||||||
modes: Counter[str] = Counter()
|
modes: Counter[str] = Counter()
|
||||||
seen: dict[str, Path] = {}
|
seen: dict[str, Path] = {}
|
||||||
|
name_rows = []
|
||||||
for split, items in splits.items():
|
for split, items in splits.items():
|
||||||
for sample in items:
|
for sample in items:
|
||||||
image = sample.image
|
image = sample.image
|
||||||
# 拍平之后不同类别目录下的同名文件会撞车,显式挡住
|
stem = image.stem if renamed is None else renamed[image]
|
||||||
if image.name in seen:
|
output_name = f"{stem}.png"
|
||||||
logger.error(
|
# 保留原名时,拍平并统一后缀后不同目录下的 a.png 与 a.bmp 会撞成同一个 a.png;
|
||||||
"拍平后文件名冲突:%s 与 %s 同名。源数据集的文件名需要全局唯一,"
|
# 重新编号时只有类别名仅大小写不同才会撞。按小写比较:macOS 默认文件系统
|
||||||
"用 scripts/export_clusters.py --prefix-parent 重新导出即可",
|
# 不区分大小写,A.png 与 a.png 也会互相覆盖。
|
||||||
seen[image.name], image,
|
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
|
return 1
|
||||||
seen[image.name] = image
|
seen[key] = image
|
||||||
modes[link_or_copy(image, out_dir / "images" / split / image.name)] += 1
|
modes[materialize_png(image, out_dir / "images" / split / output_name)] += 1
|
||||||
destination = out_dir / "labels" / split / f"{image.stem}.txt"
|
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:
|
if sample.full_box is None:
|
||||||
link_or_copy(sample.label, destination)
|
link_or_copy(sample.label, destination)
|
||||||
else:
|
else:
|
||||||
@@ -326,6 +418,12 @@ def main() -> int:
|
|||||||
|
|
||||||
write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(),
|
write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(),
|
||||||
train="images/train", val="images/val")
|
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(
|
(out_dir / "summary.json").write_text(
|
||||||
json.dumps(
|
json.dumps(
|
||||||
{
|
{
|
||||||
@@ -336,6 +434,7 @@ def main() -> int:
|
|||||||
for group, (policy, count) in sorted(empty_stats.items())
|
for group, (policy, count) in sorted(empty_stats.items())
|
||||||
},
|
},
|
||||||
"empty_policy": {"default": default_policy, **per_class_policy},
|
"empty_policy": {"default": default_policy, **per_class_policy},
|
||||||
|
"renamed": renamed is not None,
|
||||||
"val_ratio": args.val_ratio,
|
"val_ratio": args.val_ratio,
|
||||||
"seed": args.seed,
|
"seed": args.seed,
|
||||||
"split_sizes": {k: len(v) for k, v in splits.items()},
|
"split_sizes": {k: len(v) for k, v in splits.items()},
|
||||||
@@ -347,7 +446,7 @@ def main() -> int:
|
|||||||
),
|
),
|
||||||
encoding="utf-8",
|
encoding="utf-8",
|
||||||
)
|
)
|
||||||
logger.info("链接方式: %s", dict(modes))
|
logger.info("图片落盘方式: %s(hard=硬链接, copy=跨盘复制, convert=转成 PNG)", dict(modes))
|
||||||
logger.info("完成 -> %s", out_dir)
|
logger.info("完成 -> %s", out_dir)
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""
|
"""把 generate_boxes.py 的产出整理成 YOLO 训练集。"""
|
||||||
|
|
||||||
|
import csv
|
||||||
import json
|
import json
|
||||||
import sys
|
import sys
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -44,6 +45,15 @@ def source(tmp_path):
|
|||||||
return root
|
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):
|
def run(source, out, *extra):
|
||||||
argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra]
|
argv = ["prepare_yolo_dataset.py", "--source", str(source), "--out", str(out), *extra]
|
||||||
original, sys.argv = sys.argv, argv
|
original, sys.argv = sys.argv, argv
|
||||||
@@ -72,8 +82,8 @@ def test_hard_links_share_inode(source, tmp_path):
|
|||||||
"""用硬链接而不是复制,不占额外磁盘。"""
|
"""用硬链接而不是复制,不占额外磁盘。"""
|
||||||
out = tmp_path / "data" / "boxes"
|
out = tmp_path / "data" / "boxes"
|
||||||
assert run(source, out) == 0
|
assert run(source, out) == 0
|
||||||
linked = next((out / "images" / "train").iterdir())
|
original = source / "images" / "grease" / "grease_0.png"
|
||||||
original = next(p for p in (source / "images").rglob("*.png") if p.name == linked.name)
|
linked = output_of(out, "grease_0")
|
||||||
assert linked.stat().st_ino == original.stat().st_ino
|
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"))
|
summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
|
||||||
assert summary["num_images"] == 12
|
assert summary["num_images"] == 12
|
||||||
assert summary["empty_labels"]["normal"]["policy"] == "drop"
|
assert summary["empty_labels"]["normal"]["policy"] == "drop"
|
||||||
assert not (out / "labels" / "train" / "normal_0.txt").exists()
|
assert output_of(out, "normal_0") is None
|
||||||
assert not (out / "images" / "train" / "normal_0.png").exists()
|
|
||||||
|
|
||||||
|
|
||||||
def test_empty_policy_is_per_class(source, tmp_path):
|
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"
|
out = tmp_path / "data" / "boxes"
|
||||||
assert run(source, out, "--empty-policy", "normal=full") == 0
|
assert run(source, out, "--empty-policy", "normal=full") == 0
|
||||||
|
|
||||||
label = next(
|
image = output_of(out, "normal_0")
|
||||||
p for split in ("train", "val")
|
label = out / "labels" / image.parent.name / f"{image.stem}.txt"
|
||||||
for p in (out / "labels" / split).iterdir()
|
|
||||||
if p.stem == "normal_0"
|
|
||||||
)
|
|
||||||
assert label.read_text(encoding="utf-8").split() == ["1", "0.5", "0.5", "1.0", "1.0"]
|
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") == ""
|
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):
|
def test_name_collision_after_flattening_is_reported(source, tmp_path, capsys):
|
||||||
"""不同类别目录下的同名文件,拍平后会互相覆盖——必须报错而不是静默丢数据。"""
|
"""保留原名时,不同类别目录下的同名文件拍平后会互相覆盖——必须报错而不是静默丢数据。"""
|
||||||
duplicate = source / "images" / "normal" / "grease_0.png"
|
duplicate = source / "images" / "normal" / "grease_0.png"
|
||||||
Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)).save(duplicate)
|
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")
|
(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
|
assert "冲突" in capsys.readouterr().out
|
||||||
|
|
||||||
|
|
||||||
@@ -201,3 +207,183 @@ def test_dry_run_writes_nothing(source, tmp_path):
|
|||||||
out = tmp_path / "data" / "boxes"
|
out = tmp_path / "data" / "boxes"
|
||||||
assert run(source, out, "--dry-run") == 0
|
assert run(source, out, "--dry-run") == 0
|
||||||
assert not out.exists()
|
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