Shaun ZhouandClaude Opus 5
5d395b7c25
原始数据从 data/<数据集> 移到了 data/raw/<数据集>。
- 聚类配置的 dataset.root：film / pearlescent 均改为 data/raw/ 下
- 检测配置的 img_dir 与 label_dir 注释同步
- README、data/README、算法文档、流程图（含重新渲染的 PNG）与各脚本
docstring 里的路径示例一并更新
- data/README 补充目录约定：原始数据只读放 data/raw/，
export_clusters.py 的派生数据集放在 data/ 下与 raw/ 平级
- 顺手把 unlabeled 数据集加进 datasets/__init__.py：此前它没在包初始化时注册，
build_dataset({"name": "unlabeled"}) 只有在别处先 import 过该模块时才可用
迁移后重跑全量聚类，结果与迁移前逐项一致（轮廓 0.3255、稳定性 0.9641、
簇大小 [4507, 3795, 4849, 1428]）。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
105 lines
3.6 KiB
Python
Executable File
105 lines
3.6 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""按比例切分成 YOLO 的 train / val / test 目录结构。
|
|
|
|
python scripts/split_dataset.py --images data/raw/film/images --labels data/raw/film/labels \
|
|
--val-ratio 0.2 --data-yaml data/raw/film/data.yaml
|
|
|
|
默认建软链（不复制数据），产出：
|
|
|
|
```
|
|
images/train/ images/val/
|
|
labels/train/ labels/val/
|
|
```
|
|
|
|
原始文件保留在 `images/` `labels/` 根目录下不动；重复执行会先清掉旧的划分目录。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import random
|
|
import shutil
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from datasets.yolo import IMAGE_SUFFIXES, load_classes, write_data_yaml # noqa: E402
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(
|
|
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
|
)
|
|
parser.add_argument("--images", required=True, help="图片目录（划分前的平铺目录）")
|
|
parser.add_argument("--labels", default=None, help="标注目录，缺省为 images 同级的 labels")
|
|
parser.add_argument("--val-ratio", type=float, default=0.2)
|
|
parser.add_argument("--test-ratio", type=float, default=0.0)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--copy", action="store_true", help="复制文件而不是建软链")
|
|
parser.add_argument("--data-yaml", default=None, help="顺带更新/生成 data.yaml")
|
|
parser.add_argument("--classes", nargs="*", default=None, help="写 data.yaml 时的类别顺序")
|
|
args = parser.parse_args()
|
|
|
|
images_dir = Path(args.images)
|
|
labels_dir = Path(args.labels) if args.labels else images_dir.parent / "labels"
|
|
splits = ("train", "val", "test")
|
|
|
|
files = sorted(
|
|
p
|
|
for p in images_dir.iterdir()
|
|
if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES and p.parent.name not in splits
|
|
)
|
|
if not files:
|
|
print(f"{images_dir} 下没有待划分的图片（子目录里的已划分文件会被跳过）")
|
|
return 1
|
|
|
|
random.Random(args.seed).shuffle(files)
|
|
n_val = int(len(files) * args.val_ratio)
|
|
n_test = int(len(files) * args.test_ratio)
|
|
partition = {
|
|
"val": files[:n_val],
|
|
"test": files[n_val : n_val + n_test],
|
|
"train": files[n_val + n_test :],
|
|
}
|
|
|
|
for name, items in partition.items():
|
|
img_out, label_out = images_dir / name, labels_dir / name
|
|
for directory in (img_out, label_out):
|
|
if directory.exists():
|
|
shutil.rmtree(directory)
|
|
if not items:
|
|
continue
|
|
img_out.mkdir(parents=True, exist_ok=True)
|
|
label_out.mkdir(parents=True, exist_ok=True)
|
|
|
|
empty = 0
|
|
for image in items:
|
|
label = labels_dir / f"{image.stem}.txt"
|
|
place(image, img_out / image.name, args.copy)
|
|
if label.exists():
|
|
place(label, label_out / label.name, args.copy)
|
|
empty += int(not label.read_text(encoding="utf-8").strip())
|
|
else:
|
|
empty += 1 # 缺 .txt 等价于无缺陷
|
|
print(f"{name}: {len(items)} 张（其中无缺陷 {empty} 张） -> {img_out}")
|
|
|
|
if args.data_yaml:
|
|
classes = load_classes(args.classes or args.data_yaml)
|
|
path = write_data_yaml(args.data_yaml, classes, root=images_dir.parent)
|
|
print(f"data.yaml 已更新 {path}")
|
|
return 0
|
|
|
|
|
|
def place(src: Path, dst: Path, copy: bool) -> None:
|
|
if dst.exists() or dst.is_symlink():
|
|
dst.unlink()
|
|
if copy:
|
|
shutil.copy2(src, dst)
|
|
else:
|
|
dst.symlink_to(src.resolve())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|