随机分配数据集和验证集

找了很多代码 不好用调试没成功 最后逼得小白自己写了一个代码如下 需要自取

import os
import random
import shutil
'''
数据集划分： 图片在一个文件夹      需要四个文件夹 注意：图片的文件里只能有数据图片
          标签在另一个文件夹
          训练集在一个文件夹
          测试集在一个文件夹
          最后生成 train.txt train_xml.txt 文件
                  test.txt test_xml.txt 文件
            
'''
image_path = r"C:\Users\T\Desktop\rail\ceshi"  #图片的文档
image_xml_path = r"C:\Users\T\Desktop\rail\ceshi_xml"
train_path = r"C:\Users\T\Desktop\rail\ceshi_jieguo\train"
test_path = r"C:\Users\T\Desktop\rail\ceshi_jieguo\test"
train_rate = 0.8 # 训练0.8    测试0.2 划分
image_list  = os.listdir(image_path) #获取文件夹下所有内容
train_num = int(len(image_list)*train_rate) #多少张训练图片
train_list = random.sample(image_list,train_num)
test_list = list(set(image_list)-set(train_list)) #生成测试集
'''
这个shutil.copy 原地址和目的地址必须是字符串
'''
for i in test_list:
    #存入验证集图片
    store_path = test_path+'/'+str(i)
    original_path = image_path +'/'+str(i)
    shutil.copy(original_path,store_path)
    #存入验证集xml
    string_train = str(i)
    string_train = string_train.split('.', 1)[0]  #
    store_path = test_path+'/'+string_train+".xml"
    original_path = image_xml_path +'/'+string_train+".xml"
    shutil.copy(original_path, store_path)

for i in train_list:
    #存入训练集图片
    store_path = train_path+'/'+str(i)
    original_path = image_path +'/'+str(i)
    shutil.copy(original_path,store_path)
    #存入训练集xml
    string_train = str(i)
    string_train = string_train.split('.', 1)[0]  #
    store_path = train_path+'/'+string_train+".xml"
    original_path = image_xml_path +'/'+string_train+".xml"
    shutil.copy(original_path, store_path)

# print(train_list)
# print(len(image_list))
# fl=open(train_path+'\\'+"train.txt",'w')
# for i in train_list:
#     fl.write(i)
#     fl.write("\n")
# fl.close()
# print("训练集数据写入完毕")
# fl=open(train_path+'\\'+"train_xml.txt",'w')
# for i in train_list:
#     fl.write(i)
#     fl.write("\n")
# fl.close()
# print("训练集数据xml写入完毕")
#
# fl=open(train_path+'\\'+"text.txt",'w')
# for i in train_list:
#     fl.write(i)
#     fl.write("\n")
# fl.close()
# print("测试集数据写入完毕")
# fl=open(train_path+'\\'+"text_xml.txt",'w')
# for i in train_list:
#     fl.write(i)
#     fl.write("\n")
# fl.close()
# print("测试集数据xml写入完毕")