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
+114
-15
@@ -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/<name>/
|
||||
├── 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/<name>/
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user