import os
import xml.etree.ElementTree as ET
import cv2

dirpath = r'E:\Train\EM14XML'  # 原来存放xml文件的目录
newdir = r'E:\Train\EM14TXT'  # 修改label后形成的txt目录

imagepath = r"E:\Train\EM14IMG"

if not os.path.exists(newdir):
    os.makedirs(newdir)

for fp in os.listdir(dirpath):
    print(fp)
    print(fp.split('.')[0])

    imagename = fp.split('.')[0] + '.tiff'
    imagename_path = os.path.join(imagepath,imagename)
    print(imagename_path)
    image = cv2.imread(imagename_path)
    # cv2.imshow('image',image)
    # cv2.waitKey(0)

    if fp == '.DS_Store':
        continue
    root = ET.parse(os.path.join(dirpath, fp)).getroot()   #获取xml的根节点
    xmin, ymin, xmax, ymax = 0, 0, 0, 0
    sz = root.find('size')
    print(sz)
    width = float(sz[0].text)
    height = float(sz[1].text)
    filename = root.find('filename').text
    for child in root.findall('object'):  # 找到图片中的所有框
        print(child)
        sub = child.find('bndbox')  # 找到框的标注值并进行读取
        xmin = float(sub[0].text)
        ymin = float(sub[1].text)
        xmax = float(sub[2].text)
        ymax = float(sub[3].text)
        try:  # 转换成yolov3的标签格式，需要归一化到（0-1）的范围内
            x_center = (xmin + xmax) / (2 * width)
            y_center = (ymin + ymax) / (2 * height)
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height
        except ZeroDivisionError:
            print(filename, '的 width有问题')
        sub_class = str(child.find('name').text)
        if sub_class == 'EM17':
            cls = 2
        elif sub_class == 'EM14':
            cls = 0
        elif sub_class == 'EM170':
            cls = 3
        elif sub_class == 'EM18':
            cls = 1
        elif sub_class == 'EM19':
            cls = 4
        elif sub_class == 'EM190':
            cls = 5
        elif sub_class == 'EM20':
            cls = 6
        elif sub_class == 'EM200':
            cls = 7
        elif sub_class == 'EM201':
            cls = 8
        elif sub_class == 'EM202':
            cls = 9
        elif sub_class == 'EM203':
            cls = 10
        elif sub_class == 'EM180':
            cls = 11
        elif sub_class == 'EM141':
            cls = 12
        else:
            continue

            # cv.namedWindow('rect', 1)
            # cv.imshow('rect', image)
            # cv.waitKey(0)

        with open(os.path.join(newdir, fp.split('.')[0] + '.txt'), 'a+') as f:
            f.write(' '.join([str(cls), str(x_center), str(y_center), str(w), str(h) + '\n']))


        cv2.rectangle(image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255), 2)
        cv2.putText(image, str(cls), (200, 200), cv2.FONT_HERSHEY_COMPLEX, 5.0, (100, 200, 200), 2)
    cv2.imwrite('E:\\Train\\EM14Result\\' + imagename, image)