import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join


passimg = []

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(xmlname,xmlpath,txtpath):
    if not os.path.exists(txtpath):
        os.makedirs(txtpath)
    xmlfile = os.path.join(xmlpath,xmlname)
    with open(xmlfile, "r") as in_file:
      txtname = xmlname[:-4]+'.txt'
      txtfile = os.path.join(txtpath,txtname)

      tree = ET.parse(in_file)
      root = tree.getroot()
      size = root.find('size')
      w = int(size.find('width').text)
      h = int(size.find('height').text)
      if w == 0 or h == 0:
          passimg.append(xmlfile)
      else:
        with open(txtfile, "w+") as out_file:
            out_file.truncate()
            for obj in root.iter('object'):
                difficult = obj.find('difficult').text
                cls = obj.find('name').text
                if cls not in classes or int(difficult)==1:
                    continue
                cls_id = classes.index(cls)
                xmlbox = obj.find('bndbox')
                b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
                bb = convert((w,h), b)
                out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


def xml2Yolo():
    """
    xml数据格式转yolo
    :return:
    """
    list = os.listdir(xmlpath)

    for i in range(0, len(list)):
        path = os.path.join(xmlpath, list[i])
        if ('.xml' in path) or ('.XML' in path):
            convert_annotation(list[i], xmlpath,txtpath)
            print('done', i)
        else:
            print('not xml file', i)
    print("未成功转换的文件：\n", passimg)



if __name__ == "__main__":
    classes = ["111","222"]
    xmlpath = r'D:\work\WaterAffairs\labelDatas\labelV5\VOCdevkit2007\VOC2007\Annotations'
    # txtpath = r"D:\work\WaterAffairs\labelDatas\labelV3\VOCdevkit2007\VOC2007\labels"
    txtpath = None
    if txtpath is None:
        rootpath = os.path.split(xmlpath)[0]
        txtpath = os.path.join(rootpath,'labels')
    xml2Yolo()