使用 Python进行数据集分割：简洁完美的数据集划分脚本 |
您所在的位置：网站首页 › 数据集分布Python › 使用 Python进行数据集分割：简洁完美的数据集划分脚本 |
使用 Python进行数据集分割：简洁完美的数据集划分脚本
|
该脚本首先导入必要的库：**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. 识别并处理数据集中不配对图像和标签文件 |
在机器学习的世界中，正确准备数据集与您选择的算法一样重要。将数据集划分为训练集、验证集和测试集的过程是构建稳健模型的基本步骤。本文深入研究了专门为此目的设计的 Python 脚本，展示了如何有效地分割机器学习项目的数据集，特别是 YOLOv5 PyTorch 模型。【本文地址】
今日新闻 |
推荐新闻 |