按比例划分训练集、测试集和验证集

最新推荐文章于 2024-03-30 16:29:35 发布

xzz_deng

最新推荐文章于 2024-03-30 16:29:35 发布

阅读量246

点赞数

文章标签：深度学习机器学习人工智能

本文链接：https://blog.csdn.net/zunzunle/article/details/130267947

版权

这段代码使用Python的os和shutil模块，以及numpy库，将.jpg图片及其对应的.txt标签文件按照训练集、验证集和测试集的比例进行划分和复制。它首先创建了相应的目录结构，然后随机打乱文件列表，最后按照指定比例分配文件到各个数据集中。

摘要由CSDN通过智能技术生成

import os, shutil, random
import numpy as np

postfix = 'jpg'
img_path = './annotated-images'
label_path = './labels_yolo'
dataset_path = './train_yolo'
val_size, test_size = 0.1, 0.1

os.makedirs(dataset_path, exist_ok=True)
os.makedirs(f'{dataset_path}/images', exist_ok=True)
os.makedirs(f'{dataset_path}/labels', exist_ok=True)
os.makedirs(f'{dataset_path}/images/train', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/train', exist_ok=True)
os.makedirs(f'{dataset_path}/images/val', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/val', exist_ok=True)
os.makedirs(f'{dataset_path}/images/test', exist_ok=True)
os.makedirs(f'{dataset_path}/labels/test', exist_ok=True)

all_list = os.listdir(img_path)
path_list = [i for i in all_list if i.endswith('.jpg')]
path_list = np.array(path_list)
random.shuffle(path_list)

train_id = path_list[:int(len(path_list) * (1 - val_size - test_size))]
val_id = path_list[int(len(path_list) * (1 - val_size - test_size)):int(len(path_list) * (1 - test_size))]
test_id = path_list[int(len(path_list) * (1 - test_size)):]

for i in train_id:
    print(i)
    shutil.copy(f'{img_path}/{i}', f'{dataset_path}/images/train/{i}')
    shutil.copy(f'{label_path}/{i[:-4]}.txt', f'{dataset_path}/labels/train/{i[:-4]}.txt')

for i in val_id:
    print(i)
    shutil.copy(f'{img_path}/{i}', f'{dataset_path}/images/val/{i}')
    shutil.copy(f'{label_path}/{i[:-4]}.txt', f'{dataset_path}/labels/val/{i[:-4]}.txt')

for i in test_id:
    print(i)
    shutil.copy(f'{img_path}/{i}', f'{dataset_path}/images/test/{i}')
    shutil.copy(f'{label_path}/{i[:-4]}.txt', f'{dataset_path}/labels/test/{i[:-4]}.txt')