使用 Python进行数据集分割：简洁完美的数据集划分脚本

在机器学习的世界中，正确准备数据集与您选择的算法一样重要。将数据集划分为训练集、验证集和测试集的过程是构建稳健模型的基本步骤。本文深入研究了专门为此目的设计的 Python 脚本，展示了如何有效地分割机器学习项目的数据集，特别是 YOLOv5 PyTorch 模型。

---

分解脚本

该脚本首先导入必要的库：**os用于操作系统交互、shutil用于文件操作、zipfile用于处理 zip 存档以及train_test_split用于sklearn.model_selection**划分数据集。

目录结构：
它为数据集、图像和标签设置目录。这一步就像在建筑物中布置地基和房间，确保一切都在正确的位置。

数据备份：
该脚本在 zip 文件中创建数据集的备份。将此视为在对重要文件进行任何更改之前对其进行复印。

文件列表：
它编译图像文件名列表，为分割过程做好准备。这类似于图书馆员在将书籍放入特定部分之前对其进行编目。

数据集拆分：
根据预定义的比例将数据集分为训练集、验证集和测试集。此步骤对于训练模型并评估其在不同条件下的性能至关重要。

组织文件：
对于每个子集（训练、验证、测试），脚本都会创建目录并相应地移动文件。这种组织类似于厨师将不同菜肴的食材放入不同的容器中。

清理过程：
最后，脚本删除所有空目录并输出备份和新创建目录的路径。这就像做饭后清理厨房，以备下次使用一样。

import os
import shutil
import zipfile
from sklearn.model_selection import train_test_split

# 定义目录路径
dataset_directory = r'D:\Desktop\3.v2i.yolov5pytorch\train\dataset_backup'
images_directory = os.path.join(dataset_directory, 'images')
labels_directory = os.path.join(dataset_directory, 'labels')

# 定义备份 zip 文件路径
backup_zip_path = os.path.join(dataset_directory, 'dataset_backup.zip')

# 将整个数据集目录压缩备份
with zipfile.ZipFile(backup_zip_path, 'w') as backup_zip:
    for folder in [images_directory, labels_directory]:
        for root, dirs, files in os.walk(folder):
            for file in files:
                file_path = os.path.join(root, file)
                backup_zip.write(file_path, os.path.relpath(file_path, dataset_directory))

# 所有图像文件名（不带扩展名）的列表
image_filenames = [os.path.splitext(filename)[0] for filename in os.listdir(images_directory) if filename.endswith('.jpg')]

# 定义拆分比例
train_ratio = 0.80
validation_ratio = 0.19
test_ratio = 0.01

# 执行拆分
train_filenames, test_filenames = train_test_split(image_filenames, test_size=1 - train_ratio)
validation_filenames, test_filenames = train_test_split(test_filenames, test_size=test_ratio/(test_ratio + validation_ratio))

# 创建目录并移动文件的函数
def create_and_move_files(file_list, source_folder, destination_folder, file_extension):
    os.makedirs(destination_folder, exist_ok=True)
    for filename in file_list:
        shutil.move(os.path.join(source_folder, filename + file_extension),
                    os.path.join(destination_folder, filename + file_extension))

# 为每个集创建目录并移动文件
sets_directories = {
    'train': os.path.join(dataset_directory, 'train'),
    'val': os.path.join(dataset_directory, 'val'),
    'test': os.path.join(dataset_directory, 'test')
}

for set_name, file_list in zip(['train', 'val', 'test'], [train_filenames, validation_filenames, test_filenames]):
    images_set_directory = os.path.join(sets_directories[set_name], 'images')
    labels_set_directory = os.path.join(sets_directories[set_name], 'labels')
    create_and_move_files(file_list, images_directory, images_set_directory, '.jpg')
    create_and_move_files(file_list, labels_directory, labels_set_directory, '.txt')

# 删除空的images和labels文件夹
for folder in [images_directory, labels_directory]:
    if os.path.exists(folder) and not os.listdir(folder):
        os.rmdir(folder)

# 返回备份文件路径和已创建目录的路径
print(backup_zip_path, sets_directories)

进一步探索：

有兴趣深入了解吗？以下是一些建议：

• 该脚本如何适应不同类型的数据集或机器学习模型？
• 不同的训练、验证和测试比率对模型性能有什么影响？
• 如何将该脚本集成到机器学习模型的持续训练管道中？

该脚本体现了 Python 在简化和优化机器学习工作流程方面的强大功能。通过自动化数据集分割过程，不仅可以节省时间，还可以提高机器学习项目的准确性和可靠性。

---

数据集预处理脚本专栏

1. 使用 Python进行数据集分割：简洁完美的数据集划分脚本
2. 深度学习中的自动化标签转换：对数据集所有标签做映射转换
3. 识别并处理数据集中不配对图像和标签文件