43.人工智能——提取标注文件中的目标名称、生成和划分数据集

在做目标检测任务的时候，先要用labelme、labelimg、Vott等标注工具对目标进行标注。具体用法可以参看：24.人工智能：计算机视觉任务——数据格式和标注。

在训练的时候，我们还需要对数据集进行划分。这里以交通标志检测为例，看如何代码提取标注文件中的目标名称、生成和划分数据集。

数据集目录：traffic

包含两个子目录：Annotations 和 JPEGImages，分别存放xml标注文件和图像文件，主文件名分别一一对应。

实现完整代码

import xml.etree.ElementTree as ET
import os
import random

basedir="traffic"
xmldir="annotations"
imgdir="jpegimages"
classname=[]
datalst=[]

#获取标注文件中的目标名称
for f in os.listdir(os.path.join(basedir,xmldir)):
    xmlfile=os.path.join(basedir,xmldir,f)
    tree=ET.parse(xmlfile)
    root=tree.getroot() #获取根节点
    for child in root:
        for children in child:
            if children.tag=="name":
                classname.append(children.text)

classname=set(classname)
#print(classname)
labelsfile=os.path.join(basedir,"labels.txt")
with open(labelsfile,"w") as f:
    for label in classname:
        f.write(label+"\n")


              
for fimg,fxml in zip(os.listdir(os.path.join(basedir,imgdir)),os.listdir(os.path.join(basedir,xmldir))):
    imgfile=os.path.join(basedir,imgdir,fimg)
    xmlfile=os.path.join(basedir,xmldir,fxml)
    datalst.append([imgfile,xmlfile])    
       

#数据乱序，划分训练集：验证集：测试集=7：2：1
random.shuffle(datalst)
trainlst=datalst[:int(len(datalst)*0.7)]
vallst=datalst[int(len(datalst)*0.7):int(len(datalst)*0.9)] 
testlst=datalst[int(len(datalst)*0.9):]

train=os.path.join(basedir,"train_list.txt")
val=os.path.join(basedir,"val_list.txt")
test=os.path.join(basedir,"test_list.txt")

#写入训练集
with open(train,"w") as f:
    for d in trainlst:
        img=d[0]
        xml=d[1]
        f.write(img+" "+xml+"\n")
        
#写入验证集
with open(val,"w") as f:
    for d in vallst:
        img=d[0]
        xml=d[1]
        f.write(img+" "+xml+"\n")
        
#写入测试集
with open(test,"w") as f:
    for d in testlst:
        img=d[0]
        xml=d[1]
        f.write(img+" "+xml+"\n")
    

运行后，会在数据集目录下生成4个文本文件，分别是：labels.txt、train_list.txt、val_list.txt 、test_list.txt，用于后面的数据训练。