from lxml import etree
import cv2
import numpy as np


# 读取 xml 文件信息，并返回字典形式
def parse_xml_to_dict(xml):
    if len(xml) == 0:  # 遍历到底层，直接返回 tag对应的信息
        return {xml.tag: xml.text}

    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个，所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


# xml 标注文件的可视化
def xmlShow(img, xml, save=True):
    image = cv2.imdecode(np.fromfile(img, dtype=np.uint8), -1)
    with open(xml, encoding='gb18030', errors='ignore') as fid:  # 防止出现非法字符报错
        xml_str = fid.read()
    xml = etree.fromstring(xml_str)
    data = parse_xml_to_dict(xml)["annotation"]  # 读取 xml文件信息

    ob = []  # 存放目标信息
    try:
        for i in data['object']:  # 提取检测框
            name = str(i['name'])  # 检测的目标类别

            bbox = i['bndbox']
            xmin = int(bbox['xmin'])
            ymin = int(bbox['ymin'])
            xmax = int(bbox['xmax'])
            ymax = int(bbox['ymax'])

            tmp = [name, xmin, ymin, xmax, ymax]  # 单个检测框
            ob.append(tmp)
    except:
        return

    # 绘制检测框
    for name, x1, y1, x2, y2 in ob:
        cv2.rectangle(image, (x1, y1), (x2, y2), color=(255, 0, 0), thickness=2)  # 绘制矩形框
        cv2.putText(image, name, (x1, y1 - 10), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                    fontScale=0.5, thickness=1, color=(0, 0, 255))

    # 保存图像
    if save:
        cv2.imwrite('test_images/result1.jpg', image)

    # # 展示图像
    # cv2.imshow('test', image)
    # cv2.waitKey()
    # cv2.destroyAllWindows()


if __name__ == "__main__":

    # #-------------------------------- 数据集1----------------------------------------
    # img_path = 'E:/haotian/YOLO安全帽手套检测数据集(含1000张图片)+对应voc、coco和yolo三种格式标签+划分脚本+训练教程/datasets/images/233.jpg'  # 传入图片
    #
    # labels_path = 'E:/haotian/YOLO安全帽手套检测数据集(含1000张图片)+对应voc、coco和yolo三种格式标签+划分脚本+训练教程/datasets/Annotaions_PH/233.xml'
    # #-----------------------------------end1-------------------------------------

    labels_path = 'fire_smoke/annotations/smoke_000222.xml'
    img_path = 'fire_smoke/images/smoke_000222.jpg'

    xmlShow(img=img_path, xml=labels_path, save=True)