# DESC: 標(biāo)注單類(lèi)數(shù)據(jù)后，生成voc訓(xùn)練格式腳本
import os
import random
import glob
import shutil

trainval_percent = 0.66
train_percent = 0.5
base_path = r'C:\Users\wzauw\Desktop\cougar_body'
xml_path = r''
img_path = r''
total_xml = []


def mk_data_folder(*args, cur=None):
    if not cur:
        cur = base_path
    for idx, arg in enumerate(args, 1):
        cur = os.path.join(cur, arg)
        if idx == len(args):
            if not os.path.exists(cur):
                os.makedirs(cur)
    return cur


if __name__ == '__main__':
    # 創(chuàng)建文件目錄
    folder_path = {}
    voc2007 = mk_data_folder('VOCdevkit', 'VOC2007')
    for foldername in ['Annotations', 'ImageSets', 'JPEGImages']:
        folder_path[foldername] = mk_data_folder(foldername, cur=voc2007)
    for imgsetname in ['Action', 'Layout', 'Main', 'Segmentation']:
        folder_path[imgsetname] = mk_data_folder(imgsetname, cur=folder_path['ImageSets'])

    if not img_path:
        img_path = base_path
    for file in glob.glob(img_path + '/*.*[gG]'):
        if os.path.exists(os.path.join(folder_path['JPEGImages'], os.path.basename(file))):
            os.remove(file)
        else:
            shutil.move(file, folder_path['JPEGImages'])

    if not xml_path:
        xml_path = base_path
    for file in glob.glob(xml_path + '/*.xml'):
        total_xml.append(os.path.basename(file))
        if os.path.exists(os.path.join(folder_path['Annotations'], os.path.basename(file))):
            os.remove(file)
        else:
            shutil.move(file, folder_path['Annotations'])

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

    ftrainval = open(os.path.join(folder_path['Main'], 'trainval.txt'), 'w')
    ftest = open(os.path.join(folder_path['Main'], 'test.txt'), 'w')
    ftrain = open(os.path.join(folder_path['Main'], 'train.txt'), 'w')
    fval = open(os.path.join(folder_path['Main'], 'val.txt'), 'w')

    for i in range(num):
        name = total_xml[i][:-4] + '\n'
        if i in trainval:
            ftrainval.write(name)
            if i in train:
                ftrain.write(name)
            else:
                fval.write(name)
        else:
            ftest.write(name)

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