YOLO系列 -- txt2xml脚本

YOLO系列 – txt2xml脚本

一般，训练YOLO算法的时候，我们通常需要把xml转换成txt格式标签文件，可以参考我之前写的这篇：
YOLO系列 — xml2txt脚本

---

但是，有时候也需要将txt反向转换为xml格式，话不多说，直接上代码：

from xml.dom import minidom
import cv2
import os
import glob

txt_dir = r'E:\work\qx_data\clean_all_data\obj_train_data/' #txt存放路径
img_dir = r"E:\work\qx_data\clean_all_data\obj_train_data/" #图片存放路径（可以与txt存放路径相同）
save_xml_dir = r"E:\work\qx_data\clean_all_data\xml/" #生成xml所保存的路径
cls = ["person"......] #类别列表

for i in glob.glob(txt_dir+"*.txt"):
    base_name = os.path.basename(i)[:-4]
    txt_dirtory = i
    jpg_dirtory = img_dir + base_name + ".jpg"
    if not os.path.exists(jpg_dirtory):
        print(jpg_dirtory + "not exist!!!!!!!")
    else:
        img_name = base_name + ".jpg"
        floder = "VOC2007"
        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)

        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(str(w)))
        size.appendChild(width)
        height = doc.createElement('height')
        height.appendChild(doc.createTextNode(str(h)))
        size.appendChild(height)
        depth = doc.createElement('depth')
        depth.appendChild(doc.createTextNode(str(d)))
        size.appendChild(depth)
        annotation.appendChild(size)


        txtLabel = open(txt_dirtory, 'r')
        boxes = txtLabel.readlines()
        for box in boxes:
            cls_idx, g_x, g_y, g_w, g_h = box.split(' ')
            t_x1 = (float(g_x) - 0.5 * float(g_w)) * w
            t_y1 = (float(g_y) - 0.5 * float(g_h)) * h
            t_x2 = (float(g_x) + 0.5 * float(g_w)) * w
            t_y2 = (float(g_y) + 0.5 * float(g_h)) * h
            class_name = cls[int(cls_idx)]
            
            object = doc.createElement('object')
            nm = doc.createElement('name')
            nm.appendChild(doc.createTextNode(str(class_name)))
            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("Unknown"))
            object.appendChild(truncated)
            difficult = doc.createElement('difficult')
            difficult.appendChild(doc.createTextNode('0'))
            object.appendChild(difficult)

            bndbox = doc.createElement('bndbox')
            # xmin ymin
            xmin = doc.createElement('xmin')
            xmin.appendChild(doc.createTextNode(str(int(t_x1))))
            bndbox.appendChild(xmin)

            ymin = doc.createElement('ymin')
            ymin.appendChild(doc.createTextNode(str(int(t_y1))))
            bndbox.appendChild(ymin)

            # xmax ymin
            xmax = doc.createElement('xmax')
            xmax.appendChild(doc.createTextNode(str(int(t_x2))))
            bndbox.appendChild(xmax)

            ymin = doc.createElement('ymax')
            ymin.appendChild(doc.createTextNode(str(int(t_y2))))
            bndbox.appendChild(ymin)

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