将文件夹内文件按照一定比例分为训练集和测试集

在训练自己模型的时候往往需要自己准备数据集,并将数据集按照一定比例随机分为训练集和测试集,以下是一个简单的脚本,可满足此需要

# coding: utf-8
import os
import random
import shutil	
def moveFile(data_base):
	
	fileDir = data_base+"mark/"  # 原文件夹路径
    xml_dir= data_base+'/xml/'

    train_dir = data_base+'result/mark/'  # 移动到新的文件夹路径
    test_dir = data_base+'result/test/mark/'

    train_xml_dir = data_base+'result/xml/'
    test_xml_dir= data_base+'/result/test/xml/'

    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    if not os.path.exists(test_dir):
        os.makedirs(test_dir)
    if not os.path.exists(test_xml_dir):
        os.makedirs(test_xml_dir)
    if not os.path.exists(train_xml_dir):
        os.makedirs(train_xml_dir)
        
    pathDir = os.listdir(fileDir)  # 取图片的原始路径
    filenumber = len(pathDir)
    rate = 0.3  
    number = int(filenumber * rate)           # 按照rate比例从文件夹中取数据
    sample = random.sample(pathDir, number)  # 随机选取picknumber数量的数据
    #print (sample)
    for name in sample:
        print(name)
        shutil.copy(fileDir + name, test_dir + name)
        item = name[:-4]
        print(xml_dir + item + '.xml')
        print(test_xml_dir+item+'.xml')
        shutil.copy(xml_dir+item+'.xml',test_xml_dir+item+'.xml')

    for filename in os.listdir(fileDir):
        if filename not in os.listdir(test_dir):
            shutil.copy(os.path.join(fileDir,filename),os.path.join(train_dir,filename))
            shutil.copy(xml_dir + filename[:-4]+ '.xml',train_xml_dir + filename[:-4] + '.xml')
    return
if __name__ == '__main__':

    data_base = 文件夹路径
    moveFile(data_base)

你可能感兴趣的:(机器学习笔记)