数据集划分训练验证测试脚本

VisionX Lab

于 2022-03-04 16:58:07 发布

阅读量263

点赞数

CC 4.0 BY-SA版权

文章标签：深度学习 pytorch 机器学习

本文链接：https://blog.youkuaiyun.com/hasque2019/article/details/123280743

该代码用于将图片数据集按照指定比例划分为训练集、验证集和测试集。通过遍历指定目录，对每个子目录下的图片进行随机打乱并分配到相应目录。主要应用于机器学习或深度学习模型的训练过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import os
import glob
import random
import shutil

dataset_dir = './data/'
train_dir = './images/train/'
valid_dir = './images/valid/'
test_dir = './images/test/'

train_per = 0.8
valid_per = 0.1
test_per = 0.1


def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)


if __name__ == '__main__':

    for root, dirs, files in os.walk(dataset_dir):
        for sDir in dirs:
            imgs_list = glob.glob(os.path.join(root, sDir)+'/*.jpg')
            random.seed(666)
            random.shuffle(imgs_list)
            imgs_num = len(imgs_list)

            train_point = int(imgs_num * train_per)
            valid_point = int(imgs_num * (train_per + valid_per))

            for i in range(imgs_num):
                if i < train_point:
                    out_dir = train_dir + sDir + '/'
                elif i < valid_point:
                    out_dir = valid_dir + sDir + '/'
                else:
                    out_dir = test_dir + sDir + '/'

                makedir(out_dir)
                out_path = out_dir + os.path.split(imgs_list[i])[-1]
                shutil.copy(imgs_list[i], out_path)

            print('Class:{}, train:{}, valid:{}, test:{}'.format(sDir, train_point, valid_point-train_point, imgs_num-valid_point))