YOLO系列之数据集划分

YOLO系列之数据集划分（以YOLOv8为例）

前言

本博客是在我的另一篇博客labelme标注的json格式与yolo格式的相互转换 的基础上进行的，有需要可以参考

YOLO系列数据集格式

在数据集中，图像文件存放于images目录下，对应的YOLO格式的txt标签文件要存放于labels目录下，而且要在images和labels目录下分别建立train训练集、valid验证集、test
测试集（可选，但强烈建议要有！）子目录，如下展示：

dataset
	- images
		- train
			- 1.jpg
			- 5.jpg # 因为要随机划分，所以数字是随机写的，下面也一样
			-  ... 
		- valid
			- 2.jpg
			- 7.jpg
			-  ...
		- test
			- 3.jpg
			- 9.jpg
			- ...
	- labels
		- train
			- 1.txt
			- 5.txt
			-  ... 
		- valid
			- 2.txt
			- 7.txt
			-  ...
		- test
			- 3.txt
			- 9.txt
			- ...

代码实现

以下代码可以直接复制运行（只需要修改参数）

本代码需要自行修改4个参数，分别是：
img_dir:存放图像数据的目录；
label_dir:存放对应YOLO格式标签的txt文件的目录；
split_dir:要存放划分后数据集的目录；
TEST:True表示需要划分测试集，否则不。

"""
将数据集划分为训练集，验证集，测试集
"""
import os
import random
import shutil


# 创建保存数据的文件夹
def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

def split_data(img_dir, label_dir, split_dir, TEST=True):
    '''
    args:
        img_dir:原图片数据集路径
        label_dir:原yolog格式txt文件数据集路径
        split_dir:划分后数据集保存路径
        TEST:是否划分测试集
        用于将数据集划分为YOLO数据集格式的训练集,验证集,测试集
    '''
    random.seed(42)  # 随机种子
    # 1.确定原图片数据集路径
    datasetimg_dir = img_dir
    # 确定原label数据集路径
    datasetlabel_dir = label_dir

    images_dir = os.path.join(split_dir, 'images')
    labels_dir = os.path.join(split_dir, 'labels')
    dir_list = [images_dir, labels_dir]
    if TEST:
        type_label = ['train', 'valid', 'test']
    else:
        type_label = ['train', 'valid']

    for i in range(len(dir_list)):
        for j in range(len(type_label)):
            makedir(os.path.join(dir_list[i], type_label[j]))

    # 3.确定将数据集划分为训练集，验证集，测试集的比例
    train_pct = 0.8
    valid_pct = 0.1
    test_pct = 0.1
    # 4.划分
    labels = os.listdir(datasetlabel_dir)  # 展示目标文件夹下所有的文件名
    labels = list(filter(lambda x: x.endswith('.txt'), labels))  # 取到所有以.txt结尾的yolo格式文件
    random.shuffle(labels)  # 乱序路径
    label_count = len(labels)  # 计算图片数量
    train_point = int(label_count * train_pct)  # 0:train_pct
    valid_point = int(label_count * (train_pct + valid_pct))  # train_pct:valid_pct
    for i in range(label_count):
        if i < train_point:  # 保存0-train_point的图片和标签到训练集
            out_dir = os.path.join(images_dir, 'train')
            label_out_dir = os.path.join(labels_dir, 'train')

        elif train_point <= i < valid_point:  # 保存train_point-valid_point的图片和标签到验证集
            out_dir = os.path.join(images_dir, 'valid')
            label_out_dir = os.path.join(labels_dir, 'valid')
        else:  # 保存test_point-结束的图片和标签到测试集
            out_dir = os.path.join(images_dir, 'test')
            label_out_dir = os.path.join(labels_dir, 'test')

        label_target_path = os.path.join(label_out_dir, labels[i])  # 指定目标保存路径
        label_src_path = os.path.join(datasetlabel_dir, labels[i])  # 指定目标原文件路径
        img_target_path = os.path.join(out_dir, labels[i].split('.')[0] + '.JPG')  # 指定目标保存路径
        img_src_path = os.path.join(datasetimg_dir, labels[i].split('.')[0] + '.JPG')  # 指定目标原图像路径
        shutil.copy(label_src_path, label_target_path)  # 复制txt
        shutil.copy(img_src_path, img_target_path)  # 复制图片
        

    print('train:{}, valid:{}, test:{}'.format(train_point, valid_point - train_point,
                                               label_count - valid_point))


if __name__ == "__main__":
    img_dir = r'totrain\segmentdata_withlabel\jpg'
    label_dir = r'totrain\segmentdata_withlabel\txt'
    split_dir = r'totrain\segmentdata_withlabel\split'
    split_data(img_dir, label_dir, split_dir, TEST=True)
    # for dir in os.listdir(root):
    #     if not dir.endswith('.py') and not dir.endswith('txt'):
    #         split_data(dir, label_dir, root, split_dir)

同时，本代码默认按照8:1:1的比例划分训练集、验证集和测试集如果需要修改这个比例，请在代码如下部分进行修改：

    # 3.确定将数据集划分为训练集，验证集，测试集的比例
    train_pct = 0.8
    valid_pct = 0.1
    test_pct = 0.1

在我自己的数据集中，划分完成如下图所示：