python 自动划分训练集和测试集

前言

在进行深度学习的模型训练时，我们通常需要将数据进行划分，划分成训练集和测试集，若数据集太大，数据划分花费的时间太多！！！
不多说，上代码(python代码)

代码

# *_*coding: utf-8 *_*

import os
import random
import shutil
import time

def copyFile(fileDir,origion_path1,class_name):
    name = class_name
    path = origion_path1
    image_list = os.listdir(fileDir) # 获取图片的原始路径
    image_number = len(image_list)
    train_number = int(image_number * train_rate)
    train_sample = random.sample(image_list, train_number) # 从image_list中随机获取0.75比例的图像.
    test_sample = list(set(image_list) - set(train_sample))
    sample = [train_sample, test_sample]
    
    # 复制图像到目标文件夹
    for k in range(len(save_dir)):
            if os.path.isdir(save_dir[k]) and os.path.isdir(save_dir1[k]):
                for name in sample[k]:
                    name1 = name.split(".")[0] + '.xml'
                    shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k], name))
                    shutil.copy(os.path.join(path, name1), os.path.join(save_dir1[k], name1))
            else:
                os.makedirs(save_dir[k])
                os.makedirs(save_dir1[k])
                for name in sample[k]:
                    name1 = name.split(".")[0] + '.xml'
                    shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k], name))
                    shutil.copy(os.path.join(path, name1), os.path.join(save_dir1[k], name1))

if __name__ == '__main__':
    time_start = time.time()

    # 原始数据集路径
    origion_path = './JPEGImages/'
    origion_path1 = './Annotations/'

    # 保存路径
    save_train_dir = './train/JPEGImages/'
    save_test_dir = './test/JPEGImages/'
    save_train_dir1 = './train/Annotations/'
    save_test_dir1 = './test/Annotations/'
    save_dir = [save_train_dir, save_test_dir]
    save_dir1 = [save_train_dir1, save_test_dir1]

    # 训练集比例
    train_rate = 0.75

    # 数据集类别及数量
    file_list = os.listdir(origion_path)
    num_classes = len(file_list)
    for i in range(num_classes):
        class_name = file_list[i]
    copyFile(origion_path,origion_path1,class_name)
    print('划分完毕!')
    time_end = time.time()
    print('---------------')
    print('训练集和测试集划分共耗时%s!' % (time_end - time_start))

1.需要修改的地方
1)origion_path：图片路径
2)origion_path1：xml文件路径
3)train_rate:训练集比例

2.执行文件deal.py后生成
train-img：训练集图片数据
train-xml：训练集xml数据
test-img：测试集图片数据
test-xml：测试及xml数据

3.train_rate可以根据实际情况进行调整，一般train:test是3：1

注：每次划分数据都是随机的，每次执行时将之前划分好的数据保存或者重命名，不然会重复写入到4个文件夹中

总结

代码还是比较简洁的，有什么问题可以留言。