import time
import torch
from ultralytics import YOLO
import numpy as np
import cv2
import os
from lxml.etree import Element, SubElement, tostring


def create_xml(boxs,img_shape,xml_path):
    """
    创建xml文件，依次写入xml文件必备关键字
    :param boxs:   txt文件中的box
    :param img_shape:   图片信息，xml中需要写入WHC
    :return:
    """
    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'Images'
    node_filename = SubElement(node_root, 'filename')
    node_filename.text = str(img_shape[3])
    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(img_shape[1])
    node_height = SubElement(node_size, 'height')
    node_height.text = str(img_shape[0])
    node_depth = SubElement(node_size, 'depth')
    node_depth.text = str(img_shape[2])

    if len(boxs)>=1:        # 循环写入box
        for box in boxs:
            node_object = SubElement(node_root, 'object')
            node_name = SubElement(node_object, 'name')
            # if str(list_[4]) == "person":                # 根据条件筛选需要标注的标签,例如这里只标记person这类，不符合则直接跳过
            #     node_name.text = str(list_[4])
            # else:
            #     continue
            node_name.text = str(box[4])
            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(box[0])
            node_ymin = SubElement(node_bndbox, 'ymin')
            node_ymin.text = str(box[1])
            node_xmax = SubElement(node_bndbox, 'xmax')
            node_xmax.text = str(box[2])
            node_ymax = SubElement(node_bndbox, 'ymax')
            node_ymax.text = str(box[3])

    xml = tostring(node_root, pretty_print=True)   # 格式化显示，该换行的换行

    file_name = img_shape[3].split(".")[0]
    filename = xml_path+"/{}.xml".format(file_name)

    f = open(filename, "wb")
    f.write(xml)
    f.close()


def draw_bounding_box(img, class_name, confidence, x, y, x_plus_w, y_plus_h,color):
    label = f'{class_name} ({confidence:.2f})'
    cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)
    cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)


def main(weights_path,img_path,xml_path,outputs_path):
    model = YOLO(weights_path)
    arr = torch.ones(1,3,224,224)
    result_init = model(arr)[0]                    
    colors = np.random.uniform(0, 255, size=(len(result_init.names), 3))       # 初始化画图颜色

    for name in os.listdir(img_path):
        t0 = time.time()
        original_image = cv2.imread(os.path.join(img_path,name))
        img_shape = (original_image.shape[0], original_image.shape[1], original_image.shape[2], name)

        # Use the model
        results = model(original_image)[0]  # predict on an image
        boxes = results.boxes.cpu().numpy()
        CLASSES = results.names
        xyxy_cls = []

        for i in range(len(boxes.xyxy)):
            xyxy_cls.append([int(boxes.xyxy[i][0]),int(boxes.xyxy[i][1]),int(boxes.xyxy[i][2]),int(boxes.xyxy[i][3]),CLASSES[boxes.cls[i]]])
            draw_bounding_box(original_image,CLASSES[boxes.cls[i]],boxes.conf[i],int(boxes.xyxy[i][0]),int(boxes.xyxy[i][1]),int(boxes.xyxy[i][2]),int(boxes.xyxy[i][3]),colors[int(boxes.cls[i])])

        if len(xyxy_cls) >0:
            create_xml(xyxy_cls,img_shape,xml_path)      # 创建xmls
        t1 = time.time()

        print("img name: {}  infer:{:4f} ms".format(name,(t1-t0)*1000))
        cv2.imwrite(os.path.join(outputs_path, name), original_image)          # 保留输出结果图

        # cv2.imshow('image', original_image)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()



if __name__ =="__main__":

    # weights_path =r"C:\Users\Administrator\Desktop\train24\weights\best.pt"
    weights_path =r"C:\Users\Administrator\Desktop\ultralytics-main\yolov8s.pt"        #权重地址
    imgs_path = r"C:\Users\Administrator\Desktop\yolov5-label-xml-main\inference\images"  #图片地址
    xmls_path = r"C:\Users\Administrator\Desktop\yolov5-label-xml-main\inference\xmlss"  #输出xml文件地址
    outputs_path = "./"


    main(weights_path,imgs_path,xmls_path,outputs_path)