'''
author ：  Hatcher W
Email： hcwang@semi.ac.cn
use: 用于检测voc数据集是否正确，
input:（图像的文件夹路径，图像标签文件夹路径）
output：check路径，所有的可视化结果都在里面
'''
import os
import cv2
import  xml.dom.minidom
from tqdm import tqdm
from lxml import etree
image_path="./JPEGImages/"
annotation_path="/data/Hatcher/Fine_HC/anno_xml/"
save_path ="./check/"
num = 0
files_name = os.listdir(annotation_path)
class GEN_Annotations:
    def __init__(self, filename):
        self.root = etree.Element("annotation")

        child1 = etree.SubElement(self.root, "folder")
        child1.text = "Hatcher"

        child2 = etree.SubElement(self.root, "filename")
        child2.text = filename

        child3 = etree.SubElement(self.root, "source")

        child4 = etree.SubElement(child3, "annotation")
        child4.text = "MultiModel"
        child5 = etree.SubElement(child3, "database")
        child5.text = "Unknown"

        child6 = etree.SubElement(child3, "image")
        child6.text = "flickr"
        child7 = etree.SubElement(child3, "flickrid")
        child7.text = "35435"

    def set_size(self, witdh, height, channel):
        size = etree.SubElement(self.root, "size")
        widthn = etree.SubElement(size, "width")
        widthn.text = str(witdh)
        heightn = etree.SubElement(size, "height")
        heightn.text = str(height)
        channeln = etree.SubElement(size, "depth")
        channeln.text = str(channel)

    def savefile(self, filename):
        tree = etree.ElementTree(self.root)
        tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')

    def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
        object = etree.SubElement(self.root, "object")
        namen = etree.SubElement(object, "name")
        namen.text = label
        bndbox = etree.SubElement(object, "bndbox")
        xminn = etree.SubElement(bndbox, "xmin")
        xminn.text = str(xmin)
        yminn = etree.SubElement(bndbox, "ymin")
        yminn.text = str(ymin)
        xmaxn = etree.SubElement(bndbox, "xmax")
        xmaxn.text = str(xmax)
        ymaxn = etree.SubElement(bndbox, "ymax")
        ymaxn.text = str(ymax)


def make_xml(image_path,shape0,shape1,x_left, y_left, x_right, y_right):


    print(image_path)
    print("\n")
    anno = GEN_Annotations(image_path)
    anno.set_size(shape0, shape1, 3)
    anno.add_pic_attr("bird", x_left, y_left, x_right, y_right)
    anno.savefile(image_path)
for filename_ in tqdm(files_name):
    xml_path = os.path.join(annotation_path,filename_)
    try:
        dom = xml.dom.minidom.parse(xml_path)
    except:
        #os.remove(img_path)
        continue
    root = dom.documentElement
    objects=dom.getElementsByTagName("object")
    shapes = dom.getElementsByTagName("size")
    for shape in shapes:
        width = shape.getElementsByTagName('width')[0].childNodes[0].data
        height = shape.getElementsByTagName('height')[0].childNodes[0].data
    for object in objects:
        bndbox = object.getElementsByTagName('bndbox')[0]
        xmin = bndbox.getElementsByTagName('xmin')[0]
        ymin = bndbox.getElementsByTagName('ymin')[0]
        xmax = bndbox.getElementsByTagName('xmax')[0]
        ymax = bndbox.getElementsByTagName('ymax')[0]
        #print(xmin.childNodes[0].data)
        #print(type(xmin.childNodes[0].data))
        xmin_data=int(float(xmin.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
        ymin_data=int(float(ymin.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
        xmax_data=int(float(xmax.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
        ymax_data=int(float(ymax.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
        label_name=object.getElementsByTagName('name')[0].childNodes[0].data
        # cv2.rectangle(img,(xmin_data,ymin_data),(xmax_data,ymax_data),(55,255,155),1)
        # num +=1
        # cv2.putText(img,label_name,(xmin_data,ymin_data-3),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)
    make_xml(xml_path, width, height, xmin_data, ymin_data, xmax_data, ymax_data)
    flag=0
    # flag=cv2.imwrite(save_path+"{}.jpg".format(filename),img)
    # if not (flag):
    #     print(filename,"error")
print("all done ====================================")
print("number = ",num)

#(int((xmin_data+xmax_data/2)),int((ymin_data+ymax_data)/2))