File size: 2,273 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """将已转换的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()
|