import os
import random
import shutil
from sklearn.model_selection import train_test_split

# 定义路径
data_dir = 'D:\python\PycharmProjects\data_20250223\lcnn新T型板十字板增强后（复件）'  # 替换为你的数据文件夹路径
output_dir = 'D:\python\PycharmProjects\data_20250223\lcnn_20250223'  # 替换为你想要保存输出的文件夹路径

# 创建输出目录结构
images_train_dir = os.path.join(output_dir, 'images', 'train')
images_val_dir = os.path.join(output_dir, 'images', 'val')
labels_train_dir = os.path.join(output_dir, 'labels', 'train')
labels_val_dir = os.path.join(output_dir, 'labels', 'val')

os.makedirs(images_train_dir, exist_ok=True)
os.makedirs(images_val_dir, exist_ok=True)
os.makedirs(labels_train_dir, exist_ok=True)
os.makedirs(labels_val_dir, exist_ok=True)

# 获取所有图片文件名和对应的json文件名
image_files = [f for f in os.listdir(data_dir) if f.endswith('.jpg')]
json_files = {f.replace('.json', ''): f for f in os.listdir(data_dir) if f.endswith('.json')}

# 提取图片名称（不包含扩展名）以便匹配json文件
image_names = [os.path.splitext(f)[0] for f in image_files]

# 按照9:1的比例划分数据集
train_names, val_names = train_test_split(image_names, test_size=0.1, random_state=42)

# 复制文件到相应目录
for name in train_names:
    image_file = name + '.jpg'
    json_file = json_files[name]

    # 复制图片文件
    shutil.copy(os.path.join(data_dir, image_file), os.path.join(images_train_dir, image_file))

    # 复制json文件
    shutil.copy(os.path.join(data_dir, json_file), os.path.join(labels_train_dir, json_file))

for name in val_names:
    image_file = name + '.jpg'
    json_file = json_files[name]

    # 复制图片文件
    shutil.copy(os.path.join(data_dir, image_file), os.path.join(images_val_dir, image_file))

    # 复制json文件
    shutil.copy(os.path.join(data_dir, json_file), os.path.join(labels_val_dir, json_file))

# 示例调用
if __name__ == "__main__":
    # 确保输出目录和子目录存在
    os.makedirs(images_train_dir, exist_ok=True)
    os.makedirs(images_val_dir, exist_ok=True)
    os.makedirs(labels_train_dir, exist_ok=True)
    os.makedirs(labels_val_dir, exist_ok=True)

    # 执行划分和文件复制
    train_names, val_names = train_test_split(image_names, test_size=0.1, random_state=42)

    for name in train_names:
        image_file = name + '.jpg'
        json_file = json_files[name]

        # 复制图片文件
        shutil.copy(os.path.join(data_dir, image_file), os.path.join(images_train_dir, image_file))

        # 复制json文件
        shutil.copy(os.path.join(data_dir, json_file), os.path.join(labels_train_dir, json_file))

    for name in val_names:
        image_file = name + '.jpg'
        json_file = json_files[name]

        # 复制图片文件
        shutil.copy(os.path.join(data_dir, image_file), os.path.join(images_val_dir, image_file))

        # 复制json文件
        shutil.copy(os.path.join(data_dir, json_file), os.path.join(labels_val_dir, json_file))