import os
import sys

from yoloBrokedRoad import YOLOBrokenRoad
from PIL import Image
from xml.dom.minidom import *

def infer_img(yolo, img_dir,out_dir):
    # sun
    try:
        os.path.exists(img_dir)
    except:
        print("dir error", img_dir)
    imgType_list = {'jpg','bmp','png','jpeg','rgb','tif'}
    img_list = os.listdir(img_dir)
    for imgname in img_list:
        if imgname.split(".")[-1] in imgType_list:
            img_path = os.path.join(img_dir, imgname)
            save_box_path = os.path.join(out_dir, imgname[:-4] + ".xml")
            # fp = open(save_box_path, "w")
            image = Image.open(img_path)
            # try:
            #     image = Image.open(img_path)
            #     print(img_path,"has been opened!")
            # except:
            #     print(img_path,"oped error!")
            boxes, confidence, classes  = yolo.detect_image(image)  # sun

            # box_str = ""
            # for box in boxes:
            #     box_str = box_str + ",".join(str(i) for i in box) + " "
            # print("img_path:", img_path, "boxes_str:", box_str)

            box_str = ",".join([str(a) for a in boxes])
            print("img_path:", img_path, "boxes_str:", box_str)
            # fp.write(imgname[:-4] + str(box_str))
            # fp.close()
            doc=Document()
            root=doc.createElement('bndboxs')
            doc.appendChild(root)

            for i in range(len(boxes)):
                box=boxes[i]
                object=doc.createElement('object')
                bndbox=doc.createElement('bndbox')

                xmin=doc.createElement('xmin')
                xmin.appendChild(doc.createTextNode(str(box[0])))
                ymin=doc.createElement('ymin')
                ymin.appendChild(doc.createTextNode(str(box[1])))
                xmax=doc.createElement('xmax')
                xmax.appendChild(doc.createTextNode(str(box[2])))
                ymax=doc.createElement('ymax')
                ymax.appendChild(doc.createTextNode(str(box[3])))
                bndbox.appendChild(xmin)
                bndbox.appendChild(ymin)
                bndbox.appendChild(xmax)
                bndbox.appendChild(ymax)

                score=doc.createElement('score')
                score.appendChild(doc.createTextNode(str(confidence[i])))
                object.appendChild(bndbox)
                object.appendChild(score)

                root.appendChild(object)
            fp=open(save_box_path,'w',encoding='utf-8')
            doc.writexml(fp,indent='',addindent='\t',newl='\n',encoding='utf-8')
            fp.close()
    sys.stdout.write("brokenRoadRunFinished\n")
    yolo.close_session()


if __name__ == '__main__':
    img_dir = r'E:\Work\2021\BrokenRoad\Track_A08\Test'
    model_path = r'E:\Code\Python\SuperMapOpenProject\deep-sort-yolov4\model_data\brokenRoad\brokenroad.h5'
    yolo_model = YOLOBrokenRoad(model_path)
    infer_img(yolo_model, img_dir,r'E:\Work\2021\BrokenRoad\Track_A08\Test_boxes')
    yolo_model1 = YOLOBrokenRoad(model_path)
    infer_img(yolo_model1, img_dir,r'E:\Work\2021\BrokenRoad\Track_A08\Test_boxes1')