python按设定比例划分训练集和测试集代码
在利用深度神经网络模型进行图像分类时,数据集的处理至关重要,为了训练和评价模型的好坏,一般需要将数据集划分为训练集和测试集。
这里以遥感图像场景分类数据集NWPU-RESISC45为例,按照预先设置好的比例,进行训练集和测试集划分,Python代码如下:
1 # *_*coding: utf-8 *_* 2 # Author --LiMing-- 3 4 import os 5 import random 6 import shutil 7 import time 8 9 def copyFile(fileDir, class_name): 10 image_list = os.listdir(fileDir) # 获取图片的原始路径 11 image_number = len(image_list) 12 13 train_number = int(image_number * train_rate) 14 train_sample = random.sample(image_list, train_number) # 从image_list中随机获取0.8比例的图像. 15 test_sample = list(set(image_list) - set(train_sample)) 16 sample = [train_sample, test_sample] 17 18 # 复制图像到目标文件夹 19 for k in range(len(save_dir)): 20 if os.path.isdir(save_dir[k] + class_name): 21 for name in sample[k]: 22 shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k] + class_name+'/', name)) 23 else: 24 os.makedirs(save_dir[k] + class_name) 25 for name in sample[k]: 26 shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k] + class_name+'/', name)) 27 28 if __name__ == '__main__': 29 time_start = time.time() 30 31 # 原始数据集路径 32 origion_path = '/home/room/lm_other/NWPU-RESISC45/' 33 34 # 保存路径 35 save_train_dir = '/home/room/lm_other/RS_45/2_8/train/' 36 save_test_dir = '/home/room/lm_other/RS_45/2_8/test/' 37 save_dir = [save_train_dir, save_test_dir] 38 39 # 训练集比例 40 train_rate = 0.2 41 42 # 数据集类别及数量 43 file_list = os.listdir(origion_path) 44 num_classes = len(file_list) 45 46 for i in range(num_classes): 47 class_name = file_list[i] 48 image_Dir = os.path.join(origion_path, class_name) 49 copyFile(image_Dir, class_name) 50 print('%s划分完毕!' % class_name) 51 52 time_end = time.time() 53 print('---------------') 54 print('训练集和测试集划分共耗时%s!' % (time_end - time_start)