随机划分训练集和测试集并保存

 1 import os
 2 import shutil
 3 from random import randrange
 4 
 5 img_root0 = ‘E:/****/PatternNet‘       #数据集目录****这里省略了路径，自行修改
 6 
 7 #随机生成训练集和测试集，比例为7：1
 8 def ImgSet_save(images_path=img_root0):
 9     save_trainpath = ‘./data/train_images‘
10     sava_testpath = ‘./data/test_images‘
11     modnum = randrange(8)       #首先随机生成一个随机数
12     #开始划分训练集和测试集并保存
13     for temp in os.listdir(images_path):
14         image_dir = images_path +‘/‘+ temp
15         i = 0
16         print(image_dir)　　　　　　#可以输出查看下该分类下的数据集是什么
17         for imagename in os.listdir(image_dir):　　　　#获取具体的文件名称
18             src_image = image_dir + ‘/‘+ imagename
19             #print(src_image)
20             if i%8 == 0:                            #将数据集划分为7：1，每次处理完八张图片后变化一次随机数
21                 modnum = randrange(8)
22             i+=1
23             if i %8==modnum:                                 #划分为测试集
24                 sava_path = sava_testpath+‘/‘+temp+‘/‘
25                 if not os.path.exists(sava_path):
26                     os.makedirs(sava_path)
27                 shutil.copy(src=src_image,dst=sava_path) #复制文件
28             else:                                        #划分为训练集
29                 save_path = save_trainpath+‘/‘+temp+‘/‘
30                 if not os.path.exists(save_path):
31                     os.makedirs(save_path)
32                 shutil.copy(src=src_image,dst=save_path)  #
33 
34 if __name__ == ‘__main__‘:
35     ImgSet_save()
36     print("数据集划分完成，划分为7：1的训练集和测试集")