import os.path

from lxml import etree
import xml.etree.ElementTree as ET
import cv2
import numpy as np

import xml.etree.ElementTree as ET


def remove_vest_objects(root):
    """
    递归地从XML树中删除所有name='Vest'的object标签及其子标签
    :param root: XML树的根元素
    """
    # 收集要删除的elements列表，因为不能在迭代时直接修改列表
    to_remove = []
    for elem in root.findall('.//object'):
        if elem.get('name') == 'Vest':
            to_remove.append(elem)

            # 逆序删除，因为删除会影响列表的索引
    for elem in reversed(to_remove):
        elem.getparent().remove(elem)

    # 解析XML文件


def save_PH(path, save_path):

    all_xml_name = os.listdir(path)
    all_xml_path = [os.path.join(path, t) for t in all_xml_name]

    for i in range(len(all_xml_path)):

        xml = all_xml_path[i]
        tree = ET.parse(xml)
        root = tree.getroot()

        j = 0
        while j < len(root):
            f = 0
            if root[j].tag == 'object':
                for child in root[j]:
                    if child.tag == 'name' and child.text not in ['Person', 'Helmet']:
                        del root[j]
                        f = 1
                        break
            if f == 0:
                j += 1



        tree.write(os.path.join(save_path, all_xml_name[i]), encoding='utf-8', xml_declaration=False)


if __name__ == '__main__':

    path = 'E:/haotian/YOLO安全帽手套检测数据集(含1000张图片)+对应voc、coco和yolo三种格式标签+划分脚本+训练教程/datasets/Annotations'
    save_path = 'E:/haotian/YOLO安全帽手套检测数据集(含1000张图片)+对应voc、coco和yolo三种格式标签+划分脚本+训练教程/datasets/Annotaions_PH'

    save_PH(path, save_path)