【代码学习】——简单二分类数据集划分代码学习

# -*- coding: utf-8 -*-

import os
import random
import shutil

#功能：创建文件夹
#传入参数：被创建文件夹的名字
def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)


if __name__ == '__main__':

    random.seed(1)#设置随机数种子，用来和示例程序的代码生成的结果保持一致
	
	#拼接各种指点的路径
    dataset_dir = os.path.join("..", "..", "data", "RMB_data")
    split_dir = os.path.join("..", "..", "data", "rmb_split")
    train_dir = os.path.join(split_dir, "train")
    valid_dir = os.path.join(split_dir, "valid")
    test_dir = os.path.join(split_dir, "test")
	
	#划分训练验证测试的比例
    train_pct = 0.8
    valid_pct = 0.1
    test_pct = 0.1

    for root, dirs, files in os.walk(dataset_dir):
        for sub_dir in dirs:

            imgs = os.listdir(os.path.join(root, sub_dir))
            #这里的filter是用来过滤List中的元素的，lambda用来设置条件，
            imgs = list(filter(lambda x: x.endswith('.jpg'), imgs))
            #吧获取到的图片打乱
            random.shuffle(imgs)
            #记录图片的数量
            img_count = len(imgs)
			
			#划分点，80%的数据给训练，剩下的分别等分给验证和测试
            train_point = int(img_count * train_pct)
            valid_point = int(img_count * (train_pct + valid_pct))
			#在所有的图片中，划分数据集
            for i in range(img_count):
                if i < train_point:
                    out_dir = os.path.join(train_dir, sub_dir)
                elif i < valid_point:
                    out_dir = os.path.join(valid_dir, sub_dir)
                else:
                    out_dir = os.path.join(test_dir, sub_dir)

                makedir(out_dir)#创建上述的文件夹

                target_path = os.path.join(out_dir, imgs[i])
                src_path = os.path.join(dataset_dir, sub_dir, imgs[i])
				#拷贝原来文件夹中的图片放到新的需要划分的文件夹中
                shutil.copy(src_path, target_path)
			#处理的打印信息
            print('Class:{}, train:{}, valid:{}, test:{}'.format(sub_dir, train_point, valid_point-train_point,
                                                                 img_count-valid_point))