from add_xml import add_xml
from create_xml import create_xml
from xz_yolov8_atm.yolov8_det import analysis_yolov8
import glob
import cv2
import os
from ultralytics import YOLO


if __name__ == '__main__': 

    imgpath = 'E:/BANK_xz_all/imgpath'
    xmlpath = 'E:/BANK_xz_all/xz_yolov8_atm/output'

    imglist = os.listdir(imgpath)
    xmllist = os.listdir(xmlpath)

    model_all = YOLO('xz_yolov8_atm/model_files/best(3)_all.pt')
    
    for img in imglist:

        images_path = os.path.join(imgpath,img)
        print(images_path)
        xml_path = os.path.join(xmlpath,img.split('.')[0] + '.xml')
        print(xml_path)

        images = cv2.imread(images_path)
        # print(images)

        results = analysis_yolov8(images=images,
                                    model_coco=model_all,
                                    confidence=0.5)
            
        w,h,d = images.shape
        img_shape = (w,h,d,img)

        if os.path.isfile(xml_path):

            add_labels = add_xml(inforsDict=results,
                                     xmlFilePath=xml_path)
        else:
            create_new = create_xml(boxs=results,
                                        img_shape=img_shape,
                                        xml_path=xmlpath)