split_train_val.py

import os
import shutil
import random
from pathlib import Path

# 配置路径
root_dir = 'Vehicle_Model_Segmentation_Dataset/car_images'
train_dir = os.path.join(root_dir, 'train')
val_dir = os.path.join(root_dir, 'val')

# 创建train和val目录
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)

# 获取所有车型类别
categories = [d for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d)) and d not in ['train', 'val']]

split_ratio = 0.8  # 训练集比例

for category in categories:
    cat_dir = os.path.join(root_dir, category)
    images = [f for f in os.listdir(cat_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    random.shuffle(images)
    n_train = int(len(images) * split_ratio)
    train_imgs = images[:n_train]
    val_imgs = images[n_train:]

    # 创建类别子文件夹
    os.makedirs(os.path.join(train_dir, category), exist_ok=True)
    os.makedirs(os.path.join(val_dir, category), exist_ok=True)

    # 拷贝图片到train
    for img in train_imgs:
        src = os.path.join(cat_dir, img)
        dst = os.path.join(train_dir, category, img)
        shutil.copy2(src, dst)
    # 拷贝图片到val
    for img in val_imgs:
        src = os.path.join(cat_dir, img)
        dst = os.path.join(val_dir, category, img)
        shutil.copy2(src, dst)

print('数据集已按8:2比例划分为train和val。')