YOLO数据集划分（test+val+train）

Curry_Math

于 2024-08-10 20:06:40 发布

阅读量139

点赞数 4

分类专栏：神经网络与深度学习文章标签： YOLO 机器学习深度学习

本文链接：https://blog.csdn.net/qq_43060884/article/details/141095506

版权

神经网络与深度学习专栏收录该内容

4 篇文章 2 订阅

订阅专栏

YOLO数据集划分（test+val+train）

import os
import shutil
import random

random.seed(0)


def split_data(file_path, label_path, new_file_path, train_rate, val_rate, test_rate):
    # Get all image files, ignoring non-image files
    eachclass_image = [file for file in os.listdir(file_path) if file.lower().endswith(('.png', '.jpg', '.jpeg'))]
    total = len(eachclass_image)
    random.shuffle(eachclass_image)
    train_images = eachclass_image[:int(train_rate * total)]
    val_images = eachclass_image[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_images = eachclass_image[int((train_rate + val_rate) * total):]

    # Create base directories for images and labels
    images_base_path = os.path.join(new_file_path, 'images')
    labels_base_path = os.path.join(new_file_path, 'labels')
    for path in [images_base_path, labels_base_path]:
        if not os.path.exists(path):
            os.makedirs(path)

    # Copy the files into the appropriate directories
    for dataset_type, images in [('train', train_images), ('val', val_images), ('test', test_images)]:
        images_path = os.path.join(images_base_path, dataset_type)
        labels_path = os.path.join(labels_base_path, dataset_type)

        if not os.path.exists(images_path):
            os.makedirs(images_path)
        if not os.path.exists(labels_path):
            os.makedirs(labels_path)

        for image in images:
            old_image_path = os.path.join(file_path, image)
            new_image_path = os.path.join(images_path, image)
            shutil.copy(old_image_path, new_image_path)

            base_filename = os.path.splitext(image)[0]
            old_label_path = os.path.join(label_path, base_filename + '.txt')
            new_label_path = os.path.join(labels_path, base_filename + '.txt')
            if os.path.exists(old_label_path):
                shutil.copy(old_label_path, new_label_path)
            else:
                print(f"Label file for {image} does not exist, skipping.")


if __name__ == '__main__':
    #以下路径需要自己替换
    image_path = "D:/Academic/Deep Learning/yolov10-main/data/NEU/NEU-DET (1)/IMAGES"
    txt_path = 'D:/Academic/Deep Learning/yolov10-main/data/NEU/NEU-DET (1)/labels'
    new_dataset_path = "D:/Academic/Deep Learning/yolov10-main/data/NEU/divided"
    split_data(image_path, txt_path, new_dataset_path, train_rate=0.7, val_rate=0.1, test_rate=0.2)