import os
from glob import glob
from sklearn.model_selection import train_test_split


saved_path = "/home/wxd/CP/train/VOCdevkit/VOC2007/"  # 保存路径

# 2.创建要求文件夹
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")

# 3.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
# ftrainval = open(txtsavepath + '/trainval.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fval = open(txtsavepath + '/val.txt', 'w')

total_files = glob(saved_path + "Annotations/*.xml")

total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]

# for file in total_files:
#     ftrainval.write(file + "\n")

# split
train_files, val_files = train_test_split(total_files, test_size=0.15, random_state=42)

# train
for file in train_files:
    ftrain.write(file + "\n")
# val
for file in val_files:
    fval.write(file + "\n")

# ftrainval.close()
ftrain.close()
fval.close()