使用YOLOv5-Lite在树莓派4b上部署车辆检测模型（一）——UA-DETRAC车辆检测数据集的处理

目录

UA-DETRAC数据集介绍

UA-DETARC数据集处理过程

处理好的数据集（附下载链接）

参考博文

---

UA-DETRAC数据集介绍

        UA-DETRAC 是一个具有挑战性的现实世界多目标检测和多目标跟踪基准。该数据集包含使用佳能 EOS 550D 相机在中国北京和天津的 24 个不同地点拍摄的 10 小时视频。视频以每秒 25 帧 (fps) 的速度录制，分辨率为 960×540 像素。UA-DETRAC 数据集中有超过 14 万帧和 8250 辆手动标注的车辆，总共有 121 万个标记的对象边界框，其中训练集约82085张图片，测试集约56167张图片。

        官方网站：https://www.albany.edu/cnse/research/computer-vision-machine-learning-lab

        数据集官网介绍：https://wayback.archive-it.org/org-652/20231112205116/https:/detrac-db.rit.albany.edu/ 

UA-DETARC数据集处理过程

原数据集下载链接：

链接：https://pan.baidu.com/s/1QJUSodZZWaBPKxqgnGuIvg 
提取码：rucf

        下载好的数据集有如下图所示的几个文件夹。前两个文件夹内分别是测试集和训练集的xml格式的标注文件，每一个xml文件都是一个摄像头拍摄的很多张图片的全部标注。而下面两个文件夹内分别是测试集和训练集的图片，文件夹内有格式为MVI_xxxxx的子目录，每个子目录都代表一个摄像头拍摄的全部图片，子目录内即为图片。

        现在对数据集进行处理，转为yolov5要求的yolo txt格式。

        1.将voc格式转换为xml格式

        UA-DETRAC数据集提供的标注为voc格式，需要先转为xml格式。

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


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

    num = 0  # 计数
    # 读xml操作

    frame_lists = []
    output_file_name = ""
    for child in root:

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

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

            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 = child.getchildren()[0]  # 获取target_list
            # print(target_list.tag)
            object = None
            for target in target_list:
                if target.tag == "target":
                    # print(target.tag)
                    object = doc.createElement('object')
                    bndbox = doc.createElement("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":
                            vehicle_type = target_child.attrib["vehicle_type"]

                            name = doc.createElement('name')
                            pose = doc.createElement('pose')
                            truncated = doc.createElement('truncated')
                            difficult = doc.createElement('difficult')

                            name.appendChild(doc.createTextNode(str(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


'''
画方框
'''
def bboxes_draw_on_img(img, bbox, color=[255, 0, 0], thickness=2):
    # Draw bounding box...
    print(bbox)
    p1 = (int(float(bbox["xmin"])), int(float(bbox["ymin"])))
    p2 = (int(float(bbox["xmax"])), int(float(bbox["ymax"])))
    cv2.rectangle(img, p1, p2, color, thickness)


def visualization_image(image_name, xml_file_name):
    tree = ET.parse(xml_file_name)
    root = tree.getroot()

    object_lists = []
    for child in root:
        if child.tag == "folder":
            print(child.tag, child.text)
        elif child.tag == "filename":
            print(child.tag, child.text)
        elif child.tag == "size":  # 解析size
            for size_child in child:
                if size_child.tag == "width":
                    print(size_child.tag, size_child.text)
                elif size_child.tag == "height":
                    print(size_child.tag, size_child.text)
                elif size_child.tag == "depth":
                    print(size_child.tag, size_child.text)
        elif child.tag == "object":  # 解析object
            singleObject = {}
            for object_child in child:
                if object_child.tag == "name":
                    # print(object_child.tag,object_child.text)
                    singleObject["name"] = object_child.text
                elif object_child.tag == "bndbox":
                    for bndbox_child in object_child:
                        if bndbox_child.tag == "xmin":
                            singleObject["xmin"] = bndbox_child.text
                            # print(bndbox_child.tag, bndbox_child.text)
                        elif bndbox_child.tag == "ymin":
                            # print(bndbox_child.tag, bndbox_child.text)
                            singleObject["ymin"] = bndbox_child.text
                        elif bndbox_child.tag == "xmax":
                            singleObject["xmax"] = bndbox_child.text
                        elif bndbox_child.tag == "ymax":
                            singleObject["ymax"] = bndbox_child.text
            object_length = len(singleObject)
            if object_length > 0:
                object_lists.append(singleObject)
    img = cv2.imread(image_name)
    for object_coordinate in object_lists:
        bboxes_draw_on_img(img, object_coordinate)
    cv2.imshow("capture", img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # 修改绝对路径
    TrainBasePath = r"C:\...\...\...\UA-DETRAC\DETRAC-Train-Annotations-XML"
    TestBasePath = r"C:\...\...\...\UA-DETRAC\DETRAC-Test-Annotations-XML"
    TrainSaveBasePath = r"C:\...\...\...\UA-DETRAC\train-labels-xml"
    TestSaveBasePath = r"C:\...\...\...\UA-DETRAC\test-labels-xml"

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

    train_totalxml = os.listdir(TrainBasePath)
    test_totalxml = os.listdir(TestBasePath)
    
    # 转换训练集
    total_num = 0
    print("正在转换")
    # Start time
    start = time.time()
    log = open("xml_statistical_train.txt", "w")  # 分析日志，进行排错
    for xml in train_totalxml:
        file_name = os.path.join(TrainBasePath, xml)
        print(file_name)
        num = ConvertVOCXml(file_path=TrainSaveBasePath, file_name=file_name)
        print(num)
        total_num = total_num + num
        log.write(file_name + " " + str(num) + "\n")
    log.close()
    # End time
    end = time.time()
    seconds = end - start
    print("Time taken : {0} seconds".format(seconds))
    print(total_num)

    # 转换测试集
    total_num = 0
    print("正在转换")
    # Start time
    start = time.time()
    log = open("xml_statistical_test.txt", "w")  # 分析日志，进行排错
    for xml in test_totalxml:
        file_name = os.path.join(TestBasePath, xml)
        print(file_name)
        num = ConvertVOCXml(file_path=TestSaveBasePath, file_name=file_name)
        print(num)
        total_num = total_num + num
        log.write(file_name + " " + str(num) + "\n")
    log.close()
    # End time
    end = time.time()
    seconds = end - start
    print("Time taken : {0} seconds".format(seconds))
    print(total_num)

注意修改脚本中的路径，将其改为数据集所在位置的绝对路径。修改后运行该脚本即可。

        2.将xml格式转换为yolo格式

        同样需要修改脚本中的路径，修改完后运行脚本即可。

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join


def convert(size, box):
    # size=(width, height)  b=(xmin, xmax, ymin, ymax)
    # x_center = (xmax+xmin)/2        y_center = (ymax+ymin)/2
    # x = x_center / width            y = y_center / height
    # w = (xmax-xmin) / width         h = (ymax-ymin) / height

    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]

    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]

    # print(x, y, w, h)
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    # print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_file)
        root = tree.getroot()

        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in classes or int(difficult) == 1:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            # b=(xmin, xmax, ymin, ymax)
            # print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
        out_txt_f.close()


if __name__ == "__main__":
    # 把forklift_pallet的voc的xml标签文件转化为yolo的txt标签文件
    # 1、需要转化的类别
    classes = ['car', 'bus', 'van', 'others']  # 注意：这里根据自己的类别名称及种类自行更改
    # 2、voc格式的xml标签文件路径
    xml_files_train = r'C:\...\...\...\UA-DETRAC\train-labels-xml'
    xml_files_test = r'C:\...\...\...\UA-DETRAC\test-labels-xml'
    # # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files_train = r'C:\...\...\...\UA-DETRAC\train-labels-yolo'
    save_txt_files_test = r'C:\...\...\...\UA-DETRAC\test-labels-yolo'

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

    print("正在转换训练集")
    convert_annotation(xml_files_train, save_txt_files_train, classes)
    print("正在转换测试集")
    convert_annotation(xml_files_test, save_txt_files_test, classes)

        3.修改图片名称        

        由于在数据集中图片存放在形式为MVI_xxxxx的各个子目录下，在每个子目录下的图片名字都为从img00001到imgxxxxx，需要更改图片的名称，使图片名与标注的名称一一对应起来。使用下面的脚本file_rename.py即可。同样记得修改路径。

import os

# 获取要修改的文件地址
Train_path = r'C:\...\...\...\UA-DETRAC\Insight-MVT_Annotation_Train'
Test_path = r'C:\...\...\...\UA-DETRAC\Insight-MVT_Annotation_Test'
# 获取文件名列表
train_file_list = os.listdir(Train_path)
print('训练集文件列表如下：')
print(train_file_list)

# 遍历文件名，获取文件名和扩展名
for file in train_file_list:
    path_2 = Train_path + '/' + file
    file_list_inner = os.listdir(path_2)
    for filename in file_list_inner:
        pos = filename.rfind('.') - 8
        newname = file + '__' + filename[pos:-4] + '.jpg'
        # 重新命名文件
        os.rename(path_2 + '/' + filename, path_2 + '/' + newname)


test_file_list = os.listdir(Test_path)
print('测试集文件列表如下：')
print(test_file_list)

# 遍历文件名，获取文件名和扩展名
for file in test_file_list:
    path_2 = Test_path + '/' + file
    file_list_inner = os.listdir(path_2)
    for filename in file_list_inner:
        pos = filename.rfind('.') - 8
        newname = file + '__' + filename[pos:-4] + '.jpg'
        # 重新命名文件
        os.rename(path_2 + '/' + filename, path_2 + '/' + newname)

        4.对数据集进行抽帧

        数据集内图片数量太多了，而且相邻的几张图片之间差别极小，因为它们本来就是同一段视频内相邻的几帧。所以考虑对数据集进行抽帧操作，每十张图片里抽一张作为最终的数据集。

        这里提供两个脚本分别将图片和标注复制到yolov5工程的dataset目录下。image_copy.py每隔十张图片复制一张到yolov5工程目录下，labels_copy.py根据yolov5工程目录下的图片名将对应的标注文件复制到工程中。按需修改绝对路径，依次执行即可。

        image_copy.py

import os
import shutil

# 复制数据集中的图片到yolov5工程目录下，复制完训练集的图片后，修改注释再复制测试集的图片
# 每隔十张图片复制一张图片到工程目录下，完成抽帧
# 源文件夹绝对路径，按需修改
source_folder = r"C:\...\...\...\Insight-MVT_Annotation_Train"
#source_folder = r"C:\...\...\...\Insight-MVT_Annotation_Test"

# 目标文件夹绝对路径，按需修改
target_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\images"
#target_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\images"

# 遍历源文件夹中的所有子文件夹
for subfolder_name in os.listdir(source_folder):
    subfolder_path = os.path.join(source_folder, subfolder_name)

    # 检查是否是文件夹
    if os.path.isdir(subfolder_path):
        # 初始化计数器
        count = 0
        # 遍历子文件夹中的所有文件
        for file_name in os.listdir(subfolder_path):
            file_path = os.path.join(subfolder_path, file_name)
            # 检查是否是图片文件（这里假设只处理JPEG和PNG文件，可以根据实际情况修改）
            if file_name.lower().endswith(('.jpg', '.jpeg', '.png')):
                # 每隔十个图片移动一个
                if count % 10 == 0:
                    # 构建目标文件路径
                    target_file_path = os.path.join(target_folder, file_name)
                    # 移动文件
                    shutil.copy(file_path, target_file_path)
                # 增加计数器
                count += 1
print("文件移动完成。")

        labels_copy.py

import os
import shutil

# 复制数据集中的标注到yolov5工程目录下，复制完训练集的标注后，修改注释再复制测试集的标注
# 目标文件夹路径，按需修改
target_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\images"
#target_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\images"

# 文本文件源文件夹路径，按需修改
source_text_files_folder = r"C:\...\...\...\UA-DETRAC\train-labels-yolo"
#source_text_files_folder = r"C:\...\...\...\UA-DETRAC\test-labels-yolo"

# 目标文本文件文件夹路径
destination_text_files_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\labels"
#destination_text_files_folder = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\labels"


# 获取目标文件夹中的图片文件名（去除后缀）
target_image_file_names = set(os.path.splitext(file)[0] for file in os.listdir(target_folder) if file.lower().endswith(('.jpg', '.jpeg', '.png')))

# 遍历文本文件源文件夹中的所有文件
for file_name in os.listdir(source_text_files_folder):
    file_path = os.path.join(source_text_files_folder, file_name)
    # 检查是否是.txt文件且文件名与目标文件夹中的图片相同
    if file_name.lower().endswith('.txt') and os.path.splitext(file_name)[0] in target_image_file_names:
        # 构建目标文本文件路径
        destination_file_path = os.path.join(destination_text_files_folder, file_name)
        # 复制文本文件
        shutil.copy(file_path, destination_file_path)
print("文本文件复制完成。")

        5. 剔除掉无标注的图片

        数据集原始的xml标注文件中就缺失了部分图片的标注，需要将这些没有标注的图片剔除出去。在依次运行完image_copy.py、labels_copy.py两个脚本后，执行下面的脚本find_no_labels.py即可。

# 从原始训练集中抽帧，寻找抽帧得到的图片的标注时，发现找到的label文件少于图片数量
# 原因为数据集原始的xml标注文件本身就有缺失，导致部分图片缺少标注
import os
import shutil

# 绝对路径，按需修改
# 剔除掉训练集的无标注图片后，修改注释再将测试集内的无标注图片剔除
folder_image_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\images"
folder_labels_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\labels"
destination_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_train\images_with_no_label"  # 指定移动到的路径
'''
folder_image_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\images"
folder_labels_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\labels"
destination_path = r"C:\...\...\...\yolov5-6.2\datasets\UA-DETRAC_test\images_with_no_label"  # 指定移动到的路径
'''

# 获取文件夹B中.txt文件的文件名（不考虑后缀）
txt_file_names = {os.path.splitext(file)[0] for file in os.listdir(folder_labels_path) if file.endswith('.txt')}

for file in os.listdir(folder_image_path):
    if file.endswith('.jpg'):
        # 获取文件名（不考虑后缀）
        file_name = os.path.splitext(file)[0]

        # 如果文件名不在txt_file_names中，则移动文件到指定路径
        if file_name not in txt_file_names:
            source_path = os.path.join(folder_image_path, file)
            destination_file_path = os.path.join(destination_path, file)
            shutil.move(source_path, destination_file_path)
            print(f"Moved: {file} to {destination_path}")

执行完前面的几个步骤后，在yolov5工程datasets目录下就有了训练集UA-DETRAC_train，测试集UA-DETRAC_test。images目录下为图片，labels下为标注。

datasets
├── UA-DETRAC_train
│   ├── images
│   └── labels
├── UA-DETRAC_test
│   ├── images
│   └── labels
...

处理好的数据集（附下载链接）

UA-DETRAC_yolo
├── UA-DETRAC_train
│   ├── images
│   └── labels
└── UA-DETRAC_test
    ├── images
    └── labels

链接：https://pan.baidu.com/s/1H5vsla6vD4yJpGGJD3ECEA 
提取码：fz4r

参考博文

UA-DETRAC数据集转YOLO格式_ua-detrac数据集下载-CSDN博客 本文部分脚本的编写参考了该博文

使用yolov5训练UA-DETRAC车辆数据集 - 知乎 这篇文章在UA-DETARC数据集实验了多种不同处理方式对最终训练出模型的精度的影响，如果要使用UA-DETRAC数据集训练模型，可以先参考一下这篇文章。