# -*- coding:utf-8 -*-
import os
import sys
import shutil
import xml.etree.ElementTree as ET

target_classes_list = ["didi", "group", "hello", "mobike", "ofo", "other", "rider"]
origin_dir = "G:\\combine_all"
result_dir = "G:\\combine_all_bike"


def get_xml_file_path_and_name(file_dir):
    """获取文件夹下所有xml的路径"""
    file_path_list = []
    file_name_list = []
    for root, dirs, files in os.walk(file_dir):
        for file in files:
            if os.path.splitext(file)[1] == '.xml':
                file_path_list.append(os.path.join(root, file))
                file_name_list.append(file)
    return file_path_list, file_name_list


def write_xml(tree, out_path):
    """
    将xml文件写出
    tree: xml树
    out_path: 写出路径
    """
    tree.write(out_path, encoding="utf-8", xml_declaration=True)


def init_dir(origin, result):
    origin_ann_dir = os.path.join(origin, "Annotations")
    origin_img_dir = os.path.join(origin, "JPEGImages")
    if not os.path.isdir(origin_ann_dir) or not os.path.isdir(origin_img_dir):
        raise Exception("origin data is missing!")

    result_ann_dir = os.path.join(result, "Annotations")
    result_img_dir = os.path.join(result, "JPEGImages")
    # 判断文件夹是否存在，不存在则创建
    if not os.path.isdir(result_ann_dir):
        os.makedirs(result_ann_dir)
    if not os.path.isdir(result_img_dir):
        os.makedirs(result_img_dir)
    # 判断文件夹是否为空，不为空则清空
    if not os.listdir(result_ann_dir):
        shutil.rmtree(result_ann_dir)
        os.mkdir(result_ann_dir)
    if not os.listdir(result_img_dir):
        shutil.rmtree(result_img_dir)
        os.mkdir(result_img_dir)

    return origin_ann_dir, origin_img_dir, result_ann_dir, result_img_dir


def copy_img(xml_name, origin, result):
    img_name = xml_name.replace(".xml", ".jpg")
    old_path = os.path.join(origin, img_name)
    new_path = os.path.join(result, img_name)
    shutil.copyfile(old_path, new_path)


def select(file_path_list, file_name_list):
    """选择所需要的类别标签，剔除多余标签"""
    for i in range(len(file_path_list)):
        file = file_path_list[i]
        try:
            tree = ET.parse(file)
            root = tree.getroot()  # 获得根节点
        except Exception:
            print("[Error] Cannot parse file %s" % file)
            sys.exit(1)
        object_list = root.findall('object')
        # 剔除多余标签
        for target in object_list:
            if not target.find('name').text in target_classes_list:
                root.remove(target)
        # 剔除完毕后，判断当前xml中是否还有object，如果没有，则跳过，如果有，则将xml和对应的图片保存到新的位置，
        if len(root.findall('object')) == 0:
            continue
        else:
            write_xml(tree, os.path.join(result_annotation_dir, file_name_list[i]))
            copy_img(file_name_list[i], origin_image_dir, result_image_dir)


if __name__ == '__main__':
    origin_annotation_dir, origin_image_dir, result_annotation_dir, result_image_dir = init_dir(origin_dir, result_dir)

    xml_path_list, xml_name_list = get_xml_file_path_and_name(origin_annotation_dir)

    select(xml_path_list, xml_name_list)
    print('Select finished!')