import os

import random

trainval_percent = 0.8

train_percent = 0.7

imageSetsDir = './ImageSets'

xmlfilepath = './Annotations/'

if os.path.exists(imageSetsDir):

shutil.rmtree(imageSetsDir, ignore_errors=True)

os.makedirs(imageSetsDir)

if os.path.exists(imageSetsDir + '/Main'):

shutil.rmtree(imageSetsDir + '/Main', ignore_errors=True)

os.makedirs(imageSetsDir + '/Main')

total_xml = os.listdir(xmlfilepath)

num = len(total_xml)

list = range(num)

tv = int(num*trainval_percent)

tr = int(tv*train_percent)

trainval = random.sample(list,tv)

train = random.sample(trainval,tr)

ftrainval = open(imageSetsDir + '/Main/trainval.txt', 'w')

ftest = open(imageSetsDir + '/Main/test.txt', 'w')

ftrain = open(imageSetsDir + '/Main/train.txt', 'w')

fval = open(imageSetsDir + '/Main/val.txt', 'w')

for i in list:

name = total_xml[i][:-4]+'\n'

if i in trainval:

ftrainval.write(name)

if i in train:

ftrain.write(name)

else:

fval.write(name)

else:

ftest.write(name)

ftrainval.close()

ftrain.close()

fval.close()

ftest .close()

返回