Python系列 | 自定义数据集(划分训练集与验证集)

1、源数据集介绍

数据来源于TensorFlow官网:花数据集

原始数据集展示(每个文件夹内为文件名所对应花种类的图像):

在这里插入图片描述
在划分数据集时应注意LICENSE.txt文件的存在。

在这里插入图片描述

2、代码

将数据集按 7:3 的比例划分为训练集和验证集两部分。

若是读者对 os 模块不太熟悉,可见我另一篇博客:Python系列 | os模块常用命令

import os
import random
import shutil


def make_file(path):
    if os.path.exists(path):
        os.removedirs(path)
        os.makedirs(path)
    else:
        os.makedirs(path)


def main():
    path = r'.\dataset'  # 训练集和验证集的存储路径
    par_path = os.path.join(os.getcwd(), 'flower_photos')  # 原始数据集路径
    label_set = [label_i for label_i in os.listdir(par_path) if os.path.isdir(os.path.join(par_path, label_i))]  # 类别标签
    # 训练集和验证集的比例
    rate = 0.7
    # 创建训练集/测试集文件夹
    train_path = os.path.join(path, 'train')  # 训练集路径
    make_file(train_path)
    val_path = os.path.join(path, 'val')  # 验证集路径
    make_file(val_path)

    for label in label_set:
        # 创建训练集/验证集内的花种类文件
        os.mkdir(os.path.join(train_path, label))
        os.mkdir(os.path.join(val_path, label))

        root = os.path.join(par_path, label)
        image_set = os.listdir(root)
        sample_index = random.sample(image_set, k=int(rate * len(image_set)))

        # 分配数据
        for i, j in enumerate(image_set):
            origin_path = os.path.join(root, j)
            tgt_train_path = os.path.join(train_path, label)
            tgt_val_path = os.path.join(val_path, label)
            if j in sample_index:
                shutil.copy(origin_path, tgt_train_path)
            else:
                shutil.copy(origin_path, tgt_val_path)

            if i == len(image_set) - 1:
                print(f'< {label} > processing | {i + 1}/{len(image_set)} |')

    print('Finished !')


if __name__ == '__main__':
    main()

打印结果:

< daisy > processing | 633/633 |
< dandelion > processing | 898/898 |
< roses > processing | 641/641 |
< sunflowers > processing | 699/699 |
< tulips > processing | 799/799 |
Finished !

划分结果:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

  • 3
    点赞
  • 32
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值