YOLO系列笔记（十三）——数据集文件夹拆分

前言

在学习和训练涉及YOLO等神经网络的过程中，我们通常需要对收集到的大型数据集进行合理的拆分，标准的做法是将其分为训练集（train）、测试集（test）和验证集（validation）。手动对大规模的图片集进行分类不仅耗时而且容易出错，因此我编写了一个Python脚本，旨在自动化这一繁琐的过程。该脚本能够将指定文件夹中的图片按照8:1:1的比例自动分配到train、test和validation三个文件夹中，使用的是Python的os和shutil库来处理文件操作。

PS. 如果觉得该笔记对您有用的话，可以点个小小的赞吗，或者点赞收藏关注一键三连ヾ(◍’౪`◍) ~ 谢谢！！

代码

import os
import shutil
import random

# 设置源文件夹和目标文件夹路径
source_folder = './source'
train_folder = './train'
test_folder = './test'
validation_folder = './validation'

# 如果目标文件夹不存在，则创建它们
os.makedirs(train_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
os.makedirs(validation_folder, exist_ok=True)

# 读取所有图片文件
images = [f for f in os.listdir(source_folder) if f.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]

# 打乱文件顺序以随机分配
random.shuffle(images)

# 计算分配比例
total_count = len(images)
train_count = int(total_count * 0.8)
test_count = int(total_count * 0.1)
validation_count = total_count - train_count - test_count  # 保证总数正确

# 分配文件到train, test, validation文件夹
def move_files(files, destination):
    for f in files:
        shutil.move(os.path.join(source_folder, f), os.path.join(destination, f))

move_files(images[:train_count], train_folder)
move_files(images[train_count:train_count + test_count], test_folder)
move_files(images[train_count + test_count:], validation_folder)

print(f"Files distributed: {train_count} to train, {test_count} to test, {validation_count} to validation.")

说明

这个Python脚本的主要目的是将一个包含大量图片的源文件夹(source_folder)中的图片按照8:1:1的比例随机分配到三个不同的文件夹(train_folder, test_folder, validation_folder)中，用于机器学习或深度学习模型的训练、测试和验证。下面是对脚本中每个部分的详细解释：

导入模块

import os
import shutil
import random

• os: 用于处理文件和目录，如读取文件列表、创建目录等。
• shutil: 用于执行文件操作，如复制和移动文件。
• random: 用于生成随机数，这里用于打乱文件列表，确保分配的随机性。

设置文件夹路径

source_folder = './source'
train_folder = './train'
test_folder = './test'
validation_folder = './validation'

这些变量定义了源文件夹和目标文件夹的路径。这些路径可以根据实际情况进行调整。

创建目标文件夹

os.makedirs(train_folder, exist_ok=True)
os.makedirs(test_folder, exist_ok=True)
os.makedirs(validation_folder, exist_ok=True)

这些语句确保目标文件夹存在。如果目标文件夹不存在，os.makedirs 将创建它们。参数exist_ok=True表示如果文件夹已存在，不会抛出错误。

读取和准备图片列表

images = [f for f in os.listdir(source_folder) if f.endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]

这行代码读取source_folder中所有的图片文件，并将它们存储在列表images中。列表推导用于过滤出只包含图片的文件，基于文件扩展名进行筛选。os.listdir(source_folder) 这个函数是用来列出指定目录下的所有文件和目录名的。

打乱图片顺序

random.shuffle(images)

使用random.shuffle函数随机打乱images列表中元素的顺序，这是为了确保后续分配到训练、测试和验证集的随机性。

计算分配比例

total_count = len(images)
train_count = int(total_count * 0.8)
test_count = int(total_count * 0.1)
validation_count = total_count - train_count - test_count

根据总图片数，计算出应分配到训练集、测试集和验证集的图片数量。按照8:1:1的比例，训练集占80%，测试和验证集各占10%。

定义和使用移动文件函数

def move_files(files, destination):
    for f in files:
        shutil.move(os.path.join(source_folder, f), os.path.join(destination, f))

move_files(images[:train_count], train_folder)
move_files(images[train_count:train_count + test_count], test_folder)
move_files(images[train_count + test_count:], validation_folder)

• move_files 函数接受文件列表和目标文件夹作为参数，然后逐个将文件从源文件夹移动到指定的目标文件夹。
• 这里使用切片操作分割images列表，根据之前计算的数量，将文件分配到相应的文件夹。

最后打印分配结果

print(f"Files distributed: {train_count} to train, {test_count} to test, {validation_count} to validation.")

在所有操作完成后，打印出每个集合中文件的数量，以确认文件已正确分配。

总结

此脚本极大地简化了将大规模图片数据集分配到训练集、测试集和验证集的过程，不仅提高了效率，还避免了手动分配时可能出现的错误。这对于快速准备机器学习和深度学习模型的训练数据是非常有用的。