Python实现随机划分图片数据集的示例代码

下面是Python实现随机划分图片数据集的完整攻略。

1. 问题背景

在机器学习领域中，将数据集按照一定的比例划分为训练集、验证集和测试集是一项基础任务。

数据集的划分对模型训练和评估的结果都有着至关重要的影响，因此我们需要仔细地进行划分。

在本次攻略中，我们将采用Python语言实现对图片数据集的随机划分，并且按照规定的比例划分为训练集和验证集，代码也可以用来划分为训练集、验证集和测试集。

2. 需求分析

任务的输入是一个图片文件夹，输出是划分好的训练集和验证集。

我们需要定义以下参数：

• 数据集路径
• 训练集占比
• 验证集占比
• 是否打乱数据集
• 是否需要输出到文件

3. 实现流程

实现过程中，我们需要按照以下步骤进行：

3.1 读取文件夹中的所有文件

import os
import random
from shutil import copyfile

def get_file_list(path):
    file_list = []
    for root, dirs, files in os.walk(path):
        for file in files:
            file_list.append(os.path.join(root, file))
    return file_list

3.2 打乱文件列表

在获取到文件列表之后，我们可以使用Python内置的random库中的shuffle方法打乱文件列表。

def shuffle_list(file_list):
    random.shuffle(file_list)
    return file_list

3.3 划分数据集

在打乱文件列表之后，我们可以按照指定的训练集和验证集占比，将文件列表划分为训练集和验证集。

def split_data(file_list, training_ratio, validation_ratio):
    training_list = file_list[:int(len(file_list)*training_ratio)]
    validation_list = file_list[int(len(file_list)*training_ratio):int(len(file_list)*(training_ratio+validation_ratio))]
    return training_list, validation_list

3.4 复制文件到指定目录

在划分好训练集和验证集之后，我们需要复制文件到指定的目录。

def copy_files(file_list, output_path):
    for file in file_list:
        _, filename = os.path.split(file)
        output_file = os.path.join(output_path, filename)
        copyfile(file, output_file)

3.5 完整代码

将以上步骤整合起来，形成完整的代码如下：

import os
import random
from shutil import copyfile

def get_file_list(path):
    file_list = []
    for root, dirs, files in os.walk(path):
        for file in files:
            file_list.append(os.path.join(root, file))
    return file_list

def shuffle_list(file_list):
    random.shuffle(file_list)
    return file_list

def split_data(file_list, training_ratio, validation_ratio):
    training_list = file_list[:int(len(file_list)*training_ratio)]
    validation_list = file_list[int(len(file_list)*training_ratio):int(len(file_list)*(training_ratio+validation_ratio))]
    return training_list, validation_list

def copy_files(file_list, output_path):
    for file in file_list:
        _, filename = os.path.split(file)
        output_file = os.path.join(output_path, filename)
        copyfile(file, output_file)

def main(input_path, output_path, training_ratio, validation_ratio, shuffle=True, to_file=True):
    # 获取文件列表
    file_list = get_file_list(input_path)

    # 打乱文件列表
    if shuffle:
        file_list = shuffle_list(file_list)

    # 划分训练集和验证集
    training_list, validation_list = split_data(file_list, training_ratio, validation_ratio)

    # 输出到文件
    if to_file:
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        copy_files(training_list, os.path.join(output_path, "training"))
        copy_files(validation_list, os.path.join(output_path, "validation"))
    else:
        return training_list, validation_list


if __name__ == "__main__":
    input_path = "./data"
    output_path = "./output"
    training_ratio = 0.8
    validation_ratio = 0.2

    main(input_path, output_path, training_ratio, validation_ratio)

4. 示例说明

4.1 示例一

以下示例将图片数据集随机划分为训练集和测试集，训练集占比80%。

if __name__ == "__main__":
    input_path = "./data"
    output_path = "./output"
    training_ratio = 0.8
    validation_ratio = 0

    main(input_path, output_path, training_ratio, validation_ratio)

4.2 示例二

以下示例将图片数据集随机划分为训练集和验证集，训练集占比70%，验证集占比30%。

if __name__ == "__main__":
    input_path = "./data"
    output_path = "./output"
    training_ratio = 0.7
    validation_ratio = 0.3

    main(input_path, output_path, training_ratio, validation_ratio)

以上就是Python实现随机划分图片数据集的完整攻略，希望对你有所帮助。