python自动划分分类任务图像数据集，与 pytorch ImageFolder 对接

看了一下网上的划分图像数据集大多数都是存在一个文件里，按比例把这些图片分开。但是遇到多分类问题怎么办？划分后能否直接与pytorch的ImageFolder对接？这里我给出来我自己划分分类任务图像数据集的code，通过我的代码划分完数据集后，就可以直接与ImageFolder来对接训练了。

话不多说直接上代码：

def partition_dataset(label_csv_path, root_path=None, save_path=None, shuffle=False, proportion=0.20):

    """
    根据标签将图片分类，保存到不同文件夹下，文件夹名称为标签名
    :param label_csv_path: 标签存放路径
    :param root_path: 指定标签中图片的根目录
    :param save_path: 保存图片的路径
    :param proportion: 测试集所占比例
    :param shuffle: 是否shuffle
    :return: None
    """
    import random

    label_csv_list = []    # 将标签文件转化为嵌套列表 [['1', '100.PNG'], ... ,['5', '1000.PNG']]
    file_name_list = []
    label_file_dict = {}   # 嵌套列表的字典，key:标签 value:标签对应的所有样本
    label_count_dict = {}  # 每个标签对应的样本数
    label_csv_file = open(label_csv_path)
    csv_reader_line = csv.reader(label_csv_file)

    for line in csv_reader_line:
        label_csv_list.append(line[::-1])
        file_name_list.append(line[::-1][0])

    file_name_list = set(file_name_list)   # 去重，得到标签，可以作为文件名和字典的key值 {1,2,3,4,5,6}

    for item in file_name_list:
        label_file_dict.setdefault(item, [])
        label_count_dict.setdefault(item, 1)

    for item in label_csv_list:
        label_file_dict[item[0]].append(item[1])

    for item in file_name_list:
        label_count_dict[item] = len(label_file_dict[item])

    print(label_count_dict)

    # 创建数据集的目录结构
    train_path = os.path.join(save_path, "train")
    val_path = os.path.join(save_path, "val")
    if not os.path.exists(train_path):  # 如果路径不存在
        os.makedirs(train_path)
    if not os.path.exists(val_path):
        os.makedirs(val_path)

    # 在train和val目录下分别创建与标签对应的文件夹
    for item in file_name_list:
        if not os.path.exists(os.path.join(train_path, item)):
            os.makedirs(os.path.join(train_path, item))
        if not os.path.exists(os.path.join(val_path, item)):
            os.makedirs(os.path.join(val_path, item))

    # 是否shuffle
    if shuffle is True:
        for k in label_file_dict:
            random.shuffle(label_file_dict[k])

    # 开始划分
    for k in label_file_dict:
        for i in range(0, int((len(label_file_dict[k])+1)*(1.0-proportion))):
            shutil.copy(os.path.join(root_path, label_file_dict[k][i]),
                        os.path.join(train_path + "\\" + k + "\\", label_file_dict[k][i]))
            print("\033[31m copy: %s to %s \033[0m" % (
                label_file_dict[k][i], os.path.join(train_path + "\\" + k + "\\", label_file_dict[k][i])))
        for i in range(int((len(label_file_dict[k])+1)*(1.0-proportion)), len(label_file_dict[k])):
            shutil.copy(os.path.join(root_path, label_file_dict[k][i]),
                        os.path.join(val_path + "\\" + k + "\\", label_file_dict[k][i]))
            print("\033[33m copy: %s to %s \033[0m" % (
                label_file_dict[k][i], os.path.join(val_path + "\\" + k + "\\", label_file_dict[k][i])))

该方法的各种参数以及具体实现过程已经注释的很清楚了，这里我再把我的标签放上来，如果你的标签和我的格式一样，那么应该只需要修改少量的代码就可以成功运行了：

标签就长这个样子，是一个csv，没有head，左侧是图片文件名，右侧是对应的标签。

现在调用 partition_dataset 方法，进行数据集划分：

    partition_dataset(r"E:\complabel.csv",
                      r"E:\dataset\image"+"\\",
                      r"E:\comp_dataset",
                      shuffle=True,
                      proportion=0.3)

划分完成后，我们看一下划分的结果：

comp_dataset文件夹中的内容：

 train文件夹中的内容：

val文件夹中的内容：

可以看到数据集已经按标签划分为训练集和验证集了。

那我们验证一下第一类的图像是否全部被划分，并且是按照我们给定的百分比30%划分的：

首先打开train/1 文件夹，该文件夹下存放了所有标签为1的训练集图片：

 trian/1 文件夹下一共74个sample

val/1文件夹下一共31个sample：

我的标签中的第一类一共有105个sample

74+31=105   74/105≈0.704，说明划分没有问题，那么我们可以直接使用pytorch的imageFolder读取数据集进行训练了：

train_dataset = datasets.ImageFolder(root=r"E:\dataset\comp_dataset" + "\\train",
                                     transform=data_transform["train"])


train_num = len(train_dataset)

validate_dataset = datasets.ImageFolder(root=r"E:\dataset\comp_dataset" + "\\val",
                                        transform=data_transform["val"])

使用ImageFolde的好处是可以直观的看到自己数据集样本分布的情况，尤其是对一个分类任务而言，样本的均衡性是至关重要的，当然另个一好处就是调用很方便，不需要去继承Dataset类，重写读数据方法了。

最后说明一下，每个人的标签或者数据集图片格式不一样，需要自己对代码进行微调。