"""
author:Hatcher
email:hcwang@semi.ac.cn
use:将aitod小目标检测数据集按照类别分开
input：
voc
    train_images
    train_annos
    ....
    test&val
output:
voc_class
    classname1
        images
        annos
    classname2
        iamges
        annos
        .....

"""
import os
from tqdm import tqdm
import  xml.dom.minidom
import shutil
source_path_root = "/data4/Hatcher/smallobject/voc/test_annos/"
det_path_root = "/data4/Hatcher/smallobject/voc_class/"
class_names  = ["airplane"]#os.listdir(det_path_root)
for class_name in class_names:
    for source_path_anno in tqdm(os.listdir(source_path_root)):
        source_path = os.path.join(source_path_root,source_path_anno)
        dom = xml.dom.minidom.parse(source_path)
        root = dom.documentElement
        objects=dom.getElementsByTagName("object")
        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(xmin.childNodes[0].data)
            ymin_data=int(ymin.childNodes[0].data)
            xmax_data=int(xmax.childNodes[0].data)
            ymax_data=int(ymax.childNodes[0].data)
            label_name=object.getElementsByTagName('name')[0].childNodes[0].data
            if not (label_name==class_name):
                 dom.documentElement.removeChild(object)
        objects=dom.getElementsByTagName("object")
        if not (len(objects)==0):
            #存储xml，复制图像
            xml_path = os.path.join(det_path_root,class_name,"annos",source_path_anno)
            fp = open(xml_path, 'w', encoding='utf-8')
            dom.writexml(fp, indent='', addindent='\t', newl='\n', encoding='utf-8')
            fp.close()
            shutil.copy(source_path.replace("test_annos","test_images").replace(".xml",".png"),os.path.join(det_path_root,class_name,"images",source_path_anno.replace(".xml",".png")))