YOLO数据集分割训练集、测试集和验证集

记录一下自己的分割代码。

注意:

  1. 这是在windows环境,请Linux的同学们注意。
  2. 标签为txt,图像为jpg,其他的我没试过喔。

训练集、验证集、测试集(7:2:1)

复制代码
import os
import shutil
import random
from tqdm import tqdm

"""
标注文件是yolo格式(txt文件)
训练集:验证集:测试集 (7:2:1) 
"""

def split_img(img_path, label_path, split_list):
    try:
        Data = './ImageSets'
        # 创建文件夹结构
        train_img_dir = os.path.join(Data, 'images/train')
        val_img_dir = os.path.join(Data, 'images/val')
        test_img_dir = os.path.join(Data, 'images/test')

        train_label_dir = os.path.join(Data, 'labels/train')
        val_label_dir = os.path.join(Data, 'labels/val')
        test_label_dir = os.path.join(Data, 'labels/test')

        # 创建所有目录
        os.makedirs(train_img_dir, exist_ok=True)
        os.makedirs(val_img_dir, exist_ok=True)
        os.makedirs(test_img_dir, exist_ok=True)
        os.makedirs(train_label_dir, exist_ok=True)
        os.makedirs(val_label_dir, exist_ok=True)
        os.makedirs(test_label_dir, exist_ok=True)

    except Exception as e:
        print(f'创建目录时出错: {e}')

    train, val, test = split_list
    # 仅处理存在对应标签文件的图像
    all_img_path = []
    for img in os.listdir(img_path):
        img_full = os.path.join(img_path, img)
        label_full = toLabelPath(img_full, label_path)
        if os.path.isfile(label_full):
            all_img_path.append(img_full)
        else:
            print(f"跳过无标签的图像: {img_full}")

    # 随机打乱所有图像路径
    random.shuffle(all_img_path)
    total = len(all_img_path)
    train_num = int(total * train)
    val_num = int(total * val)
    test_num = total - train_num - val_num

    # 划分数据集
    train_img = all_img_path[:train_num]
    val_img = all_img_path[train_num:train_num + val_num]
    test_img = all_img_path[train_num + val_num:]

    # 复制训练集
    for img in tqdm(train_img, desc='训练集', unit='img'):
        label = toLabelPath(img, label_path)
        _copy(img, train_img_dir)
        _copy(label, train_label_dir)

    # 复制验证集
    for img in tqdm(val_img, desc='验证集', unit='img'):
        label = toLabelPath(img, label_path)
        _copy(img, val_img_dir)
        _copy(label, val_label_dir)

    # 复制测试集
    for img in tqdm(test_img, desc='测试集', unit='img'):
        label = toLabelPath(img, label_path)
        _copy(img, test_img_dir)
        _copy(label, test_label_dir)

def _copy(from_path, to_dir):
    """复制文件到目标目录,并确保目标目录存在"""
    try:
        os.makedirs(to_dir, exist_ok=True)
        shutil.copy(from_path, to_dir)
    except Exception as e:
        print(f"复制 {from_path} 到 {to_dir} 失败: {e}")

def toLabelPath(img_path, label_path):
    """根据图片路径生成对应的标签路径"""
    img_filename = os.path.basename(img_path)
    base = os.path.splitext(img_filename)[0]  # 去除扩展名
    label_filename = base + '.txt'
    return os.path.join(label_path, label_filename)

if __name__ == '__main__':
    # 使用原始字符串避免转义问题
    img_path = r'D:\文件\PlantVillage_for_object_detection\Dataset\images'    #自己的图片文件地址
    label_path = r'D:\文件\PlantVillage_for_object_detection\Dataset\labels'    #自己的标签文件地址
    split_list = [0.7, 0.2, 0.1]  # 训练集:验证集:测试集
    split_img(img_path, label_path, split_list)

over

相关推荐
lianyinghhh17 分钟前
yolo8实现目标检测
yolo·目标检测·macos
PyAIExplorer29 分钟前
图像处理中的霍夫变换:直线检测与圆检测
图像处理·人工智能
野指针1213835 分钟前
【使用Flask基于PaddleOCR3.0开发一个接口 调用时报错RuntimeError: std::exception】
python·flask
千宇宙航1 小时前
闲庭信步使用图像验证平台加速FPGA的开发:第七课——获取RAW图像
图像处理·计算机视觉·fpga开发
格林威1 小时前
Baumer工业相机堡盟工业相机如何通过DeepOCR模型识别判断数值和字符串的范围和相似度(C#)
开发语言·人工智能·python·数码相机·计算机视觉·c#·视觉检测
赵英英俊1 小时前
Python Day8
python
DAWN_T171 小时前
Python编译器(Pycharm Jupyter)
python·机器学习·jupyter·pycharm
Draina1 小时前
在pycharm中运行sagemath脚本的配置过程
ide·python·安全·pycharm·密码学
presenttttt2 小时前
用Python和OpenCV从零搭建一个完整的双目视觉系统(五)
开发语言·python·opencv·计算机视觉