# -*- coding: utf-8 -*-
"""
根据txt生成xml文件，txt格式为class xmin xmax ymin ymax
@author: bai
"""
import argparse
import os,cv2

try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET

class gen_xml():
    def __init__(self,root='annotation'):
        '''
        最多3层
        '''
        self.root_a = ET.Element(root)
        
        self.sub_root_a = None
        self.sub_sub_root_a = None
        
    def set_sub_node(self,last,sub_node,val):#last = root','sub_root' or 'sub_sub_root'
        if last == 'root':
            b = ET.SubElement(self.root_a, sub_node)
            b.text = val
        elif last == 'sub_root':
            b = ET.SubElement(self.sub_root_a, sub_node)
            b.text = val
        elif last == 'sub_sub_root':
            b = ET.SubElement(self.sub_sub_root_a, sub_node)
            b.text = val
            
    def set_sub_root(self,last,sub_root):#last = root','sub_root'
        if last == 'root':
            self.sub_root_a = ET.SubElement(self.root_a, sub_root)
        elif last == 'sub_root':
            self.sub_sub_root_a = ET.SubElement(self.sub_root_a, sub_root)
    def out(self,filename):
        fp = open(filename,'wb')
        tree = ET.ElementTree(self.root_a)
        tree.write(fp)
        fp.close()
        
def parse_args():
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description='txt2xml demo')
    parser.add_argument('--JPEGImages','-J', dest='image_path',type=str, help='JPEGImages path to read',
                        default='None')
    parser.add_argument('--txt','-T', dest='txt_path', help='txt path to read',
                        default='None')
    parser.add_argument('--Annotations','-A', dest='xml_path', help='xml path to save',
                        default='None')
    args = parser.parse_args()

    return args

def check_bbox(width,height,xmin,ymin,xmax,ymax):
    if xmin < 0:
        xmin = 0
    if xmin > width:
        xmin = width
    if xmax < 0:
        xmax = 0
    if xmax > width:
        xmax = width

    
    if ymin < 0:
        ymin = 0
    if ymin > height:
        ymin = height
    if ymax < 0:
        ymax = 0
    if ymax > height:
        ymax = height
    return xmin,ymin,xmax,ymax
def main(args):
    list_txt=os.listdir(args.txt_path)
#    list_need=['filename','size','name','part']
#    list_sub_need=['width','height','depth','xmin','ymin','xmax','ymax']
    if len(list_txt)==0:
        return
    
    for txt_name in list_txt:
        print(txt_name)
        txt=open(os.path.join(args.txt_path,txt_name),'r')
        image_name=os.path.join(args.image_path,txt_name.replace('.txt','.jpg'))
        img=cv2.imread(image_name)
        height=img.shape[0]
        width=img.shape[1]
        depth=img.shape[2]
        
        my_xml = gen_xml('annotation')
        my_xml.set_sub_node('root','filename','%s'%image_name)
        my_xml.set_sub_root('root','size')
        my_xml.set_sub_node('sub_root','%s'%'height','%s'%height)
        my_xml.set_sub_node('sub_root','%s'%'width','%s'%width)
        my_xml.set_sub_node('sub_root','%s'%'depth','%s'%depth)
        
        lines=txt.readlines()
        for line in lines:
            line=line.replace('\n','')
            class_name_coor=line.split(' ')
            xmin=int(float(class_name_coor[1]))
            ymin=int(float(class_name_coor[2]))
            xmax=int(float(class_name_coor[3]))
            ymax=int(float(class_name_coor[4]))
            
            xmin,ymin,xmax,ymax = check_bbox(width,height,xmin,ymin,xmax,ymax)
            if xmax - xmin <= 30 and ymax - ymin <= 30:
                continue
                
             
            if len(class_name_coor)>=5:
                my_xml.set_sub_root('root','object')
                my_xml.set_sub_node('sub_root','name',class_name_coor[0])
                my_xml.set_sub_root('sub_root','bndbox')
                my_xml.set_sub_node('sub_sub_root','%s'%'xmin','%d'%xmin)
                my_xml.set_sub_node('sub_sub_root','%s'%'ymin','%d'%ymin)
                my_xml.set_sub_node('sub_sub_root','%s'%'xmax','%d'%xmax)
                my_xml.set_sub_node('sub_sub_root','%s'%'ymax','%d'%ymax)
                
        if os.path.exists(image_name):
            my_xml.out(os.path.join(args.xml_path,txt_name.replace('.txt','.xml')))

if __name__ == '__main__':
    args = parse_args()
    if args.image_path == 'None' or args.xml_path == 'None':
        exit(0)
    main(args)
    print('end')