实用脚本六：python划分YOLOv8-Seg数据集训练集测试集

YOLOv8-Seg数据集的格式如下：

coco8-seg---images---train---...png

                                ---val

                ---labels   ---train---...txt

                                ---val

其中txt文本的格式为

<class-index> <x1> <y1> <x2> <y2> ... <xn> <yn>

注意x1、y1这些坐标需要根据图像大小进行归一化，如下45是类别index，后面是归一化后的坐标点

45 0.782016 0.986521 0.937078 0.874167 0.957297 0.782021 0.950562 0.739333 0.825844 0.561792 0.714609 0.420229 0.657297 0.391021 0.608422 0.4 0.0303438 0.750562 0.0016875 0.811229 0.003375 0.889896 0.0320156 0.986521

我们最开始的数据集是一个文件夹里有图像和标签文件，然后有一些图像没有标签，所以选择按照标签进行划分

# -*- coding: utf-8 -*-
"""
将数据集划分为训练集，验证集，测试集
"""

import os
import random
import shutil

# 创建保存图像的文件夹
def makedir(new_dir):
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)
random.seed(1) # 随机种子

# 1.确定原图像数据集路径
dataset_dir = "data/"  ##原始数据集路径

# 2.确定数据集划分后保存的路径
split_dir = "dataset/images"
label_dir = 'dataset/labels'##划分后保存路径
train_dir = os.path.join(split_dir, "train")
train_label_dir = os.path.join(label_dir, "train")
valid_dir = os.path.join(split_dir, "val")
valid_label_dir = os.path.join(label_dir, "val")
test_dir = os.path.join(split_dir, "test")
test_label_dir = os.path.join(label_dir, "test")

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

# 4.基于img划分
# for root, dirs, files in os.walk(dataset_dir):
#     for sub_dir in dirs: # 遍历0，1，2，3，4，5...9文件夹
#         imgs = os.listdir(os.path.join(root, sub_dir)) # 展示目标文件夹下所有的文件名
#         imgs = list(filter(lambda x: x.endswith('.png'), imgs)) # 取到所有以.png结尾的文件，如果改了图片格式，这里需要修改
#         random.shuffle(imgs)  # 乱序图片路径
#         img_count = len(imgs)  # 计算图片数量
#         train_point = int(img_count * train_pct)  # 0:train_pct
#         valid_point = int(img_count * (train_pct + valid_pct))  # train_pct:valid_pct
#
#         for i in range(img_count):
#             if i < train_point:  # 保存0-train_point的图片到训练集
#                 out_dir = train_dir
#                 out_label_dir = train_label_dir
#             elif i < valid_point:  # 保存train_point-valid_point的图片到验证集
#                 out_dir = valid_dir
#                 out_label_dir = valid_label_dir
#             else:  #  保存valid_point-结束的图片到测试集
#                 out_dir = test_dir
#                 out_label_dir = test_label_dir
#             makedir(out_dir) # 创建文件夹
#             target_path = os.path.join(out_dir, imgs[i]) # 指定目标保存路径
#             src_path = os.path.join(dataset_dir, sub_dir, imgs[i])#指定目标原图像路径
#             label_name = os.path.splitext(imgs[i])[0]+'.txt'
#             label_path = os.path.join(dataset_dir, sub_dir, label_name)
#             target_label_path = os.path.join(out_label_dir, label_name)
#             shutil.copy(src_path, target_path)
#             shutil.copy(label_path, target_label_path)
#
#         print('Class:{}, train:{}, valid:{}, test:{}'.format(sub_dir, train_point, valid_point-train_point,
#                                                              img_count-valid_point))
#基于label划分
for root, dirs, files in os.walk(dataset_dir):
    for sub_dir in dirs: # 遍历0，1，2，3，4，5...9文件夹
        labs = os.listdir(os.path.join(root, sub_dir)) # 展示目标文件夹下所有的文件名
        labs = list(filter(lambda x: x.endswith('.txt'), labs)) # 取到所有以.txt结尾的文件，如果改了图片格式，这里需要修改
        random.shuffle(labs)  # 乱序标签路径
        labs_count = len(labs)  # 计算标签数量
        train_point = int(labs_count * train_pct)  # 0:train_pct
        valid_point = int(labs_count * (train_pct + valid_pct))  # train_pct:valid_pct

        for i in range(labs_count):
            if i < train_point:  # 保存0-train_point的图片到训练集
                out_dir = train_dir
                out_label_dir = train_label_dir
            elif i < valid_point:  # 保存train_point-valid_point的图片到验证集
                out_dir = valid_dir
                out_label_dir = valid_label_dir
            else:  #  保存valid_point-结束的图片到测试集
                out_dir = test_dir
                out_label_dir = test_label_dir
            if not os.path.exists(out_dir):
                os.makedirs(out_dir)
            if not os.path.exists(out_label_dir):
                os.makedirs(out_label_dir)

            src_path = os.path.join(dataset_dir, sub_dir, labs[i])#指定目标原标签路径
            img_name = os.path.splitext(labs[i])[0]+'.png'#标签对应的图像名
            img_path = os.path.join(dataset_dir, sub_dir, img_name)#标签对应的图像路径
            target_img_path = os.path.join(out_dir, img_name)#保存的图像路径
            shutil.copy(src_path, out_label_dir)
            shutil.copy(img_path, target_img_path)

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

划分完之后在目标文件里获得YOLOv8-seg格式的分割数据集