UA-DETRAC数据集划分

UA-DETRAC数据集作为世界上最大的车辆检测数据集之一，经常被用来进行对象检测和识别任务方面的训练，但其官方提供的文件格式跟当下流行的Pascal VOC格式不太一致，所以需要对其进行重新划分和xml文件内容的重塑，以下为具体代码

一、将xml文件转换为Pascal VOC格式

导入相关的库

import xml.etree.ElementTree as ET
from xml.dom.minidom import Document
import os
import cv2
import time

函数定义并解析输入XML文件

def ConvertVOCXml(file_path="", file_name=""):
    tree = ET.parse(file_name)
    root = tree.getroot()

这个函数接受两个参数：file_path（输出XML文件的目录路径）和file_name（输入的XML文件名），并使用xml.etree.ElementTree模块解析输入的XML文件，并获取其根节点

初始化变量

num = 0  # 计数
output_file_name = ""

num用于计数转换了多少个frame，output_file_name用于存储输出XML文件的名称

遍历根节点的子节点

for child in root:

假设根节点的直接子节点代表不同的frame

创建DOM文档

        if (child.tag == "frame"):
            # 创建dom文档
            doc = Document()
            # 创建根节点
            annotation = doc.createElement('annotation')
            # 根节点插入dom树
            doc.appendChild(annotation)

使用xml.dom.minidom模块创建一个新的DOM文档，并添加根节点<annotation>

填充输出XML文件内容

            pic_id = child.attrib["num"].zfill(5)
            output_file_name = root.attrib["name"] + "__img" + pic_id + ".xml"

            folder = doc.createElement("folder")
            folder.appendChild(doc.createTextNode("VOC2007"))
            annotation.appendChild(folder)

            filename = doc.createElement("filename")
            pic_name = "img" + pic_id + ".jpg"
            filename.appendChild(doc.createTextNode(pic_name))
            annotation.appendChild(filename)

            sizeimage = doc.createElement("size")
            imagewidth = doc.createElement("width")
            imageheight = doc.createElement("height")
            imagedepth = doc.createElement("depth")

            imagewidth.appendChild(doc.createTextNode("960"))
            imageheight.appendChild(doc.createTextNode("540"))
            imagedepth.appendChild(doc.createTextNode("3"))

            sizeimage.appendChild(imagedepth)
            sizeimage.appendChild(imagewidth)
            sizeimage.appendChild(imageheight)
            annotation.appendChild(sizeimage)

            target_list = list(child)[0]  # 获取target_list
            object = None
            for target in target_list:
                if (target.tag == "target"):
                    object = doc.createElement('object')
                    bndbox = doc.createElement("bndbox")

文件夹名称：固定为"VOC2007"

文件名：根据frame的编号生成

图片尺寸：固定宽度为960，高度为540，深度（通道数）为3

目标对象：遍历frame中的目标对象，为每个对象创建一个<object>节点，并填充边界框（<bndbox>）和属性（如名称、姿态、截断标志和难度标志）

处理边界框

                    for target_child in target:
                        if (target_child.tag == "box"):
                            xmin = doc.createElement("xmin")
                            ymin = doc.createElement("ymin")
                            xmax = doc.createElement("xmax")
                            ymax = doc.createElement("ymax")
                            xmin_value = int(float(target_child.attrib["left"]))
                            ymin_value = int(float(target_child.attrib["top"]))
                            box_width_value = int(float(target_child.attrib["width"]))
                            box_height_value = int(float(target_child.attrib["height"]))
                            xmin.appendChild(doc.createTextNode(str(xmin_value)))
                            ymin.appendChild(doc.createTextNode(str(ymin_value)))
                            if (xmin_value + box_width_value > 960):
                                xmax.appendChild(doc.createTextNode(str(960)))
                            else:
                                xmax.appendChild(doc.createTextNode(str(xmin_value + box_width_value)))
                            if (ymin_value + box_height_value > 540):
                                ymax.appendChild(doc.createTextNode(str(540)))
                            else:
                                ymax.appendChild(doc.createTextNode(str(ymin_value + box_height_value)))

对于每个目标对象的边界框，代码计算其左上角和右下角的坐标。如果计算出的坐标超出了图片尺寸，则将其限制在图片边界内

处理属性

                        if (target_child.tag == "attribute"):
                            name = doc.createElement('name')
                            pose = doc.createElement('pose')
                            truncated = doc.createElement('truncated')
                            difficult = doc.createElement('difficult')

                            name.appendChild(doc.createTextNode(str((target_child.attrib["vehicle_type"]))))

                            pose.appendChild(doc.createTextNode("Left"))  # 随意指定
                            truncated.appendChild(doc.createTextNode("0"))  # 随意指定
                            difficult.appendChild(doc.createTextNode("0"))  # 随意指定

                            object.appendChild(name)
                            object.appendChild(pose)
                            object.appendChild(truncated)
                            object.appendChild(difficult)

                    bndbox.appendChild(xmin)
                    bndbox.appendChild(ymin)
                    bndbox.appendChild(xmax)
                    bndbox.appendChild(ymax)
                    object.appendChild(bndbox)
                    annotation.appendChild(object)

            file_path_out = os.path.join(file_path, output_file_name)
            f = open(file_path_out, 'w')
            f.write(doc.toprettyxml(indent=' ' * 4))
            f.close()
            num = num + 1
    return num

对于每个目标对象，从attribute标签中读取vehicle_type作为对象名称，并随意指定姿态、截断标志和难度标志。

输出XML文件

            file_path_out = os.path.join(file_path, output_file_name)
            f = open(file_path_out, 'w')
            f.write(doc.toprettyxml(indent=' ' * 4))
            f.close()

将生成的DOM文档转换为漂亮的XML字符串，并写入到指定的输出路径

计数并返回

            num = num + 1
    return num

每处理一个frame，计数器num增加1，最后返回处理了多少个frame

主函数

if (__name__ == "__main__"):
    basePath = "UA-DETRAC/Annotations/Test-XML/DETRAC-Test-Annotations-XML"
    totalxml = os.listdir(basePath)
    total_num = 0
    flag = False
    print("正在转换")
    saveBasePath = "xml_train"
    if os.path.exists(saveBasePath) == False:  # 判断文件夹是否存在
        os.makedirs(saveBasePath)
    # Start time
    start = time.time()
    log = open("xml_statistical.txt", "w")  # 分析日志，进行排错
    for xml in totalxml:
        file_name = os.path.join(basePath, xml)
        print(file_name)
        num = ConvertVOCXml(file_path=basePath, file_name=file_name)
        print(num)
        total_num = total_num + num
        log.write(file_name + " " + str(num) + "\n")
    # End time
    end = time.time()
    seconds = end - start
    print("Time taken : {0} seconds".format(seconds))
    print(total_num)

    log.write(str(total_num) + "\n")

至此，数据集中的xml文件格式以成功转换为Pascal VOC格式，但还需对图片文件夹进行处理

二、将图片文件提取出来

导入必要的库

import os
import shutil

定义路径

# xml路径的地址
XmlPath = r'UA-DETRAC\Annotations\Test-XML'
# 原图片的地址
pictureBasePath = r"UA-DETRAC\images\Test-Data\Insight-MVT_Annotation_Test"
# 保存图片的地址
saveBasePath = r"E:UA-DETRAC\images\Test-Data"

获取xml文件列表

total_xml = os.listdir(XmlPath)
num = len(total_xml)
list = range(num)

使用os.listdir(XmlPath)获取XmlPath目录下所有文件的列表，并存储在total_xml变量中，并计算这些XML文件的总数并存储在num变量中。

检查并创建保存路径

if os.path.exists(saveBasePath) == False:  # 判断文件夹是否存在
    os.makedirs(saveBasePath)

使用os.path.exists(saveBasePath)检查保存路径是否存在，如果不存在，使用os.makedirs(saveBasePath)创建该路径。

遍历XML文件

for xml in total_xml:
    xml_temp = xml.split("__")
    folder = xml_temp[0]
    filename = xml_temp[1].split(".")[0] + ".jpg"
    # print(folder)
    # print(filename)
    temp_pictureBasePath = os.path.join(pictureBasePath, folder)
    filePath = os.path.join(temp_pictureBasePath, filename)
    # print(filePath)
    newfile = xml.split(".")[0] + ".jpg"
    newfile_path = os.path.join(saveBasePath, newfile)
    print(newfile_path)
    shutil.copyfile(filePath, newfile_path)

对于total_xml列表中的每一个XML文件（变量名为xml）：

使用split("__")方法将XML文件名分割成两部分，分别代表文件夹名和剩余的文件名部分（xml_temp）。

从xml_temp中提取文件夹名（folder）和图片文件名（不包括扩展名，filename）。

构建原始图片文件的完整路径（filePath）。

根据XML文件名生成新的图片文件名（不包含XML扩展名，只保留基础名并添加.jpg扩展名，newfile）。

构建新图片文件的完整保存路径（newfile_path）。

打印新图片文件的保存路径（用于调试或记录）。

使用shutil.copyfile(filePath, newfile_path)将原始图片复制到新的保存路径。

打印XML文件总数

print("xml file total number", num)