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

txt格式，每行从左往右分别<bbox_left>,<bbox_top>,<bbox_width>,<bbox_height>,<score>,<object_category>,<truncation>,<occlusion> 分类：

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 = "
" + level * "	"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "	"
        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, "
")