YOLO目标检测数据集划分代码

原创

于 2025-02-06 16:38:31 发布 · 319 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #目标检测 #人工智能

使用自制数据集训练YOLO目标检测算法前，需要对数据集进行划分，以下代码可以将数据集的图片和标签分别保存

import os
import shutil
import numpy as np
from tqdm import tqdm


def split_dataset(images_dir, labels_dir, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1, random_seed=None):
    # # 检查比例是否合法
    # if train_ratio + val_ratio + test_ratio!= 1.0:
    #     raise ValueError("The sum of train_ratio, val_ratio, and test_ratio must be 1.0")
    # 获取所有图片文件的名称列表
    image_files = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
    num_images = len(image_files)
    # 使用 numpy 的随机排列函数
    if random_seed is not None:
        np.random.seed(random_seed)
    # 先将图片文件列表按照比例分成三个部分，而不是先随机打乱
    train_indices = np.random.choice(num_images, size=int(num_images * train_ratio), replace=False)
    remaining_indices = np.setdiff1d(np.arange(num_images), train_indices)
    val_indices = np.random.choice(remaining_indices, size=int(num_images * val_ratio), replace=False)
    test