随机划分训练集和测试集并保存
将指定的数据集路径输入到函数中,函数将创建对应路径,并复制指定路径下的文件到当前所创建的路径下:
其中划分比例暂时还没有写出自定义,也可以自己添加修改
代码中主要应用了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个类别的图像: