2.YOLO-TXT标签等比例划分数据集（附完整代码）

作用：yolo训练前需要自己等比例划分数据集
条件：标注文件是TXT格式

设置目录路径：
“”"
训练集：验证集：测试集 （第73行，自己定比例）
自建yolo原始图片的文件夹（第71行）
自建yolo原始标签xtx的文件夹（第70行）
另存放划分比例后的文件夹（第15行）
“”"
主要是修改这四个地方即可，话不多说 ，咱们直接附代码

"""
条件为标注文件是TXT格式
训练集：验证集 （第73行，自己定比例）
自建yolo原始图片的文件夹（第70行）
自建yolo原始标签TXT的文件夹（第71行）
另存放划分比例8:2后的文件夹（第17行）
"""

import os
import shutil
import random
from tqdm import tqdm

def split_img(img_path, label_path, split_list):
    try:  # 创建数据集文件夹
        # Data = 'F:/data7_potato/dataest'  # 另存放划分比例后的文件夹
        Data = 'F:/data16/dataest'  # 另存放划分比例后的文件夹

        train_img_dir = Data + '/images/train'
        val_img_dir = Data + '/images/val'
        test_img_dir = Data + '/images/test'

        train_label_dir = Data + '/labels/train'
        val_label_dir = Data + '/labels/val'
        test_label_dir = Data + '/labels/test'

        # 创建文件夹
        os.makedirs(train_img_dir, exist_ok=True)
        os.makedirs(train_label_dir, exist_ok=True)
        os.makedirs(val_img_dir, exist_ok=True)
        os.makedirs(val_label_dir, exist_ok=True)
        os.makedirs(test_img_dir, exist_ok=True)
        os.makedirs(test_label_dir, exist_ok=True)

    except:
        print('文件目录已存在')

    total_imgs = os.listdir(img_path)
    total_img_path = [os.path.join(img_path, img) for img in total_imgs]
    total_labels = [toLabelPath(img, label_path) for img in total_img_path]

    total_size = len(total_imgs)
    train_size = int(split_list[0] * total_size)
    val_size = int(split_list[1] * total_size)
    test_size = total_size - train_size - val_size

    train_img_idx = random.sample(range(total_size), train_size)
    val_img_idx = random.sample([i for i in range(total_size) if i not in train_img_idx], val_size)
    test_img_idx = [i for i in range(total_size) if i not in train_img_idx + val_img_idx]

    for idx in tqdm(train_img_idx, desc='train ', ncols=80, unit='img'):
        shutil.copy(total_img_path[idx], train_img_dir)
        shutil.copy(total_labels[idx], train_label_dir)

    for idx in tqdm(val_img_idx, desc='val ', ncols=80, unit='img'):
        shutil.copy(total_img_path[idx], val_img_dir)
        shutil.copy(total_labels[idx], val_label_dir)

    for idx in tqdm(test_img_idx, desc='test ', ncols=80, unit='img'):
        shutil.copy(total_img_path[idx], test_img_dir)
        shutil.copy(total_labels[idx], test_label_dir)

def toLabelPath(img_path, label_path):
    img = os.path.basename(img_path)
    label = img.split('.jpg')[0] + '.txt'
    return os.path.join(label_path, label)

if __name__ == '__main__':
    # img_path = 'F:/data7_potato/images'   # 自建yolo原始图片的文件夹
    img_path = 'F:/data16/images'
    label_path = 'F:/data16/labels'  # 自建yolo原始标签TXT的文件夹
    # label_path = 'F:/data7_potato/labels'  # 自建yolo原始标签TXT的文件夹
    split_list = [0.8, 0.1]  # 训练集：验证集：测试集（8:1:1）自己定数据集的划分比例[train:val:test]
    split_img(img_path, label_path, split_list)