Shaun ZhouandClaude Opus 5
4b1147ed8f
- convert_annotations: labelme / VOC -> YOLO txt，顺带生成 data.yaml
- split_dataset: 切分成 images/{train,val} 与 labels/{train,val}，默认建软链不复制数据
- export_onnx: 导出部署用 ONNX，CPU 可跑，缺依赖时提示 `uv sync --group export`
- profile_inference: torch.profiler 测单帧延迟/FPS、CUDA 显存峰值与算子级明细
- check_env: Python / PyTorch / Nvidia GPU 自检
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/film/images --labels data/film/labels \
|
|
--val-ratio 0.2 --data-yaml data/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())
|