yolov5训练数据数据集切分

        上个学期，我参加了一个竞赛是关于检测车牌并识别车牌内容的，这无疑要用到yolov来进行训练数据，对数据集切分便是训练数据的前提，我也从网上找了很多代码，也并没有找到和自己预期效果一样的，于是便自己改编，以下代码我将解释他的最终切分效果和它怎么使用。

        之前我写过一篇相同的文章，但当时一时疏忽，到今天用的时候发现出现了很多问题，我对此给大家道歉，也用了1个小时进行改正，如果之前看过我写的文章，对大家造成困扰，我为此向看过文章的所有人表达深深的歉意，下面是改好的（试过了没有什么问题），真诚的希望能给大家带来帮助。

        

import os
import shutil
import random
def extract_train_images(new_dir, image_dir, new_image_dir):#该函数的作用是根据new_dir筛选出对应的图片放到new_image_dir中
    # 确保新目录存在
    if not os.path.exists(new_image_dir):
        os.makedirs(new_image_dir)
    # 遍历每个不含NULL的txt文件
    for txt_filename in os.listdir(new_dir):
        # 获取对应的图片编号
        image_num = os.path.splitext(txt_filename)[0]
        # 构造图片文件名
        image_filename = f'{image_num}.jpg'
        image_path = os.path.join(image_dir, image_filename)
        # 检查图片是否存在
        print(image_path)
        if os.path.exists(image_path):
            # 将图片复制到新目录
            new_image_path = os.path.join(new_image_dir, image_filename)
            shutil.copy(image_path, new_image_path)
            print(f'Copied {image_filename} to new directory')
        else:
            print(f'Image {image_filename} not found')
def split_traindataset(new_image_dir,new_dir, train, train1, text, text1, val,val1):#这个是具体进行切分的函数
    txt_files = os.listdir(new_dir)
    split_point = int(0.8 * len(txt_files))
    split_point1= int(0.9 * len(txt_files))#按照811的比例分成训练集，测试集，验证集
    train_txt_files = txt_files[:split_point]
    test_txt_files = txt_files[split_point:split_point1]
    val_txt_files = txt_files[split_point1:]#进行分割
    for txt in train_txt_files:
        shutil.copy(os.path.join(new_dir, txt), "D:/Date/train_txt/")
    extract_train_images(train1, new_image_dir, train)
    # 将测试集复制到新目录
    for txt in test_txt_files:
        shutil.copy(os.path.join(new_dir, txt), "D:/Date/test_txt/")
    extract_train_images(text1, new_image_dir, text)
    for txt in val_txt_files:
        shutil.copy(os.path.join(new_dir, txt), "D:/Date/val_txt/")
    extract_train_images(val1, new_image_dir, val)
    print("数据集切分完成")
if __name__ == '__main__':
    #写出所有所需文件夹的路径
    new_dir = "D:/Date/yolo_txt"#所有标注好的图片对应的txt文档
    new_image_dir = "D:/Date/images"#所有的数据集
    train = "D:/Date/train"#训练集图片
    train1 = "D:/Date/train_txt"#训练集图片对应的txt文档
    text = "D:/Date/test"#测试集图片
    text1 = "D:/Date/test_txt"#测试集图片对应的txt文档
    val = "D:/Date/val"#验证集图片
    val1 = "D:/Date/val_txt"#验证集图片对应的txt文档
    split_traindataset(new_image_dir,new_dir,train,train1,text,text1,val,val1)#调用函数

        这里，大家有个点一定要注意，就是主函数下定义的new_image_dir和函数extract_train_images里的new_image_dir是不同的。

        在一些特殊情况，比如说对车牌进行训练，有些图片没有车牌，我们自然无法进行标注，就不会生成txt文档，但在数据切分时，需要的两个文件夹必须是所有数据集图片和它对应标注好的txt文档一一对应的，这是我们要先调用一下函数extract_train_images进行处理。比如说，文件夹a为·所有标注好的图片生成的txt文档，文件夹b为所有的数据集，文件夹c为与文件夹a一一对应的数据集中的图片（调用函数前为空），我们就可以调用函数extract_train_images（a,b,c）此时c就可以作为我们数据集切分的new_image_dir。

       

import os
import shutil
def extract_train_images(new_dir, image_dir, new_image_dir):#该函数的作用是根据new_dir筛选出对应的图片放到new_image_dir中
    # 确保新目录存在
    if not os.path.exists(new_image_dir):
        os.makedirs(new_image_dir)
    # 遍历每个不含NULL的txt文件
    for txt_filename in os.listdir(new_dir):
        # 获取对应的图片编号
        image_num = os.path.splitext(txt_filename)[0]
        # 构造图片文件名
        image_filename = f'{image_num}.jpg'
        image_path = os.path.join(image_dir, image_filename)
        # 检查图片是否存在
        if os.path.exists(image_path):
            # 将图片复制到新目录
            new_image_path = os.path.join(new_image_dir, image_filename)
            shutil.copy(image_path, new_image_path)
            print(f'Copied {image_filename} to new directory')
        else:
            print(f'Image {image_filename} not found')
new_dir = "D:/Date/yolo_txt"#所有标注好的图片对应的txt文档
new_image_dir = "D:/Date/images"#所有的数据集
train = "D:/Date/image"#与new_dir一一对应的数据集中的图片
extract_train_images(new_dir,new_image_dir,train)

        如果还有问题，希望大家多多指教。