YOLOv5 or YOLOv8 快速劃分訓練集 驗證集 測試集

該腳本實作了将原始資料集自動劃分為yolo訓練資料集排列形式

import os
import random
import shutil
"""
該腳本實作了将原始資料集自動劃分為yolo訓練資料集排列形式 (目錄排序如下)

old_root:
    dataset
        images
        yolo_labels
        
new_root:
    dataset
        test
            images
            labels
        train
            images
            labels  
        val
            images
            labels          
"""


# 檔案的遞歸删除(清空path檔案夾下的所有檔案，但不會删除空檔案夾)
def del_dir(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)

        # 如果是一個檔案就删除，不然繼續遞歸檔案夾。
        if os.path.isfile(path_file):
            os.remove(path_file)

        elif os.path.isdir(path_file):
            file = os.listdir(path_file)
            if len(file) == 0:
                os.rmdir(path_file)
            else:
                del_dir(path_file)
                os.rmdir(path_file)


def remake_dir(path):
    if os.path.exists(path):
        os.mkdir(os.path.join(path, "images"))
        os.mkdir(os.path.join(path, "labels"))
    else:
        os.mkdir(path)
        os.mkdir(os.path.join(path, "images"))
        os.mkdir(os.path.join(path, "labels"))


def copy_files(old_root, new_root):

    # 保證随機可複現
    # 這保證了每次調用copy_files函數時的随機值是相同的
    random.seed(0)

    file_lists = os.listdir(old_root)

    all_files = []
    train_files = []
    val_files = []
    test_files = []

    for file in file_lists:
        all_files.append(file)

    L = len(all_files)
    train_num = int(0.8 * L)
    val_num = int(0.1 * L)
    test_num = L - train_num - val_num

    # 随機選取訓練資料
    for i in range(train_num):
        temp = random.choice(all_files)
        train_files.append(temp)
        all_files.remove(temp)
    # 随機選取驗證資料
    for i in range(val_num):
        temp = random.choice(all_files)
        val_files.append(temp)
        all_files.remove(temp)
    # 随機選取測試資料
    for i in range(test_num):
        temp = random.choice(all_files)
        test_files.append(temp)
        all_files.remove(temp)

    # 拷貝對應的訓練集資料
    for file_name in train_files:

        if (".png" or ".jpg") in file_name:
            t_name = str("train\\images")
        else:
            t_name = str("train\\labels")

        old_file_path = os.path.join(old_root + str("\\") + file_name)
        new_file_path = os.path.join(new_root + str("\\") + t_name + str("\\") + file_name)

        shutil.copy(old_file_path, new_file_path)

    # 拷貝對應的驗證集資料
    for file_name in val_files:

        if (".png" or ".jpg") in file_name:
            t_name = str("val\\images")
        else:
            t_name = str("val\\labels")

        old_file_path = os.path.join(old_root + str("\\") + file_name)
        new_file_path = os.path.join(new_root + str("\\") + t_name + str("\\") + file_name)

        shutil.copy(old_file_path, new_file_path)

    # 拷貝對應的測試集資料
    for file_name in test_files:

        if (".png" or ".jpg") in file_name:
            t_name = str("test\\images")
        else:
            t_name = str("test\\labels")

        old_file_path = os.path.join(old_root + str("\\") + file_name)
        new_file_path = os.path.join(new_root + str("\\") + t_name + str("\\") + file_name)

        shutil.copy(old_file_path, new_file_path)

    return


def main():

    old_root = r"D:\datasets\UCAS_AOD\CAR"

    new_root = r"D:\datasets\UCAS_AOD\train_yolo"

    path_list = ["train", "val", "test"]

    # 該部分循環實作了若檔案夾存在則删除檔案夾下的所有檔案，并建立"images"和"labels"
    # 若檔案夾不存在，則建立該檔案，并建立子檔案夾"images"和"labels"
    # 注釋該部分則程式不會對檔案是否存在進行判斷，直接進行copy操作
    for path in path_list:
        path = os.path.join(new_root, path)
        if os.path.exists(path):
            del_dir(path)
            remake_dir(path)

        else:
            remake_dir(path)

    # 該部分實作對原始資料集的劃分
    order_list = ["images", "yolo_labels"]
    for root in order_list:

        order_root = os.path.join(old_root, root)

        copy_files(order_root, new_root)


if __name__ == '__main__':
    main()