python小工具 txt脚本转成xml数据集

import cv2
from pathlib import Path
from xml.dom.minidom import Document

img_dir = '//'  # 图片文件路径
txt_dir = '/t/'  # txt文件路径
xml_dir = '//'  # xml文件路径
labels = {'0': 'smog', '1': 'car'}  # 标签名

files = Path(txt_dir).glob('*txt')

for file in files:
    print(file)
    xmlBuilder = Document()
    annotation = xmlBuilder.createElement('annotation')  # 创建annotation标签
    xmlBuilder.appendChild(annotation)
    txtFile = open(file)
    txtList = txtFile.readlines()
    img = cv2.imread(img_dir + file.stem + '.jpg')
    h, w, d = img.shape

    filename = xmlBuilder.createElement('filename')  # filename标签
    filenamecontent = xmlBuilder.createTextNode(img_dir + file.stem + '.jpg')
    filename.appendChild(filenamecontent)
    annotation.appendChild(filename)  # filename标签结束

    size = xmlBuilder.createElement('size')  # size标签
    width = xmlBuilder.createElement('width')  # size子标签width
    widthcontent = xmlBuilder.createTextNode(str(w))
    width.appendChild(widthcontent)
    size.appendChild(width)  # size子标签width结束

    height = xmlBuilder.createElement('height')  # size子标签height
    heightcontent = xmlBuilder.createTextNode(str(h))
    height.appendChild(heightcontent)
    size.appendChild(height)  # size子标签height结束

    depth = xmlBuilder.createElement('depth')  # size子标签depth
    depthcontent = xmlBuilder.createTextNode(str(d))
    depth.appendChild(depthcontent)
    size.appendChild(depth)  # size子标签depth结束

    annotation.appendChild(size)  # size标签结束

    for line in txtList:
        oneline = line.strip().split(' ')
        object = xmlBuilder.createElement('object')  # object 标签
        picname = xmlBuilder.createElement('name')  # name标签
        namecontent = xmlBuilder.createTextNode(labels[oneline[0]])
        picname.appendChild(namecontent)
        object.appendChild(picname)  # name标签结束

        bndbox = xmlBuilder.createElement('bndbox')  # bndbox标签
        xmin = xmlBuilder.createElement('xmin')  # xmin标签
        mathData = int(((float(oneline[1])) * w) - (float(oneline[3])) * 0.5 * w)
        xminContent = xmlBuilder.createTextNode(str(mathData))
        xmin.appendChild(xminContent)
        bndbox.appendChild(xmin)  # xmin标签结束

        ymin = xmlBuilder.createElement('ymin')  # ymin标签
        mathData = int(((float(oneline[2])) * h) - (float(oneline[4])) * 0.5 * h)
        yminContent = xmlBuilder.createTextNode(str(mathData))
        ymin.appendChild(yminContent)
        bndbox.appendChild(ymin)  # ymin标签结束

        xmax = xmlBuilder.createElement('xmax')  # xmax标签
        mathData = int(((float(oneline[1])) * w) + (float(oneline[3])) * 0.5 * w)
        xmaxContent = xmlBuilder.createTextNode(str(mathData))
        xmax.appendChild(xmaxContent)
        bndbox.appendChild(xmax)  # xmax标签结束

        ymax = xmlBuilder.createElement('ymax')  # ymax标签
        mathData = int(((float(oneline[2])) * h) + (float(oneline[4])) * 0.5 * h)
        ymaxContent = xmlBuilder.createTextNode(str(mathData))
        ymax.appendChild(ymaxContent)
        bndbox.appendChild(ymax)  # ymax标签结束

        object.appendChild(bndbox)  # bndbox标签结束

        annotation.appendChild(object)  # object标签结束

    f = open(xml_dir + file.stem + '.xml', 'w')
    xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
    f.close()