PASCAL VOC格式的目标检测数据集生成ImageSets/Main中的各类txt文件

import os
import random
import xml
from xml.dom import minidom


VOC_CLASSES = ['face']


############################################
xml_file_path = "D:\dataset\cityscapes\cityscape_clean_car\Annotations\\" 
save_Path = ""
trainval_percent = 0
train_percent = 0
############################################


total_xml = os.listdir(xml_file_path) 
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
print("train and val size", tv)
print("train size", tr)
##############################################################
"""
将信息写入test.txt、train.txt、val.txt、trainval.txt
"""
ftrainval = open(os.path.join(save_Path, 'trainval.txt'), 'w')
ftest = open(os.path.join(save_Path, 'test.txt'), 'w')
ftrain = open(os.path.join(save_Path, 'train.txt'), 'w')
fval = open(os.path.join(save_Path, 'val.txt'), 'w')
for i in list:  # 第i个xml文件
        xml_name = total_xml[i][:-4]
        if i in trainval:
            ftrainval.write(xml_name + "\n")
            if i in train:
                ftrain.write(xml_name + "\n")
            else:
                fval.write(xml_name + "\n")
        else:
            ftest.write(xml_name + "\n")
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()


######################################################################
"""
将信息写入(class_name)_test.txt、(class_name)_train.txt、(class_name)_val.txt、(class_name)_trainval.txt
"""
for idx in range(len(VOC_CLASSES)):  # 每一个类单独处理
        class_name = VOC_CLASSES[idx]
        # 创建txt
        class_trainval = open(os.path.join(save_Path, str(class_name) + '_trainval.txt'), 'w')
        class_test = open(os.path.join(save_Path, str(class_name) + '_test.txt'), 'w')
        class_train = open(os.path.join(save_Path, str(class_name) + '_train.txt'), 'w')
        class_val = open(os.path.join(save_Path, str(class_name) + '_val.txt'), 'w')
        for k in list:
            xml_name = total_xml[k][:-4]  # xml的名称
            print(xml_name)
            xml_path = os.path.join(xml_file_path, xml_name + '.xml')
            ##################################################
            # 将获取的xml文件名送入到dom解析
            dom = xml.dom.minidom.parse(xml_path)  # 输入xml文件具体路径
            root = dom.documentElement
            # 获取xml object标签<name>
            object_name = root.getElementsByTagName('name')
            if len(object_name) > 0 and xml_name in object_name:  # 存在object（矩形框并且class_name在object_name列表中
                if k in trainval:
                    class_trainval.write(xml_name + ' ' + str(1) + "\n")
                    if k in train:
                        class_train.write(xml_name + ' ' + str(1) + "\n")
                    else:
                        class_val.write(xml_name + ' ' + str(1) + "\n")
                else:
                    class_test.write(xml_name + ' ' + str(1) + "\n")
            else:
                if k in trainval:
                    class_trainval.write(xml_name + ' ' + str(-1) + "\n")
                    if k in train:
                        class_train.write(xml_name + ' ' + str(-1) + "\n")
                    else:
                        class_val.write(xml_name + ' ' + str(-1) + "\n")
                else:
                    class_test.write(xml_name + ' ' + str(-1) + "\n")
        class_trainval.close()
        class_test.close()
        class_train.close()
        class_val.close()  # 1类的.txt编辑好了