关于图像分类任务中划分数据集，并且生成分类类别的josn字典文件-Toy模板网

这篇具有很好参考价值的文章主要介绍了关于图像分类任务中划分数据集，并且生成分类类别的josn字典文件。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

1. 前言

在做图像分类任务的时候，数据格式是文件夹格式，相同文件夹下存放同一类型的类别

不少网上的数据，没有划分数据集，虽然代码简单，每次重新编写还是颇为麻烦，这里记录一下

如下，有的数据集这样摆放：

关于图像分类任务中划分数据集，并且生成分类类别的josn字典文件,# 关于 classification,分类,数据挖掘,人工智能

可以看出这是个三分类任务，不过没有划分测试集、验证集

代码存放位置：和数据集dataset 同一路径

关于图像分类任务中划分数据集，并且生成分类类别的josn字典文件,# 关于 classification,分类,数据挖掘,人工智能

2. 完整代码

如下：

import random
import os
import shutil
from tqdm import tqdm
import json


def split_data(root, test_rate, flag=True):
    # 待分类数据的当前目录
    classes_directory = [i for i in os.listdir(root) if os.path.isdir(os.path.join(root, i))]

    # 建立生成后的目录，方便拷贝
    for i in classes_directory:
        os.makedirs(os.path.join('./data/train', i))  # 训练集
        os.makedirs(os.path.join('./data/test', i))  # 测试集

    # 是否生成类别的 json 字典文件，默认生成
    if flag:
        class_indices = dict((k, v) for v, k in enumerate(classes_directory))
        json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
        with open('class_indices.json', 'w') as json_file:
            json_file.write(json_str)

    # 遍历每个文件夹下的文件
    for cla in classes_directory:
        cla_path = os.path.join(root, cla)  # 每个文件夹的路径
        images_path = [os.path.join(root, cla, i) for i in os.listdir(cla_path)]

        # 按比例随机采样测试集样本
        test_split_path = random.sample(images_path, k=int(len(images_path) * test_rate))

        # 划分数据
        for i in tqdm(images_path, desc=cla):
            if i in test_split_path:
                shutil.copy(i, os.path.join('./data/test', cla))
            else:
                shutil.copy(i, os.path.join('./data/train', cla))


if __name__ == '__main__':
    rawDataSet = './dataset'  # 原始数据的路径

    if os.path.exists('./data'):  # 如果之前有，那么删除
        shutil.rmtree('./data')

    os.makedirs('./data/train')
    os.makedirs('./data/test')

    # 划分数据
    split_data(root=rawDataSet, test_rate=0.2)

运行代码过程：

关于图像分类任务中划分数据集，并且生成分类类别的josn字典文件,# 关于 classification,分类,数据挖掘,人工智能