txt2xml.py文件

from xml.dom import minidom
import cv2
import os

path = '/home/hy/work/OCR_Detection_Simple/data/label'
setname = set([i.split('.')[0] for i in os.listdir(path)])
for i in setname:
    txt_dirtory = os.path.join(path, i + '.txt')
    jpg_dirtory = os.path.join(path, i + '.jpg')
    img_name = jpg_dirtory.split('/')[-1]
    floder = jpg_dirtory.split('/')[-2]
    im = cv2.imread(jpg_dirtory)
    w = im.shape[1]
    h = im.shape[0]
    d = im.shape[2]
    # print w,h,d
    print('i=', i)
    doc = minidom.Document()  # 创建DOM树对象

    annotation = doc.createElement('annotation')  # 创建子节点
    doc.appendChild(annotation)  # annotation作为doc树的子节点

    folder = doc.createElement('folder')
    folder.appendChild(doc.createTextNode(floder))  # 文本节点作为floder的子节点
    annotation.appendChild(folder)  # folder作为annotation的子节点

    filename = doc.createElement('filename')
    filename.appendChild(doc.createTextNode(img_name))
    annotation.appendChild(filename)

    filename = doc.createElement('path')
    filename.appendChild(doc.createTextNode(jpg_dirtory))
    annotation.appendChild(filename)

    source = doc.createElement('source')
    database = doc.createElement('database')
    database.appendChild(doc.createTextNode("Unknown"))
    source.appendChild(database)
    annotation.appendChild(source)


    size = doc.createElement('size')
    width = doc.createElement('width')
    width.appendChild(doc.createTextNode("%d" % w))
    size.appendChild(width)
    height = doc.createElement('height')
    height.appendChild(doc.createTextNode("%d" % h))
    size.appendChild(height)
    depth = doc.createElement('depth')
    depth.appendChild(doc.createTextNode("%d" % d))
    size.appendChild(depth)
    annotation.appendChild(size)

    segmented = doc.createElement('segmented')
    segmented.appendChild(doc.createTextNode("0"))
    annotation.appendChild(segmented)

    txtLabel = open(txt_dirtory, 'r')
    boxes = txtLabel.readlines()
    for box in boxes:
        # print(box)
        box = box.replace('\n', '')
        mark = box.split(' ')[0]
        box = box.split(' ')[-1].split(',')
        # print(mark,box)
        # print(box)
        object = doc.createElement('object')
        nm = doc.createElement('name')
        nm.appendChild(doc.createTextNode(mark))
        object.appendChild(nm)
        pose = doc.createElement('pose')
        pose.appendChild(doc.createTextNode("Unspecified"))
        object.appendChild(pose)
        truncated = doc.createElement('truncated')
        # truncated.appendChild(doc.createTextNode("1"))
        truncated.appendChild(doc.createTextNode(mark))
        object.appendChild(truncated)
        difficult = doc.createElement('difficult')
        difficult.appendChild(doc.createTextNode('0'))
        object.appendChild(difficult)

        #给定两个点的坐标  simple

        bndbox = doc.createElement('bndbox')
        # xmin ymin
        xmin = doc.createElement('xmin')
        xmin.appendChild(doc.createTextNode(box[0]))
        bndbox.appendChild(xmin)

        ymin = doc.createElement('ymin')
        ymin.appendChild(doc.createTextNode(box[1]))
        bndbox.appendChild(ymin)

        # xmax ymin
        xmax = doc.createElement('xmax')
        xmax.appendChild(doc.createTextNode(box[2]))
        bndbox.appendChild(xmax)

        ymin = doc.createElement('ymin')
        ymin.appendChild(doc.createTextNode(box[1]))
        bndbox.appendChild(ymin)

        # xmax ymax
        xmax = doc.createElement('xmax')
        xmax.appendChild(doc.createTextNode(box[2]))
        bndbox.appendChild(xmax)

        ymax = doc.createElement('ymax')
        yma = str(int(box[1]) + int(box[3]))
        ymax.appendChild(doc.createTextNode(yma))
        bndbox.appendChild(ymax)


        # xmin ymax
        xmin = doc.createElement('xmin')
        xmin.appendChild(doc.createTextNode(box[0]))
        bndbox.appendChild(xmin)

        ymax = doc.createElement('ymax')
        yma = str(int(box[1]) + int(box[3]))
        ymax.appendChild(doc.createTextNode(yma))
        bndbox.appendChild(ymax)

        #+++++++++++++++++++++++++++++++++++++++++++++++++
        object.appendChild(bndbox)
        annotation.appendChild(object)
        savefile = open(os.path.join(path, i + '.xml'), 'w')
        savefile.write(doc.toprettyxml())
        savefile.close()