orisys
520 lines
20 KiB
Python
520 lines
20 KiB
Python
"""YOLO 数据集划分：可创建输出文件夹，自定义 train/val/test 百分比（带界面）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import random
|
||
import shutil
|
||
import tkinter as tk
|
||
from pathlib import Path
|
||
from tkinter import filedialog, messagebox, ttk
|
||
|
||
import yaml
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"}
|
||
|
||
|
||
def iter_images(folder: Path) -> list[Path]:
|
||
if not folder.is_dir():
|
||
return []
|
||
return sorted(p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS)
|
||
|
||
|
||
def find_label_for(img: Path, src: Path, split_hint: str | None = None) -> Path | None:
|
||
"""按 stem 在常见 labels 位置查找同名 txt。"""
|
||
candidates: list[Path] = []
|
||
if split_hint:
|
||
candidates.append(src / "labels" / split_hint / f"{img.stem}.txt")
|
||
candidates.extend(
|
||
[
|
||
src / "labels" / "all" / f"{img.stem}.txt",
|
||
src / "labels" / f"{img.stem}.txt",
|
||
img.with_suffix(".txt"),
|
||
img.parent.parent / "labels" / img.parent.name / f"{img.stem}.txt",
|
||
img.parent.parent / "labels" / f"{img.stem}.txt",
|
||
]
|
||
)
|
||
for cand in candidates:
|
||
if cand.exists():
|
||
return cand
|
||
return None
|
||
|
||
|
||
def find_image_label_pairs(src: Path) -> list[tuple[Path, Path | None]]:
|
||
"""收集源目录全部样本。
|
||
|
||
支持:
|
||
- images/{all|train|val|test} + labels/...
|
||
- images/ 扁平目录
|
||
- 根目录直接放图片
|
||
若已有 train/val，会合并后再重新划分。
|
||
"""
|
||
pairs_map: dict[str, tuple[Path, Path | None]] = {}
|
||
|
||
split_dirs = []
|
||
for name in ("all", "train", "val", "test"):
|
||
d = src / "images" / name
|
||
if d.is_dir() and any(d.iterdir()):
|
||
split_dirs.append((name if name != "all" else None, d))
|
||
|
||
if split_dirs:
|
||
for hint, img_dir in split_dirs:
|
||
for img in iter_images(img_dir):
|
||
pairs_map[img.stem] = (img, find_label_for(img, src, hint))
|
||
else:
|
||
flat_candidates = [src / "images", src]
|
||
img_dir = next((p for p in flat_candidates if iter_images(p)), None)
|
||
if img_dir is None:
|
||
raise FileNotFoundError(f"找不到图片目录，请检查: {src}")
|
||
for img in iter_images(img_dir):
|
||
pairs_map[img.stem] = (img, find_label_for(img, src, None))
|
||
|
||
pairs = [pairs_map[k] for k in sorted(pairs_map)]
|
||
if not pairs:
|
||
raise FileNotFoundError(f"在 {src} 下没有找到图片")
|
||
return pairs
|
||
|
||
|
||
def validate_pairs(pairs: list[tuple[Path, Path | None]]) -> list[str]:
|
||
warnings: list[str] = []
|
||
missing = [img.name for img, lb in pairs if lb is None]
|
||
if missing:
|
||
warnings.append(f"有 {len(missing)} 张图缺少标签，例如: {', '.join(missing[:3])}")
|
||
bad = 0
|
||
for img, lb in pairs:
|
||
if lb is None:
|
||
continue
|
||
text = lb.read_text(encoding="utf-8").strip()
|
||
if not text:
|
||
continue
|
||
for line in text.splitlines():
|
||
parts = line.split()
|
||
if len(parts) != 5:
|
||
bad += 1
|
||
continue
|
||
try:
|
||
vals = list(map(float, parts[1:]))
|
||
except ValueError:
|
||
bad += 1
|
||
continue
|
||
if any(v < 0 or v > 1 for v in vals):
|
||
bad += 1
|
||
if bad:
|
||
warnings.append(f"有 {bad} 行标签格式异常（应为: class x y w h，且坐标 0~1）")
|
||
return warnings
|
||
|
||
|
||
|
||
def split_indices(n: int, train_r: float, val_r: float, test_r: float, seed: int) -> dict[str, list[int]]:
|
||
total = train_r + val_r + test_r
|
||
if abs(total - 1.0) > 1e-6:
|
||
raise ValueError(f"train+val+test 必须等于 100%，当前为 {total * 100:.1f}%")
|
||
if n < 2:
|
||
raise ValueError("样本太少，至少需要 2 张图")
|
||
if test_r > 0 and n < 3:
|
||
raise ValueError("划分测试集至少需要 3 张图")
|
||
|
||
idxs = list(range(n))
|
||
random.Random(seed).shuffle(idxs)
|
||
|
||
n_test = int(round(n * test_r)) if test_r > 0 else 0
|
||
n_val = int(round(n * val_r))
|
||
n_train = n - n_val - n_test
|
||
if n_train < 1:
|
||
raise ValueError("训练集为空，请调大训练比例")
|
||
if n_val < 1:
|
||
raise ValueError("验证集为空，请调大验证比例")
|
||
if test_r > 0 and n_test < 1:
|
||
raise ValueError("测试集为空，请调大测试比例")
|
||
|
||
return {
|
||
"train": idxs[:n_train],
|
||
"val": idxs[n_train : n_train + n_val],
|
||
"test": idxs[n_train + n_val :] if n_test else [],
|
||
}
|
||
|
||
|
||
def transfer(src: Path, dst: Path, move: bool) -> None:
|
||
dst.parent.mkdir(parents=True, exist_ok=True)
|
||
if move:
|
||
shutil.move(str(src), str(dst))
|
||
else:
|
||
shutil.copy2(src, dst)
|
||
|
||
|
||
def write_data_yaml(yaml_path: Path, dataset_root: Path, has_test: bool, names: dict[int, str]) -> None:
|
||
try:
|
||
rel = dataset_root.resolve().relative_to(ROOT.resolve()).as_posix()
|
||
except ValueError:
|
||
rel = dataset_root.resolve().as_posix()
|
||
|
||
data = {
|
||
"path": rel,
|
||
"train": "images/train",
|
||
"val": "images/val",
|
||
"test": "images/test" if has_test else None,
|
||
"names": names,
|
||
}
|
||
if data["test"] is None:
|
||
data.pop("test")
|
||
|
||
yaml_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with yaml_path.open("w", encoding="utf-8") as f:
|
||
f.write("# 由 split_dataset.py 自动生成\n")
|
||
f.write("# 训练: yolo train cfg=train.yaml\n\n")
|
||
yaml.safe_dump(data, f, allow_unicode=True, sort_keys=False)
|
||
|
||
|
||
def load_or_build_names(yaml_path: Path, names_arg: str) -> dict[int, str]:
|
||
if names_arg.strip():
|
||
return {i: n.strip() for i, n in enumerate(names_arg.split(",")) if n.strip()}
|
||
if yaml_path.exists():
|
||
with yaml_path.open(encoding="utf-8") as f:
|
||
old = yaml.safe_load(f) or {}
|
||
old_names = old.get("names")
|
||
if isinstance(old_names, dict):
|
||
return {int(k): str(v) for k, v in old_names.items()}
|
||
if isinstance(old_names, list):
|
||
return {i: str(v) for i, v in enumerate(old_names)}
|
||
return {0: "object"}
|
||
|
||
|
||
def run_split(
|
||
src: Path,
|
||
dst: Path,
|
||
train_pct: float,
|
||
val_pct: float,
|
||
test_pct: float,
|
||
seed: int,
|
||
move: bool,
|
||
names_arg: str,
|
||
yaml_path: Path,
|
||
) -> str:
|
||
if not src.exists():
|
||
raise FileNotFoundError(f"源目录不存在: {src}")
|
||
|
||
dst.mkdir(parents=True, exist_ok=True)
|
||
pairs = find_image_label_pairs(src)
|
||
warnings = validate_pairs(pairs)
|
||
missing = sum(1 for _, lb in pairs if lb is None)
|
||
splits = split_indices(len(pairs), train_pct / 100.0, val_pct / 100.0, test_pct / 100.0, seed)
|
||
|
||
same_dir = src.resolve() == dst.resolve()
|
||
stage = dst / "_split_stage"
|
||
if stage.exists():
|
||
shutil.rmtree(stage)
|
||
stage_img = stage / "images"
|
||
stage_lbl = stage / "labels"
|
||
stage_img.mkdir(parents=True)
|
||
stage_lbl.mkdir(parents=True)
|
||
|
||
staged: list[tuple[Path, Path | None]] = []
|
||
for img, label in pairs:
|
||
new_img = stage_img / img.name
|
||
transfer(img, new_img, move=False)
|
||
new_lbl = None
|
||
if label is not None:
|
||
new_lbl = stage_lbl / label.name
|
||
transfer(label, new_lbl, move=False)
|
||
staged.append((new_img, new_lbl))
|
||
|
||
if move:
|
||
# 源文件已暂存后，再删除源侧原文件
|
||
for img, label in pairs:
|
||
img.unlink(missing_ok=True)
|
||
if label is not None:
|
||
label.unlink(missing_ok=True)
|
||
|
||
for split in ("train", "val", "test"):
|
||
for kind in ("images", "labels"):
|
||
out = dst / kind / split
|
||
if out.exists():
|
||
shutil.rmtree(out)
|
||
out.mkdir(parents=True, exist_ok=True)
|
||
|
||
for split, idxs in splits.items():
|
||
if not idxs:
|
||
continue
|
||
img_out = dst / "images" / split
|
||
lbl_out = dst / "labels" / split
|
||
for i in idxs:
|
||
img, label = staged[i]
|
||
transfer(img, img_out / img.name, move=True)
|
||
if label is not None:
|
||
transfer(label, lbl_out / label.name, move=True)
|
||
else:
|
||
(lbl_out / f"{img.stem}.txt").write_text("", encoding="utf-8")
|
||
|
||
shutil.rmtree(stage, ignore_errors=True)
|
||
|
||
names = load_or_build_names(yaml_path, names_arg)
|
||
write_data_yaml(yaml_path, dst, has_test=bool(splits["test"]), names=names)
|
||
|
||
lines = [
|
||
f"源目录: {src}",
|
||
f"输出目录: {dst}" + ("（原地重划分）" if same_dir else ""),
|
||
f"总样本: {len(pairs)} 无标签: {missing}",
|
||
]
|
||
for split, idxs in splits.items():
|
||
if idxs:
|
||
lines.append(f" {split}: {len(idxs)}")
|
||
for w in warnings:
|
||
lines.append(f"警告: {w}")
|
||
lines.append(f"已生成: {yaml_path}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
class SplitApp(tk.Tk):
|
||
def __init__(self) -> None:
|
||
super().__init__()
|
||
self.title("YOLO 数据集划分")
|
||
self.geometry("760x520")
|
||
self.minsize(700, 480)
|
||
|
||
self.var_src = tk.StringVar(value=str(ROOT / "datasets" / "phonedata"))
|
||
self.var_dst = tk.StringVar(value=str(ROOT / "datasets" / "phonedata"))
|
||
self.var_yaml = tk.StringVar(value=str(ROOT / "data.yaml"))
|
||
self.var_names = tk.StringVar(value="phone")
|
||
self.var_train = tk.DoubleVar(value=80)
|
||
self.var_val = tk.DoubleVar(value=20)
|
||
self.var_test = tk.DoubleVar(value=0)
|
||
self.var_use_test = tk.BooleanVar(value=False)
|
||
self.var_seed = tk.IntVar(value=42)
|
||
self.var_move = tk.BooleanVar(value=False)
|
||
self.var_sum = tk.StringVar(value="合计: 100%")
|
||
|
||
self._build()
|
||
self.on_test_toggle()
|
||
self.refresh_sum()
|
||
|
||
def _build(self) -> None:
|
||
pad = {"padx": 8, "pady": 6}
|
||
frm = ttk.Frame(self, padding=12)
|
||
frm.pack(fill="both", expand=True)
|
||
frm.columnconfigure(1, weight=1)
|
||
|
||
ttk.Label(frm, text="源数据文件夹").grid(row=0, column=0, sticky="w", **pad)
|
||
ttk.Entry(frm, textvariable=self.var_src).grid(row=0, column=1, sticky="ew", **pad)
|
||
ttk.Button(frm, text="浏览", command=self.browse_src).grid(row=0, column=2, **pad)
|
||
|
||
ttk.Label(frm, text="输出文件夹").grid(row=1, column=0, sticky="w", **pad)
|
||
ttk.Entry(frm, textvariable=self.var_dst).grid(row=1, column=1, sticky="ew", **pad)
|
||
btns = ttk.Frame(frm)
|
||
btns.grid(row=1, column=2, **pad)
|
||
ttk.Button(btns, text="浏览", command=self.browse_dst).pack(side="left")
|
||
ttk.Button(btns, text="新建", command=self.create_dst).pack(side="left", padx=(6, 0))
|
||
|
||
ttk.Label(frm, text="data.yaml").grid(row=2, column=0, sticky="w", **pad)
|
||
ttk.Entry(frm, textvariable=self.var_yaml).grid(row=2, column=1, sticky="ew", **pad)
|
||
ttk.Button(frm, text="浏览", command=self.browse_yaml).grid(row=2, column=2, **pad)
|
||
|
||
ttk.Label(frm, text="类别名").grid(row=3, column=0, sticky="w", **pad)
|
||
ttk.Entry(frm, textvariable=self.var_names).grid(row=3, column=1, sticky="ew", **pad)
|
||
ttk.Label(frm, text="逗号分隔，如 phone,cable").grid(row=3, column=2, sticky="w", **pad)
|
||
|
||
ratio = ttk.LabelFrame(frm, text="划分百分比（合计必须 100%）", padding=10)
|
||
ratio.grid(row=4, column=0, columnspan=3, sticky="ew", **pad)
|
||
for c in range(6):
|
||
ratio.columnconfigure(c, weight=1)
|
||
|
||
ttk.Label(ratio, text="训练 %").grid(row=0, column=0, sticky="w")
|
||
train_entry = ttk.Entry(ratio, textvariable=self.var_train, width=8)
|
||
train_entry.grid(row=0, column=1, sticky="w")
|
||
ttk.Label(ratio, text="验证 %").grid(row=0, column=2, sticky="w")
|
||
val_entry = ttk.Entry(ratio, textvariable=self.var_val, width=8)
|
||
val_entry.grid(row=0, column=3, sticky="w")
|
||
ttk.Checkbutton(ratio, text="启用测试集", variable=self.var_use_test, command=self.on_test_toggle).grid(
|
||
row=0, column=4, sticky="w"
|
||
)
|
||
self.test_entry = ttk.Entry(ratio, textvariable=self.var_test, width=8)
|
||
self.test_entry.grid(row=0, column=5, sticky="w")
|
||
ttk.Label(ratio, textvariable=self.var_sum).grid(row=1, column=0, columnspan=3, sticky="w", pady=(8, 0))
|
||
ttk.Label(ratio, text="快捷:").grid(row=1, column=3, sticky="e", pady=(8, 0))
|
||
ttk.Button(ratio, text="80/20", command=lambda: self.set_ratio(80, 20, 0)).grid(row=1, column=4, sticky="w", pady=(8, 0))
|
||
ttk.Button(ratio, text="70/20/10", command=lambda: self.set_ratio(70, 20, 10)).grid(
|
||
row=1, column=5, sticky="w", pady=(8, 0)
|
||
)
|
||
|
||
for var in (self.var_train, self.var_val, self.var_test):
|
||
var.trace_add("write", lambda *_: self.refresh_sum())
|
||
|
||
opts = ttk.Frame(frm)
|
||
opts.grid(row=5, column=0, columnspan=3, sticky="ew", **pad)
|
||
ttk.Label(opts, text="随机种子").pack(side="left")
|
||
ttk.Entry(opts, textvariable=self.var_seed, width=8).pack(side="left", padx=(6, 16))
|
||
ttk.Checkbutton(opts, text="移动文件（默认复制）", variable=self.var_move).pack(side="left")
|
||
|
||
ttk.Button(frm, text="扫描源目录", command=self.scan_src).grid(row=6, column=0, columnspan=1, sticky="ew", pady=10)
|
||
ttk.Button(frm, text="开始划分", command=self.run).grid(row=6, column=1, columnspan=2, sticky="ew", pady=10)
|
||
|
||
self.log = tk.Text(frm, height=12, wrap="word")
|
||
self.log.grid(row=7, column=0, columnspan=3, sticky="nsew")
|
||
frm.rowconfigure(7, weight=1)
|
||
self.log.insert("end", "说明：\n")
|
||
self.log.insert("end", "1. 源目录可为未划分数据，或已有 train/val（会合并后重新划分）\n")
|
||
self.log.insert("end", "2. 输出目录可新建；结构: images/train|val|test + labels/...\n")
|
||
self.log.insert("end", "3. 百分比自定义，启用测试集时 train+val+test=100\n")
|
||
self.log.insert("end", "4. 当前 phonedata 格式正确时可直接点「扫描源目录」确认数量\n")
|
||
|
||
def browse_src(self) -> None:
|
||
path = filedialog.askdirectory(initialdir=self.var_src.get() or str(ROOT))
|
||
if path:
|
||
self.var_src.set(path)
|
||
|
||
def browse_dst(self) -> None:
|
||
path = filedialog.askdirectory(initialdir=self.var_dst.get() or str(ROOT))
|
||
if path:
|
||
self.var_dst.set(path)
|
||
|
||
def create_dst(self) -> None:
|
||
parent = filedialog.askdirectory(title="选择新建文件夹的父目录", initialdir=str(ROOT / "datasets"))
|
||
if not parent:
|
||
return
|
||
dialog = tk.Toplevel(self)
|
||
dialog.title("新建输出文件夹")
|
||
dialog.geometry("420x140")
|
||
dialog.transient(self)
|
||
dialog.grab_set()
|
||
name_var = tk.StringVar(value="split")
|
||
ttk.Label(dialog, text=f"父目录: {parent}").pack(anchor="w", padx=12, pady=(12, 4))
|
||
row = ttk.Frame(dialog)
|
||
row.pack(fill="x", padx=12)
|
||
ttk.Label(row, text="文件夹名").pack(side="left")
|
||
ttk.Entry(row, textvariable=name_var).pack(side="left", fill="x", expand=True, padx=8)
|
||
|
||
def ok() -> None:
|
||
name = name_var.get().strip()
|
||
if not name:
|
||
messagebox.showwarning("新建", "请输入文件夹名", parent=dialog)
|
||
return
|
||
path = Path(parent) / name
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
self.var_dst.set(str(path))
|
||
dialog.destroy()
|
||
|
||
ttk.Button(dialog, text="创建", command=ok).pack(pady=12)
|
||
|
||
def browse_yaml(self) -> None:
|
||
path = filedialog.asksaveasfilename(
|
||
initialdir=str(ROOT),
|
||
initialfile="data.yaml",
|
||
defaultextension=".yaml",
|
||
filetypes=[("YAML", "*.yaml;*.yml"), ("All", "*.*")],
|
||
)
|
||
if path:
|
||
self.var_yaml.set(path)
|
||
|
||
def set_ratio(self, train: float, val: float, test: float) -> None:
|
||
self.var_use_test.set(test > 0)
|
||
self.on_test_toggle()
|
||
self.var_train.set(train)
|
||
self.var_val.set(val)
|
||
self.var_test.set(test)
|
||
self.refresh_sum()
|
||
|
||
def on_test_toggle(self) -> None:
|
||
if self.var_use_test.get():
|
||
self.test_entry.state(["!disabled"])
|
||
if float(self.var_test.get() or 0) <= 0:
|
||
self.var_test.set(10)
|
||
self.var_train.set(70)
|
||
self.var_val.set(20)
|
||
else:
|
||
self.var_test.set(0)
|
||
self.test_entry.state(["disabled"])
|
||
self.refresh_sum()
|
||
|
||
def refresh_sum(self) -> None:
|
||
try:
|
||
train = float(self.var_train.get())
|
||
val = float(self.var_val.get())
|
||
test = float(self.var_test.get()) if self.var_use_test.get() else 0.0
|
||
total = train + val + test
|
||
ok = abs(total - 100.0) < 0.05
|
||
self.var_sum.set(f"合计: {total:.1f}% {'✓' if ok else '（必须等于 100%）'}")
|
||
except (tk.TclError, ValueError, TypeError):
|
||
self.var_sum.set("合计: --")
|
||
|
||
def scan_src(self) -> None:
|
||
try:
|
||
pairs = find_image_label_pairs(Path(self.var_src.get()))
|
||
warnings = validate_pairs(pairs)
|
||
missing = sum(1 for _, lb in pairs if lb is None)
|
||
msg = f"扫描完成: 共 {len(pairs)} 张图，缺标签 {missing}\n"
|
||
for w in warnings:
|
||
msg += f"警告: {w}\n"
|
||
if not warnings and missing == 0:
|
||
msg += "图文配对与标签格式检查通过。\n"
|
||
except Exception as exc:
|
||
messagebox.showerror("扫描失败", str(exc))
|
||
self.log.insert("end", f"\n扫描错误: {exc}\n")
|
||
return
|
||
self.log.insert("end", f"\n{msg}")
|
||
self.log.see("end")
|
||
messagebox.showinfo("扫描结果", msg)
|
||
|
||
def run(self) -> None:
|
||
try:
|
||
train = float(self.var_train.get())
|
||
val = float(self.var_val.get())
|
||
test = float(self.var_test.get()) if self.var_use_test.get() else 0.0
|
||
if abs(train + val + test - 100.0) > 0.05:
|
||
raise ValueError("百分比合计必须等于 100")
|
||
msg = run_split(
|
||
src=Path(self.var_src.get()),
|
||
dst=Path(self.var_dst.get()),
|
||
train_pct=train,
|
||
val_pct=val,
|
||
test_pct=test,
|
||
seed=int(self.var_seed.get()),
|
||
move=self.var_move.get(),
|
||
names_arg=self.var_names.get(),
|
||
yaml_path=Path(self.var_yaml.get()),
|
||
)
|
||
except Exception as exc:
|
||
messagebox.showerror("划分失败", str(exc))
|
||
self.log.insert("end", f"\n错误: {exc}\n")
|
||
return
|
||
self.log.insert("end", f"\n{msg}\n")
|
||
self.log.see("end")
|
||
messagebox.showinfo("完成", msg)
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="YOLO 数据集划分")
|
||
parser.add_argument("--cli", action="store_true", help="命令行模式（默认打开界面）")
|
||
parser.add_argument("--src", type=Path, default=ROOT / "datasets" / "phonedata")
|
||
parser.add_argument("--dst", type=Path, default=ROOT / "datasets" / "phonedata")
|
||
parser.add_argument("--train", type=float, default=80, help="训练集百分比，如 80")
|
||
parser.add_argument("--val", type=float, default=20, help="验证集百分比，如 20")
|
||
parser.add_argument("--test", type=float, default=0, help="测试集百分比，如 10")
|
||
parser.add_argument("--seed", type=int, default=42)
|
||
parser.add_argument("--move", action="store_true")
|
||
parser.add_argument("--names", type=str, default="phone")
|
||
parser.add_argument("--yaml", type=Path, default=ROOT / "data.yaml")
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
if args.cli:
|
||
msg = run_split(
|
||
src=args.src if args.src.is_absolute() else ROOT / args.src,
|
||
dst=args.dst if args.dst.is_absolute() else ROOT / args.dst,
|
||
train_pct=args.train,
|
||
val_pct=args.val,
|
||
test_pct=args.test,
|
||
seed=args.seed,
|
||
move=args.move,
|
||
names_arg=args.names,
|
||
yaml_path=args.yaml if args.yaml.is_absolute() else ROOT / args.yaml,
|
||
)
|
||
print(msg)
|
||
return
|
||
SplitApp().mainloop()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|