制作VOC格式数据集的train.txt、val.txt文件

文章目录[隐藏]

一、前言
二、txt生成程序

一、前言

在以前的文章中已经聊过 VOC数据集的组织结构 和 VOC格式数据集转yolo(darknet)格式。

当我们按照组织结构将自己的图片和xml标注文件放在指定文件夹下之后，在转换时我们还需要使用到VOC格式数据集中的train.txt、val.txt文件，这两个文件分别保存的是用来训练和验证的图片的名称，注意，仅仅是图片名。

二、`txt`生成程序

在程序中只需要设置好数据集文件夹的路径即可，也即saved_path变量，最后会在ImageSets/Main路径下生成所需的txt文件！！！

import os
from glob import glob
from sklearn.model_selection import train_test_split

#1.标签路径

saved_path = "./helmetVOC/"                #保存路径

#2.创建要求文件夹
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")
    

#3.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')

total_files = glob(saved_path + "Annotations/*.xml")

total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]

for file in total_files:
    ftrainval.write(file + "\n")

#split
train_files,val_files = train_test_split(total_files,test_size=0.15,random_state=42)

#train
for file in train_files:
    ftrain.write(file + "\n")
#val
for file in val_files:
    fval.write(file + "\n")

ftrainval.close()
ftrain.close()
fval.close()