划分数据集用于yolov8训练

将数据集划分用于yolov8训练

图片的名称和标签的名称一致
1.jpg ------> 1.txt

images_dir=‘/path/to/Images’ # 图片存放的位置
labels_dir=‘/path/to/Labels’, #图片对应label存放的位置
target_dir=‘/path/to/save/target_img_labels’, # 输出划分数据的root路径

代码1

import os
import random
from shutil import copyfile

def split_dataset(images_dir, labels_dir, target_dir, train_ratio=0.7, val_ratio=0.3, seed=None):
    """
    将数据集按照指定比例划分为训练集和验证集。

    参数:
    - images_dir: 包含图像文件的目录路径
    - labels_dir: 包含标签文件的目录路径
    - target_dir: 目标数据集文件夹路径，将包含划分后的 train 和 val 子目录
    - train_ratio: 训练集比例
    - val_ratio: 验证集比例
    - seed: 随机种子，可选参数，用于确保随机性可复现

    返回:
    无返回值，划分后的数据集将保存在 target_dir 中
    """
    # 获取图像文件列表
    image_files = [file for file in os.listdir(images_dir) if file.endswith('.jpg')]

    # 随机打乱文件顺序
    random.seed(seed)
    random.shuffle(image_files)

    # 计算划分数据集的索引
    total_files = len(image_files)
    train_split = int(train_ratio * total_files)
    val_split = int(val_ratio * total_files)

    # 创建目标文件夹结构
    os.makedirs(os.path.join(target_dir, 'train', 'images'), exist_ok=True)
    os.makedirs(os.path.join(target_dir, 'train', 'labels'), exist_ok=True)
    os.makedirs(os.path.join(target_dir, 'val', 'images'), exist_ok=True)
    os.makedirs(os.path.join(target_dir, 'val', 'labels'), exist_ok=True)

    # 划分数据集并复制到目标文件夹
    for file in image_files[:train_split]:
        src_image = os.path.join(images_dir, file)
        src_label = os.path.join(labels_dir, file.replace('.jpg', '.txt'))
        dst_image = os.path.join(target_dir, 'train', 'images', file)
        dst_label = os.path.join(target_dir, 'train', 'labels', file.replace('.jpg', '.txt'))
        copyfile(src_image, dst_image)
        copyfile(src_label, dst_label)

    for file in image_files[train_split:train_split + val_split]:
        src_image = os.path.join(images_dir, file)
        src_label = os.path.join(labels_dir, file.replace('.jpg', '.txt'))
        dst_image = os.path.join(target_dir, 'val', 'Images', file)
        dst_label = os.path.join(target_dir, 'val', 'Labels', file.replace('.jpg', '.txt'))
        copyfile(src_image, dst_image)
        copyfile(src_label, dst_label)

# 调用函数进行数据集划分
split_dataset(
    images_dir='/path/to/Images',
    labels_dir='/path/to/Labels',
    target_dir='/path/to/save/target_img_labels',
    train_ratio=0.7,
    val_ratio=0.3,
    seed=42  # 可选，设置随机种子以确保可复现性
)

代码2

import os
import shutil
import random
from tqdm import tqdm

# 定义原始数据文件夹和目标数据集文件夹
data_dir = '/path/to/Images'
target_dir = '/path/to/labels'

# 定义数据集划分比例
train_split_ratio = 0.7
val_split_ratio = 0.3

# 创建目标数据集文件夹及其子目录结构
os.makedirs(target_dir, exist_ok=True)
os.makedirs(os.path.join(target_dir, 'meta'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'val'), exist_ok=True)

# 获取原始数据文件夹下的子目录列表
categories = os.listdir(data_dir)

# 遍历每个子目录
for category in categories:
    # 获取该类别下的所有文件
    files = os.listdir(os.path.join(data_dir, category))

    # 随机打乱文件顺序
    random.shuffle(files)

    # 计算划分数据集的索引
    total_files = len(files)
    train_split = int(train_split_ratio * total_files)
    val_split = int(val_split_ratio * total_files)

    # 划分数据集并复制到目标文件夹，使用tqdm添加进度条
    for file in tqdm(files[:train_split], desc=f'Copying train data for {category}'):
        src = os.path.join(data_dir, category, file)
        dst = os.path.join(target_dir, 'train', category)
        os.makedirs(dst, exist_ok=True)
        shutil.copy(src, os.path.join(dst, file))

    for file in tqdm(files[train_split:train_split + val_split], desc=f'Copying validation data for {category}'):
        src = os.path.join(data_dir, category, file)
        dst = os.path.join(target_dir, 'val', category)
        os.makedirs(dst, exist_ok=True)
        shutil.copy(src, os.path.join(dst, file))

# 创建标注文件（train.txt、val.txt、test.txt）
with open(os.path.join(target_dir, 'meta', 'train.txt'), 'w') as train_txt:
    for category in categories:
        train_files = os.listdir(os.path.join(target_dir, 'train', category))
        for file in train_files:
            train_txt.write(f'{os.path.join("train", category, file)} {category}\n')

with open(os.path.join(target_dir, 'meta', 'val.txt'), 'w') as val_txt:
    for category in categories:
        val_files = os.listdir(os.path.join(target_dir, 'val', category))
        for file in val_files:
            val_txt.write(f'{os.path.join("val", category, file)} {category}\n')

print("数据集划分完成！")