【深度学习数据集的自动切分，分别生成训练，验证，测试的txt文件】

最新推荐文章于 2024-07-27 09:45:00 发布

小小小MaYi

最新推荐文章于 2024-07-27 09:45:00 发布

阅读量578

点赞数 2

文章标签：计算机视觉人工智能深度学习 opencv python

本文链接：https://blog.csdn.net/qq_37652891/article/details/130965341

版权

数据集的自动切分，分别生成训练，验证，测试的txt文件。

对文件夹中的图像数据进行切分，训练集：验证：测试集比例=8:1:1

1、图像数据（格式统一）

在这里插入图片描述

2、生成的切分txt文件(train、val、test、trainval)

在这里插入图片描述

3、txt文件中的每一行为图像文件的名称

在这里插入图片描述

4、代码如下

需要修改图像数据集和生成txt文件的路径，以及切分比例。

import os
import random
import argparse

parser = argparse.ArgumentParser()
#图像数据集路径
parser.add_argument('--xml_path', default=r'F:\nanana\522\labels', type=str, help='input xml label path')
#生成txt文件路径
parser.add_argument('--txt_path', default=r'F:\nanana\522', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 0.9
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
random.seed(2021)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()