制作自己的VOC数据集------生成ImageSets/Main txt文件

import os
import random
import xml
VOC_CLASSES = ['NG1', 'NG2', 'NG3', 'NG4', 'NG5', 'NG6',
               'NG7', 'NG8', 'NG9', 'NG10', 'NG11']

def generate_train_val_test_txt():
    xml_file_path = "D:/fruit_detection/NG_data/VOC2012/Annotations"  # xml文件路径
    save_Path = "D:/fruit_detection/NG_data/VOC2012/ImageSets/Main"
    ############################################3
    trainval_percent = 0.9
    train_percent = 0.8
    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(xmls_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编辑好了
    #################################################


#xmls_file_path = "D:/fruit_detection/NG_data/VOC2012/Annotations"  # xml文件路径
#save_Path = "D:/fruit_detection/NG_data/VOC2012/ImageSets/Main"

generate_train_val_test_txt()