import os
import random
import shutil

def split_dataset(images_dir, labels_dir, train_ratio=0.8):
    # 创建 train 和 val 文件夹
    train_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'images')
    val_images_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'images')
    train_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/train', 'labels')
    val_labels_dir = os.path.join('/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/val', 'labels')

    os.makedirs(train_images_dir, exist_ok=True)
    os.makedirs(val_images_dir, exist_ok=True)
    os.makedirs(train_labels_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_index = int(len(image_files) * train_ratio)

    # 划分文件
    train_files = image_files[:split_index]
    val_files = image_files[split_index:]

    # 移动图像文件到对应的文件夹
    for file in train_files:
        shutil.move(os.path.join(images_dir, file), os.path.join(train_images_dir, file))
        label_file = file.replace('.jpg', '.txt')  # 假设标签文件与图像文件同名，但扩展名为 .txt
        shutil.move(os.path.join(labels_dir, label_file), os.path.join(train_labels_dir, label_file))

    for file in val_files:
        shutil.move(os.path.join(images_dir, file), os.path.join(val_images_dir, file))
        label_file = file.replace('.jpg', '.txt')  # 假设标签文件与图像文件同名，但扩展名为 .txt
        shutil.move(os.path.join(labels_dir, label_file), os.path.join(val_labels_dir, label_file))

    print(f"数据集已成功划分为训练集和验证集，训练集占比 {train_ratio * 100}%")

# 使用示例
images_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/images'  # 替换为你的 images 文件夹路径
labels_dir = '/home/admin-root/haotian/python哈汽锻8安全帽识别/安全帽头数据集20250428/2200/images/labels'  # 替换为你的 labels 文件夹路径
train_ratio = 0.8  # 训练集占比，默认为 80%

split_dataset(images_dir, labels_dir, train_ratio)