将KITTI中的2D检测数据集转换为YOLOV5能训练的格式

下载KITTI数据集中的2d检测图片

关于该数据集的介绍，参考: kitti目标检测2D数据集
注意：

1. 2d数据集的连接
2. 下载数据需要自己注册一个kitti的账号
3. 下数据集和标注文件（详见该部分的参考）

去除DontCare标注

DontCare 标签表示物体未被标记的区域，例如因为它们离激光扫描仪太远。为防止此类对象干扰检测结果，我们直接忽略数据集中的 DontCare 区域。此部分及其后续部分参考：用yolov5训练kitti数据集，并在该博客的基础上加以修改。

代码实现：

import glob
import os
import string


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))
    unique_category = set(category_list)
    return unique_category


def process_labels(txt_list):
    for item in txt_list:
        new_txt = []
        new_txtname = str(item).replace('label_2', 'new_label')
        try:
            with open(item, 'r') as r_tdf:
                res = r_tdf.readlines()
                # print(res)
                for row in res:
                    category = str(row).split(' ')[0]
                    # print(category)
                    if category == 'DontCare':  # 忽略Dontcare类
                        continue
                    else:
                        new_txt.append(row)  # 将其余类别重新写入新的txt文件
            # print(new_txt)
            if not os.path.exists('new_label'):
                os.mkdir('new_label')
            with open(new_txtname, 'w+') as w_tdf:  # w+是打开原文件将内容删除，另写新内容进去
                for temp in new_txt:
                    w_tdf.write(str(temp))

        except IOError as ioerr:
            print('File error:' + str(ioerr))


if __name__ == '__main__':
    txt_list = glob.glob('label_2/*.txt')
    # classes = show_category(txt_list)
    # print(classes)  # 输出集合
    # txt_list = txt_list[:10]
    process_labels(txt_list)

转换为XML文件

原数据集的标注格式并不适合在yolov5中使用。故在此，我们先将其转换为xml格式的标注，之后再转换为归一化坐标的txt格式标注。
代码实现：

# kitti_txt_to_xml.py
# encoding:utf-8
# 根据一个给定的XML Schema，使用DOM树的形式从空白文件生成一个XML
from xml.dom.minidom import Document
import cv2
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写入文件
    if not os.path.exists('Annotations'):
        os.mkdir('Annotations')
    f = open('Annotations/'+name+'.xml', 'w')
    f.write(doc.toprettyxml(indent = ''))
    f.close()


if __name__ == '__main__':
    class_ind = ('Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc')
    cur_dir = os.getcwd()
    labels_dir = os.path.join(cur_dir, 'new_label')
    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('/home/dcr/kitti/training/image_2', img_name) # 路径需要自行修改
            img_size = cv2.imread(img_path).shape
            generate_xml(name, split_lines, img_size, class_ind)

    print('all txts has converted into xmls')

xml转换为txt并划分标签的训练集和验证集

此部分先将标签转换为txt格式，然后划分为训练集和验证集，数据集一共7481张图片，其中划分出验证集1500张图片（测试集在下载的原始数据解压之后就有，不需要单独划分）。
代码实现：

# xml_to_yolo_txt.py
import glob
import os
import random
import xml.etree.ElementTree as ET
import numpy as np


# 转换一个xml文件为txt
def single_xml_to_txt(xml_file, save_path: str):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    print(xml_file)
    txt_file = str(xml_file.split('.')[1]).replace('/Annotations', save_path) + '.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):
    # 保存的txt文件路径
    if not os.path.exists('labels/train'):
        os.makedirs('labels/train')  # 递归创建
    if not os.path.exists('labels/val'):
        os.makedirs('labels/val')

    xml_list = glob.glob(path + '*.xml')
    random.shuffle(xml_list)
    xml_val = xml_list[:1500]  # 验证数据
    xml_train = xml_list[1500:]  # 训练数据
    for xml_file in xml_val:
        single_xml_to_txt(xml_file,'labels/val')
    for xml_file in xml_train:
        single_xml_to_txt(xml_file,'labels/train')


if __name__ == '__main__':
    # 这里的类名为我们xml里面的类名，顺序现在不需要考虑
    class_names = ['Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc']
    # xml文件路径
    path = './Annotations/'

    dir_xml_to_txt(path)
    print("Done")

根据标签的划分结果将图片分开

将图片从原始数据中分别复制到指定的文件夹。
代码实现：

import glob
import os
import shutil
import re


def img_split(origin_path, target_path, label_path):
    if not os.path.exists(target_path):
        os.makedirs(target_path)
    origin_img_path = origin_path  # 原图片路径 'training/image_2/'
    train_list = glob.glob(label_path + '*.txt')  # 标签文件的路径'label_train/*.txt’
    # print(train_list)
    pattern = re.compile(r'\d+')  # 使用正则表达式得到标签中的文件名（6个数字）
    img_list = pattern.findall(str(train_list))

    for img in img_list:
        img_file = origin_img_path + img + '.png'
        img_path = target_path + img + '.png'   # 目标文件路径'img_train/*.png'
        # print(img_file)
        shutil.copy(img_file, img_path)  # 从原路径复制图片到目标路径


if __name__ == '__main__':
    # 复制训练集图片
    img_split('/home/dcr/kitti/training/image_2/', 'images/train/', 'labels/train/')
    # 复制验证集图片
    img_split('/home/dcr/kitti/training/image_2/', 'images/val/', 'labels/val/')

数据集处理完成

按顺序执行完上述步骤之后，其中的images和labels文件夹就可以直接在 yolov5 中使用。整个过程中主要参考：

1. https://zhuanlan.zhihu.com/p/57493587
2. https://blog.csdn.net/qq_45978858/article/details/119718516