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

  将指定的数据集路径输入到函数中,函数将创建对应路径,并复制指定路径下的文件到当前所创建的路径下:

  其中划分比例暂时还没有写出自定义,也可以自己添加修改

 代码中主要应用了os下的lstdir函数和shutil下的函数,具体代码可参考如下

 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的训练集和测试集")

我这PatternNet目录下是38个类别的图像:

原文地址:https://www.cnblogs.com/2020zxc/p/14631703.html