import os
import random
import shutil

# 设置数据集路径
images_dir = 'images'
labels_dir = 'labels'

# 创建train和val的子文件夹
train_images_dir = 'train/images'
train_labels_dir = 'train/labels'
val_images_dir = 'val/images'
val_labels_dir = 'val/labels'

os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(val_images_dir, exist_ok=True)
os.makedirs(val_labels_dir, exist_ok=True)

# 获取所有的图片文件名
image_files = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]

# 打乱文件列表
random.shuffle(image_files)

# 定义划分比例
split_ratio = 0.8  # 80% train, 20% val

# 计算分割点
split_point = int(len(image_files) * split_ratio)

# 分割数据集
train_files = image_files[:split_point]
val_files = image_files[split_point:]

# 复制图片和标签文件到各自的文件夹
for file in train_files:
    shutil.copy(os.path.join(images_dir, file), os.path.join(train_images_dir, file))
    label_file = file.split('.')[0] + '.txt'  # 假设标签文件名为图片文件名加上.txt
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(train_labels_dir, label_file))

for file in val_files:
    shutil.copy(os.path.join(images_dir, file), os.path.join(val_images_dir, file))
    label_file = file.split('.')[0] + '.txt'  # 假设标签文件名为图片文件名加上.txt
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(val_labels_dir, label_file))

print("Data split and copied successfully.")