YOLOV5——將圖片和標注數據按比例切分為訓練集和測試集

將圖片和標注數據按比例切分后存儲至新的路徑下

# 將圖片和標注數據按比例切分為 訓練集和測試集

import os
from shutil import copy2
# 原始路徑
image_original_path = "../image_data/seed/images/"
label_original_path = "../image_data/seed/labels/"
# 上級目錄
parent_path = os.path.dirname(os.getcwd())
# 訓練集路徑
train_image_path = os.path.join(parent_path, "image_data/seed/train/images/")
train_label_path = os.path.join(parent_path, "image_data/seed/train/labels/")
# 測試集路徑
test_image_path = os.path.join(parent_path, 'image_data/seed/test/images/')
test_label_path = os.path.join(parent_path, 'image_data/seed/test/labels/')

# 檢查文件夾是否存在
def mkdir():
    if not os.path.exists(train_image_path):
        os.makedirs(train_image_path)
    if not os.path.exists(train_label_path):
        os.makedirs(train_label_path)

    if not os.path.exists(test_image_path):
        os.makedirs(test_image_path)
    if not os.path.exists(test_label_path):
        os.makedirs(test_label_path)


def main():
    mkdir()
    # 復制移動圖片數據
    all_image = os.listdir(image_original_path)
    for i in range(len(all_image)):
        if i % 10 != 0:
            copy2(os.path.join(image_original_path, all_image[i]), train_image_path)
        else:
            copy2(os.path.join(image_original_path, all_image[i]), test_image_path)

    # 復制移動標注數據
    all_label = os.listdir(label_original_path)
    for i in range(len(all_label)):
        if i % 10 != 0:
            copy2(os.path.join(label_original_path, all_label[i]), train_label_path)
        else:
            copy2(os.path.join(label_original_path, all_label[i]), test_label_path)


if __name__ == '__main__':
    main()