# # ## # # # # # # # # # #using: 将yolo txt label 转换成xml标签
# #-*- coding: utf-8 -*-

import cv2
import os

xml_head = '''<annotation>
    <folder>VOC2007</folder>
    <filename>{}</filename>.
    <source>
    </source>
    <owner>
        <flickrid>null</flickrid>
        <name>null</name>
    </owner>
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>{}</depth>
    </size>
    <segmented>0</segmented>
    '''
xml_obj = '''
    <object>
        <name>{}</name>
        <pose>Rear</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
    '''
xml_end = '''
</annotation>'''

# labels = ['mouse','snake','lizard']  # label for datasets
# labels = ['mouse','snake','lizard','fire','smoke','open','close','R_0','R_1','G_0','G_1','Y_0','Y_1','mask','person','hat','reflect']  #  需要训练的所有类，label for datasets
# labels = ['1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16','17','18','19','20','21','22','23','24','25','26','27','28','29','30']
labels = ['1', '2', '3', '4']
cnt = 0
txt_path = 'txt2xml/labels/'
xml_path = 'txt2xml/xmls/'
img_path = 'txt2xml/images/'
os.makedirs(xml_path, exist_ok=True)
for yolo_file in os.listdir(txt_path):
    jpg = img_path + yolo_file.replace('.txt', '.png')  # image path
    txt = txt_path + yolo_file  # yolo label txt path
    xml = xml_path + yolo_file.replace('.txt', '.xml')  # xml save path
    print(jpg)
    obj = ''

    img = cv2.imread(jpg)
    img_h, img_w = img.shape[0], img.shape[1]
    head = xml_head.format(str(jpg), str(img_w), str(img_h), str(3))
    with open(txt, 'r') as f:
        for line in f.readlines():
            yolo_datas = line.strip().split(' ')
            label = int(float(yolo_datas[0].strip()))
            center_x = round(float(str(yolo_datas[1]).strip()) * img_w)
            center_y = round(float(str(yolo_datas[2]).strip()) * img_h)
            bbox_width = round(float(str(yolo_datas[3]).strip()) * img_w)
            bbox_height = round(float(str(yolo_datas[4]).strip()) * img_h)

            xmin = str(int(center_x - bbox_width / 2))
            ymin = str(int(center_y - bbox_height / 2))
            xmax = str(int(center_x + bbox_width / 2))
            ymax = str(int(center_y + bbox_height / 2))

            obj += xml_obj.format(labels[label], xmin, ymin, xmax, ymax)
    with open(xml, 'w') as f_xml:
        f_xml.write(head + obj + xml_end)
    cnt += 1
    print(cnt)