"""
author: Hatcher
email:hcwang@semi.ac.cn
input:
Annotations_Original
JPEGImages_Original
w h 
output:
Annotations
Annotations
use: 将目标检测数据集中所有的图像resize为新的目标大小，并对应修改标注文件中的信息
"""
import os
import cv2
import  xml.dom.minidom
from tqdm import tqdm
from creat_xml import GEN_Annotations
anno_pathRoot = "/data4/Hatcher/airplane-03/small_object/test_data/800*800/Annotations_Original/"
image_pathRoot = anno_pathRoot.replace("/Annotations_Original/","/JPEGImages_Original/")
w = 800
h = 800

def DetResizeForTrain(anno_pathRoot,image_pathRoot,w,h):
    annoResultPath = anno_pathRoot.replace("/Annotations_Original/","/Annotations/")
    annoJPEGPath = image_pathRoot.replace("/JPEGImages_Original/", "/JPEGImages/")
    if not os.path.exists(annoResultPath):
        os.makedirs(annoResultPath)
    if not os.path.exists(annoJPEGPath):
        os.makedirs(annoJPEGPath)
    for annopath in tqdm(os.listdir(anno_pathRoot)):
        anno = GEN_Annotations(annopath.replace(".xml", ".jpg"))
        #print(os.path.join(anno_pathRoot, annopath))
        dom = xml.dom.minidom.parse(os.path.join(anno_pathRoot, annopath))
        size = dom.getElementsByTagName("size")
        # print(size[0].getElementsByTagName('width')[0].childNodes[0].data)
        width = size[0].getElementsByTagName('width')[0].childNodes[0].data
        height = size[0].getElementsByTagName('height')[0].childNodes[0].data
        depth = size[0].getElementsByTagName('depth')[0].childNodes[0].data
        anno.set_size(int(w), int(h), int(depth))
        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]

            xmin_data = int(int(xmin.childNodes[0].data)*(w/int(width)))
            ymin_data = int(int(ymin.childNodes[0].data)*(h/int(height)))
            xmax_data = int(int(xmax.childNodes[0].data)*(w/int(width)))
            ymax_data = int(int(ymax.childNodes[0].data)*(h/int(height)))
            label_name = object.getElementsByTagName('name')[0].childNodes[0].data
            anno.add_pic_attr(label=label_name, xmin=xmin_data, xmax=xmax_data, ymin=ymin_data, ymax=ymax_data)
        result_annoPath = os.path.join(annoResultPath, annopath.replace("-","_"))
        anno.savefile(result_annoPath)
        image = cv2.imread(os.path.join(image_pathRoot,annopath.replace(".xml",".jpg")))
        imageResize = cv2.resize(image,(w,h))
        cv2.imwrite(os.path.join(annoJPEGPath,annopath.replace(".xml",".png").replace("-","_")),imageResize)

if __name__=="__main__":
    DetResizeForTrain(anno_pathRoot,image_pathRoot,w,h)