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