import xml.etree.ElementTree as ET
import os
import cv2

#{'hangAluminum', 'lackEdge', 'lackPole', 'dirty'}
rootdir = r'D:\doc\WeChat Files\wxid_krokt2g4mldb22\FileStorage\File\2023-03\PTC-IMG-XML\XML'
width = 2448
height = 2048
for file in os.listdir(rootdir):
    fw = open(r'D:\doc\WeChat Files\wxid_krokt2g4mldb22\FileStorage\File\2023-03\PTC-IMG-XML\txt\{}.txt'.format(file.split('.')[0]),'w')
    tree = ET.ElementTree(file=os.path.join(rootdir,file))
    tree.getroot()
    root = tree.getroot()
    for child_of_root in root:
        if child_of_root.tag == 'object':
            for child_of_root_ in child_of_root:
                if child_of_root_.tag == 'name':
                    if child_of_root_.text == 'hangAluminum':
                        fw.write('0 ')
                    if child_of_root_.text == 'lackEdge':
                        fw.write('1 ')
                    if child_of_root_.text == 'lackPole':
                        fw.write('2 ')
                    if child_of_root_.text == 'dirty':
                        fw.write('3 ')
                if child_of_root_.tag == 'bndbox':
                    for x in child_of_root_:
                        if x.tag == 'xmin':
                            xmin = int(x.text)
                        elif x.tag == 'ymin':
                            ymin = int(x.text)
                        elif x.tag == 'xmax':
                            xmax = int(x.text)
                        elif x.tag == 'ymax':
                            ymax = int(x.text)
                    w = xmax - xmin
                    h = ymax - ymin
                    x = (xmin + xmax) / 2
                    y = (ymin + ymax) / 2
                    x = str(x / width)
                    y = str(y / height)
                    w = str(w / width)
                    h = str(h / height)
                    fw.write(x + ' ' + y + ' ' + w + ' ' + h + '\n')
    fw.close()