import os
import cv2
from lxml.etree import Element, SubElement, tostring
import xml.etree.ElementTree as ET


# 解析原始XML标注文件，获取人物的坐标信息
def parse_xml(xml_path,labels):

    tree = ET.parse(xml_path)
    root = tree.getroot()

    labeks_list = []
    for obj in root.findall('object'):

        try:

            name_labels = str(obj.find('name').text)
        
            if name_labels == labels:

                labels_dict = {}
                labels_dict['name'] = name_labels
                labels_dict['xmin'] = int(obj.find('bndbox/xmin').text)
                labels_dict['ymin'] = int(obj.find('bndbox/ymin').text)
                labels_dict['xmax'] = int(obj.find('bndbox/xmax').text)
                labels_dict['ymax'] = int(obj.find('bndbox/ymax').text)
                labeks_list.append(labels_dict)
        except:

            continue
    #如果有需要的信息，就返回列表
    if labeks_list:
        return labeks_list


# 筛选手部和头部信息
def select_infors(xml_path,person):

    infors_list = []

    # 分别获得头部和手部信息
    hands = parse_xml(xml_path=xml_path,labels='phone')
    head = parse_xml(xml_path=xml_path,labels='head')

    # 分别处理头部和手部信息
    if hands:

        for h,hand in enumerate(hands):

            hand_result = change_bbox(infor=hand,contrast=person)

            if hand_result:

                infors_list.append(hand_result) 

    if head:

        for h,head_infor in enumerate(head):

            head_result = change_bbox(infor=head_infor,contrast=person)

            if head_result:
                infors_list.append(head_result)
    else:
        pass       
    
    return infors_list


# 转换坐标值
def change_bbox(infor,contrast):
        
        if contrast['ymin'] <= infor['ymin'] and contrast['ymax'] >= infor['ymax'] and contrast['xmin'] <= infor['xmin'] and contrast['xmax'] >= infor['xmax']:

            # 截取图片后的相对坐标

            infor['ymin'] = infor['ymin'] - contrast['ymin']
            infor['ymax'] = infor['ymax']  - contrast['ymin']
            infor['xmin'] = infor['xmin'] - contrast['xmin']
            infor['xmax'] = infor['xmax']  - contrast['xmin']

            return infor

        else:
            pass


# 根据坐标值截图并保存
def cut_img_bbox(image_path,infors,num):

    # # 读取大图
    image = cv2.imread(image_path)

    # 取大一些范围的截图尺寸
    infors.update(ymin=infors['ymin'] - 25)
    infors.update(ymax=infors['ymax'] + 25)
    infors.update(xmin=infors['xmin'] - 25)
    infors.update(xmax=infors['xmax'] + 25)

    # 截取人物图片
    person_image = image[infors['ymin']:infors['ymax'], infors['xmin']:infors['xmax']]

    # 保存人物图片
    person_filename = os.path.basename(image_path).split('.')[0] +'_'+ f"person_{num}.jpg"
    # print(person_filename)
    person_image_path = os.path.join(output_img_folder, person_filename)
    cv2.imwrite(person_image_path, person_image)

    return person_filename




# 读取大图并截取人物图片
def extract_people(image_path, xml_path, output_img_folder, output_xml_folder):

    # 解析XML标注文件、获得人的标注信息
    people = parse_xml(xml_path=xml_path,labels='person')

    # 逐个截取人物图片
    for i, person in enumerate(people):

        # 截图保存并返回图片名
        person_filename = cut_img_bbox(image_path=image_path,infors=person,num=i)

        # 写xml
        node_root = Element('annotation')
        node_folder = SubElement(node_root, 'folder')
        node_folder.text = 'Images'
        node_filename = SubElement(node_root, 'filename')
        node_filename.text = str(person_filename)
        node_size = SubElement(node_root, 'size')
        node_width = SubElement(node_size, 'width')
        node_width.text = str(person['xmax'] - person['xmin'])
        node_height = SubElement(node_size, 'height')
        node_height.text = str(person['ymax'] - person['ymin'])
        node_depth = SubElement(node_size, 'depth')
        node_depth.text = str(3)        

        # 获得小图上的相对坐标
        select_infor_list = select_infors(xml_path=xml_path,person=person)

        if len(select_infor_list)>=1:        # 循环写入box
            for infors in select_infor_list:
                node_object = SubElement(node_root, 'object')
                node_name = SubElement(node_object, 'name')
                node_name.text = str(infors['name'])
                node_difficult = SubElement(node_object, 'difficult')
                node_difficult.text = '0'
                node_bndbox = SubElement(node_object, 'bndbox')
                node_xmin = SubElement(node_bndbox, 'xmin')
                node_xmin.text = str(int(infors['xmin']))
                node_ymin = SubElement(node_bndbox, 'ymin')
                node_ymin.text = str(int(infors['ymin']))
                node_xmax = SubElement(node_bndbox, 'xmax')
                node_xmax.text = str(int(infors['xmax']))
                node_ymax = SubElement(node_bndbox, 'ymax')
                node_ymax.text = str(int(infors['ymax']))
        else:
            pass
        xml = tostring(node_root, pretty_print=True)   # 格式化显示，该换行的换行

        # # 保存新的XML标注文件
        new_xml_filename = person_filename.split('.')[0] + '.xml'
        new_xml_path = os.path.join(output_xml_folder, new_xml_filename)
        new_xml = open(new_xml_path, "wb")
        new_xml.write(xml)
        new_xml.close()

        print(f"Saved person image {person_filename} and corresponding annotation {new_xml_filename}")


# 读取文件
def main(image_path,xml_path,output_img_folder,output_xml_folder):

    # 检查保存的文件路径
    print("frame image save path：{}".format(output_img_folder,output_xml_folder))
    os.makedirs(output_img_folder, exist_ok=True)
    os.makedirs(output_xml_folder, exist_ok=True)

    img_file_list = os.listdir(image_path)
    xml_file_list = os.listdir(xml_path)

    for img_file in img_file_list:

        img_file_full = os.path.join(image_path,img_file)

        xml_file_combination_name = img_file.split('.')[0] + '.xml'

        xml_file_combination = os.path.join(xml_path,xml_file_combination_name)\

        if xml_file_combination_name in xml_file_list:

            # 遍历文件处理数据
            extract_people(image_path=img_file_full, 
                           xml_path=xml_file_combination,
                           output_img_folder=output_img_folder,
                           output_xml_folder=output_xml_folder
                           )
        else:
            pass



if __name__ == '__main__':  

    # 原始大图像路径
    image_path = "E:/Images_Data/XZ/bank2/yolov8/images"

    # 原始XML标注文件路径
    xml_path = "E:/Images_Data/XZ/bank2/yolov8/XML_collect"

    # 输出文件夹路径
    output_img_folder = "E:/Images_Data/XZ/bank2/yolov8/output_files/images"
    output_xml_folder = 'E:/Images_Data/XZ/bank2/yolov8/output_files/Annotations'

    # parse_xml('E:/Images_Data/XZ/bank2/yolov8/XML_collect/Test_1_000014.xml')

    main(image_path=image_path,
                  xml_path=xml_path,
                  output_img_folder=output_img_folder,
                  output_xml_folder=output_xml_folder
                  )