YOLOv5训练KITTI数据集

先把态度摆正

1763人浏览 · 2025-04-17 18:41:38

先把态度摆正 · 2025-04-17 18:41:38 发布

一、下载KITTI数据集

这一部分网上有很多教程，这里直接给出网盘链接：
链接: https://pan.baidu.com/s/1dsTyLELeJ3_ZxpvgomhFNw?pwd=kitt 提取码: kitt

由于原始数据集过于庞大，因此我们只选取部分数据进行训练，这里我选择的是000000-000399，000901-001000，共计500张图片。

其中80%用于训练集，20%用于验证集。

新建VOCdevkit文件夹，组成如下形式，将下载好的数据对应放在文件夹下：

VOCdevkit
|__VOC2007
   |__Annotations
   |__JPEGImages
   |__YOLOLabels
   |__predefined_classes.txt

其中JPEGImages文件夹下放的是*.png图片，YOLOLabels文件夹下放的是*.txt文本。

直接训练的话会报错，原因就是下载的数据集是KITTI独有的数据格式，不适用于我们的YOLOv5网络，所以我们得转换一下。

二、KITTI数据集转换成可运行的YOLOv5格式

这一部分参考了kitti数据集转换成可运行的YOLOv5格式这篇文章。

我们打开下载好的KITTI数据集里标签中的一个内容，比如000000.txt，

Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 1.89 0.48 1.20 1.84 1.47 8.41 0.01

需要经过下面的步骤将其转化为YOLOv5格式，比如000000.txt，

2 0.6217320261437909 0.6081081081081081 0.08006535947712418 0.44324324324324327

注意： 从这里我们可以看出Pedestrian转化后的标签为2，那么在yolov5-5.0/data/****.yaml中也要与之对应，其他两类同理，即：

'Car','Cyclist','Pedestrian'按照0，1，2的顺序排好，否则在后续检测的时候会输出相反的类别标签。

第一个脚本：

将原始数据集里的类别归为三类：Car、Cyclist、Pedestrian

# modify_annotations_txt.py
import glob

import string

txt_list = glob.glob('C:/Users/lenovo/Desktop/VOCdevkit/VOC2007/YOLOLabels/*.txt')  # 存储Labels文件夹所有txt文件路径 ；注意要保留/*

print(txt_list)


def show_category(txt_list):
    category_list = []
    for item in txt_list:
        try:
            with open(item) as tdf:
                for each_line in tdf:
                    labeldata = each_line.strip().split(' ')  # 去掉前后多余的字符并把其分开
                    category_list.append(labeldata[0])  # 只要第一个字段，即类别
        except IOError as ioerr:
            print('File error:' + str(ioerr))
    print(set(category_list))  # 输出集合


def merge(line):
    each_line = ''
    for i in range(len(line)):
        if i != (len(line) - 1):
            each_line = each_line + line[i] + ' '
        else:
            each_line = each_line + line[i]  # 最后一条字段后面不加空格
    each_line = each_line + '\n'
    return (each_line)


print('before modify categories are:\n')
show_category(txt_list)

for item in txt_list:
    new_txt = []
    try:
        with open(item, 'r') as r_tdf:
            for each_line in r_tdf:
                labeldata = each_line.strip().split(' ')
                if labeldata[0] in ['Truck','Van','Tram']: # 合并汽车类
                    labeldata[0] = labeldata[0].replace(labeldata[0],'Car')
                if labeldata[0] == 'Person_sitting':  # 合并行人类
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'Pedestrian')
                if labeldata[0] == 'DontCare':  # 忽略Dontcare类
                    continue
                if labeldata[0] == 'Misc':  # 忽略Misc类
                    continue
                new_txt.append(merge(labeldata))  # 重新写入新的txt文件
        with open(item, 'w+') as w_tdf:  # w+是打开原文件将内容删除，另写新内容进去
            for temp in new_txt:
                w_tdf.write(temp)
    except IOError as ioerr:
        print('File error:' + str(ioerr))

print('\nafter modify categories are:\n')
show_category(txt_list)

第二个脚本：

将.txt转化为.xml并存放在Annotations下

# kitti_txt_to_xml.py
# encoding:utf-8
# 根据一个给定的XML Schema，使用DOM树的形式从空白文件生成一个XML
from xml.dom.minidom import Document
import cv2
import glob
import os


def generate_xml(name, split_lines, img_size, class_ind):
    doc = Document()  # 创建DOM文档对象
    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)
    title = doc.createElement('folder')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    annotation.appendChild(title)
    img_name = name + '.png'
    title = doc.createElement('filename')
    title_text = doc.createTextNode(img_name)
    title.appendChild(title_text)
    annotation.appendChild(title)
    source = doc.createElement('source')
    annotation.appendChild(source)
    title = doc.createElement('database')
    title_text = doc.createTextNode('The KITTI Database')
    title.appendChild(title_text)
    source.appendChild(title)
    title = doc.createElement('annotation')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    source.appendChild(title)
    size = doc.createElement('size')
    annotation.appendChild(size)
    title = doc.createElement('width')
    title_text = doc.createTextNode(str(img_size[1]))
    title.appendChild(title_text)
    size.appendChild(title)
    title = doc.createElement('height')
    title_text = doc.createTextNode(str(img_size[0]))
    title.appendChild(title_text)
    size.appendChild(title)
    title = doc.createElement('depth')
    title_text = doc.createTextNode(str(img_size[2]))
    title.appendChild(title_text)
    size.appendChild(title)
    for split_line in split_lines:
        line = split_line.strip().split()
        if line[0] in class_ind:
            object = doc.createElement('object')
            annotation.appendChild(object)
            title = doc.createElement('name')
            title_text = doc.createTextNode(line[0])
            title.appendChild(title_text)
            object.appendChild(title)
            bndbox = doc.createElement('bndbox')
            object.appendChild(bndbox)
            title = doc.createElement('xmin')
            title_text = doc.createTextNode(str(int(float(line[4]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymin')
            title_text = doc.createTextNode(str(int(float(line[5]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('xmax')
            title_text = doc.createTextNode(str(int(float(line[6]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymax')
            title_text = doc.createTextNode(str(int(float(line[7]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
    # 将DOM对象doc写入文件
    f = open(r'C:/Users/lenovo/Desktop/VOCdevkit/VOC2007/Annotations/' + name + '.xml', 'w')  # xml要改这里
    f.write(doc.toprettyxml(indent=''))
    f.close()


if __name__ == '__main__':
    class_ind = ('Pedestrian', 'Car', 'Cyclist')
    cur_dir = os.getcwd()
    labels_dir = os.path.join(cur_dir, r'C:/Users/lenovo/Desktop/VOCdevkit/VOC2007/YOLOLabels')  # 要改这里
    # labels_dir=os.path.join(cur_dir,'label_2')
    for parent, dirnames, filenames in os.walk(labels_dir):  # 分别得到根目录，子目录和根目录下文件
        for file_name in filenames:
            full_path = os.path.join(parent, file_name)  # 获取文件全路径
            f = open(full_path)
            split_lines = f.readlines()  # 以行为单位读
            name = file_name[:-4]  # 后四位是扩展名.txt，只取前面的文件名
            img_name = name + '.png'
            img_path = os.path.join(r'C:/Users/lenovo/Desktop/VOCdevkit/VOC2007/JPEGImages',
                                    img_name)  # 路径需要自行修改    改这里
            img_size = cv2.imread(img_path).shape
            generate_xml(name, split_lines, img_size, class_ind)
            print('txts has converted into xmls')

可以使用labelimg查看*.xml文件：

第三个脚本：

把*.xml转化为适合于yolo训练的*.txt标签模式（该脚本要与待处理的数据放在同一目录下）

# xml_to_yolo_txt.py
# 此代码和VOC_KITTI文件夹同目录
import glob
import xml.etree.ElementTree as ET

# 这里的类名为我们xml里面的类名，顺序现在不需要考虑
class_names = ['Car', 'Cyclist', 'Pedestrian']
# xml文件路径
path = r'C:/Users/lenovo/Desktop/VOCdevkit/VOC2007/sign/'  # 要改这里

# 转换一个xml文件为txt
def single_xml_to_txt(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    # 保存的txt文件路径
    txt_file = xml_file.split('.')[0] + '.txt'
    with open(txt_file, 'w') as txt_file:
        for member in root.findall('object'):
            # filename = root.find('filename').text
            picture_width = int(root.find('size')[0].text)
            picture_height = int(root.find('size')[1].text)
            class_name = member[0].text
            # 类名对应的index
            class_num = class_names.index(class_name)

            box_x_min = int(member[1][0].text)  # 左上角横坐标
            box_y_min = int(member[1][1].text)  # 左上角纵坐标
            box_x_max = int(member[1][2].text)  # 右下角横坐标
            box_y_max = int(member[1][3].text)  # 右下角纵坐标
            print(box_x_max, box_x_min, box_y_max, box_y_min)
            # 转成相对位置和宽高
            x_center = float(box_x_min + box_x_max) / (2 * picture_width)
            y_center = float(box_y_min + box_y_max) / (2 * picture_height)
            width = float(box_x_max - box_x_min) / picture_width
            height = float(box_y_max - box_y_min) / picture_height
            print(class_num, x_center, y_center, width, height)
            txt_file.write(str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(
                height) + '\n')


# 转换文件夹下的所有xml文件为txt！！！！！！！！！！！！！！！！！
def dir_xml_to_txt(path):
    for xml_file in glob.glob(path + '*.xml'):
        single_xml_to_txt(xml_file)

dir_xml_to_txt(path)

在运行第三个脚本时，到第73张图片时程序中断了，出现了报错：

"int(member[1][0].text),IndexError: child index out of range"

分别打开000073.xml文件和000000.xml文件，如下：

对比两幅图可得，原因在于000073.xml文件里的bndbox是object的第5个参数，从0算起的话对应数字应该是4，而转换代码里写的member[1][0].text，所以应该把member[1][0].text到member[1][3].text中的1全部换成4，即：

box_x_min = int(member[4][0].text)  # 左上角横坐标
box_y_min = int(member[4][1].text)  # 左上角纵坐标
box_x_max = int(member[4][2].text)  # 右下角横坐标
box_y_max = int(member[4][3].text)  # 右下角纵坐标

可是这样一来，虽然可将000073.xml转化为000073.txt，其他的******.xml却无法转化。

这里我的办法是把单独处理000073.xml，新建一个文件夹，里面放入剪切过来的000073.xml和复制的一份第三个脚本py文件，并对该py文件做上述修改，即可处理000073.xml

注：后续做了第二次实验（扩充KITTI数据集为7841张图片），我在运行上述脚本3时没有报错，我也不清楚是什么原因。

三、划分训练集与验证集

方法一：

最无脑，直接手动划分数据集如下：

VOCdevkit
|__images
   |__train
   |__val
|__labels
   |__train
   |__val
|__VOC2007
   |__Annotations
   |__JPEGImages
   |__YOLOLabels
   |__predefined_classes.txt

方法二：

新建 huafen.py 文件，即如下代码，随机划分数据集：

# huafen.py

import os, random, shutil


def moveimg(fileDir, tarDir):
    pathDir = os.listdir(fileDir)  # 取图片的原始路径
    filenumber = len(pathDir)
    rate = 0.2  # 自定义抽取图片的比例，比方说100张抽20张，那就是0.2 ->这里是验证集0.2
    picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量图片
    sample = random.sample(pathDir, picknumber)  # 随机选取picknumber数量的样本图片
    print(sample)
    for name in sample:
        shutil.move(fileDir + name, tarDir + "\\" + name)
    return


def movelabel(file_list, file_label_train, file_label_val):
    for i in file_list:
        if i.endswith('.png'):
            # filename = file_label_train + "\\" + i[:-4] + '.xml'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            filename = file_label_train + "\\" + i[:-4] + '.txt'  # 可以改成xml文件将’.txt‘改成'.xml'就可以了
            if os.path.exists(filename):
                shutil.move(filename, file_label_val)
                print(i + "处理成功！")


if __name__ == '__main__':
    fileDir = r"VOCdevkit\VOC2007\JPEGImages" + "\\"  # 源图片文件夹路径
    tarDir = r'VOCdevkit\images\val'  # 图片移动到新的文件夹路径
    moveimg(fileDir, tarDir)
    file_list = os.listdir(tarDir)
    file_label_train = r"VOCdevkit\VOC2007\YOLOLabels"  # 源图片标签路径
    file_label_val = r"VOCdevkit\labels\val"  # 标签
    # 移动到新的文件路径
    movelabel(file_list, file_label_train, file_label_val)

tips：需要预先创建如下目录

使用如上代码，可以将一个文件夹里面的图片按一定的比例（修改第7行中的rate的值）抽取照片放入到一个新的文件夹里面（需要新建一个文件夹）。这样源文件夹就变成了训练集，新建的为验证集。然后代码会将抽取照片对应的标签文件放入到一个新建的文件夹中（该文件夹需要新建）。这样我们就有图片和标签文件的训练集和验证集。