VisDrone数据集txt格式标注更改xml

txt格式，每行从左往右分别,,,,,,,
分类：

pedestrian (1), people (2), bicycle (3), car (4), van (5),truck (6), tricycle (7), awning-tricycle (8), bus (9), motor (10), others (11)

代码，从一个文件夹中读取jpg文件，另一个文件夹读标注，汇总成xml格式

# -*- coding:utf-8 -*-
import os
import xml.etree.ElementTree as ET
import cv2


# 增加换行符
def __indent(elem, level=0):
    i = "\n" + level * "\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            __indent(elem, level + 1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


def __makeXML(shape, realName, f, savePath = "E:/dataset/UAV_for_counting/"):
    root = ET.Element('annotation')  # 创建节点
    tree = ET.ElementTree(root)  # 创建文档

    filename = ET.Element('filename')
    filename.text = txt_path
    root.append(filename)
    # 图像尺寸
    size = ET.Element('size')
    width = ET.Element('width')
    height = ET.Element('height')
    depth = ET.Element('depth')
    width.text = str(shape[0])
    height.text = str(shape[1])
    depth.text = '3'
    size.append(width)
    size.append(height)
    size.append(depth)
    root.append(size)

    # 逐行读txt
    while True:
        line = f.readline()
        if line:
            # object节点
            object = ET.Element('object')
            name = ET.Element('name')
            pose = ET.Element('pose')
            difficult = ET.Element('Difficult')
            bndbox = ET.Element('bndbox')
            xmin = ET.Element('xmin')
            ymin = ET.Element('ymin')
            xmax = ET.Element('xmax')
            ymax = ET.Element('ymax')
            truncated = ET.Element('truncated')

            s = line.split(',')

            if s[5] == '1' or s[5] == '2':
                name.text = 'pedestrian'
            elif s[5] == '4' or s[5] == '5' or s[5] == '6' or s[5] == '9':
                name.text = 'vehicle'
            else:
                continue

            xmin.text = s[0]
            ymin.text = s[1]
            xmax.text = str(int(s[0]) + int(s[2]))
            ymax.text = str(int(s[1]) + int(s[3]))

            pose.text = 'Unspecified'
            difficult.text = '0'
            truncated.text = s[6]

            bndbox.append(xmin)
            bndbox.append(ymin)
            bndbox.append(xmax)
            bndbox.append(ymax)
            object.append(name)
            object.append(pose)
            object.append(difficult)
            object.append(truncated)
            object.append(bndbox)
            root.append(object)

        else:
            break

    __indent(root)  # 增加换行符
    tree.write(savePath + realName + ".xml", encoding='utf-8', xml_declaration=False)


img_path = "E:/dataset/UAV_for_counting/"
txt_path = "E:/数据集/VisDrone2019-DET-train/annotations/"
filelist = os.listdir(img_path)

for i, filename in enumerate(filelist):
    if filename.endswith(".jpg"):
        img = cv2.imread(img_path + filename)
        size = img.shape

        realName = filename.split(".")[0]
        f = open(txt_path + realName + ".txt", encoding='utf-8')
        __makeXML(size, realName, f)
        f.close()
        print(i, "\n")