| """将已转换的YOLO标注和对应图片划分为train/val集""" |
| import shutil |
| import random |
| import argparse |
| from pathlib import Path |
|
|
| CAMERAS = ["EastLeft", "EastRight", "WestRight", "Westleft"] |
| PERIODS = ["Day", "Night"] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("labeled_images_root", help="Labeled_images目录") |
| parser.add_argument("yolo_labels_root", help="yolo_labels_tmp目录") |
| parser.add_argument("out_root", help="Detection_dataset目录") |
| parser.add_argument("--val-ratio", type=float, default=0.1) |
| parser.add_argument("--seed", type=int, default=42) |
| args = parser.parse_args() |
|
|
| random.seed(args.seed) |
| labeled_root = Path(args.labeled_images_root) |
| labels_root = Path(args.yolo_labels_root) |
| out_root = Path(args.out_root) |
|
|
| train_img = out_root / "images" / "train" |
| val_img = out_root / "images" / "val" |
| train_lbl = out_root / "labels" / "train" |
| val_lbl = out_root / "labels" / "val" |
| for d in [train_img, val_img, train_lbl, val_lbl]: |
| d.mkdir(parents=True, exist_ok=True) |
|
|
| stats = {"train": 0, "val": 0} |
|
|
| for cam in CAMERAS: |
| for period in PERIODS: |
| img_dir = labeled_root / cam / period |
| lbl_dir = labels_root / cam / f"{period}Label" |
| if not img_dir.exists() or not lbl_dir.exists(): |
| print(f"跳过 {cam}/{period}（目录不存在）") |
| continue |
|
|
| imgs = sorted(img_dir.glob("*.jpg")) |
| random.shuffle(imgs) |
| n_val = max(1, int(len(imgs) * args.val_ratio)) |
| val_set = set(img.stem for img in imgs[:n_val]) |
|
|
| for img_path in imgs: |
| lbl_path = lbl_dir / (img_path.stem + ".txt") |
| if not lbl_path.exists(): |
| print(f"警告：找不到标注 {lbl_path}") |
| continue |
| split = "val" if img_path.stem in val_set else "train" |
| shutil.copy2(img_path, out_root / "images" / split / img_path.name) |
| shutil.copy2(lbl_path, out_root / "labels" / split / (img_path.stem + ".txt")) |
| stats[split] += 1 |
|
|
| print(f"划分完成: train={stats['train']}, val={stats['val']}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
