基于MMDetection的atss模型在kitti数据集上进行训练

基于MMDetection的atss模型在kitti数据集上进行训练

kitti数据集国内下载地址：https://opendatalab.com/KITTI_Object/download

这个网站里面有很多数据集可以下载，强烈安利给大家

kitti数据集官网: https://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d

代码我放在了gihutb : https://github.com/kouyuanbo/kitti2coco.git，欢迎star✨✨

最近需要在kittii数据集上做些实验，选择了使用MMDetection框架，而MMDetection框架没有提供对kitti数据集的处理方案。因此，在MMDetection上训练kitti数据集需要一些额外的处理。MMDetection官方文档给出了使用MMDetection在自定义数据集上进行训练的三种方式：

MMDetection 一共支持三种形式应用新数据集：

1. 将数据集重新组织为 COCO 格式。
2. 将数据集重新组织为一个中间格式。（这种方法官方有实例，也可以参考这里）
3. 实现一个新的数据集。

本文主要使用第一种方法（最简单），将数据集转换成COCO格式。

😋tips:本实验使用的MMDetection 版本是3.0.0，这个版本是2023年新发布的，有一些代码已经跟之前的不一样了，因此，如果遇到问题**强烈建议先去查文档。强烈建议先去查文档。强烈建议先去查文档。**文档已经写的非常详细了，如果查一些之前的博客反而可能会遇到问题。

如果过程中遇到问题，可以评论哦，看到会回复的~如果你觉得这篇博客对你有点用，可以动动小手，点个👍嘛。

预备知识

COCO数据集标注格式介绍

首先先对COCO数据集的格式做一个介绍，coco数据集的目录结构如下：

coco
├── annotations
├── test2017
├── train2017
└── val2017

coco格式的核心是一个json文件，里面包含三部分：

• images：一个list，里面每个元素是一个字典，包含file_name,height,width,id等内容
• annotations：一个list，里面每个元素也是一个字典，包含segmentation, area, iscrowd, image_id, bbox, category_id, id等内容。其中对于目标检测任务有用的字段有id, image_id, category_id, bbox, area。其中最关键的是bbox的格式，COCO数据集标注的bbox格式是(x, y, w, h)，其中的(x, y)是左上角的坐标！🤡麻了，我一直以为是中心点的坐标，结果卡了好长时间
• categories：目标类别，标号->类别名称的映射关系

'images': [
 {
     'file_name': 'COCO_val2014_000000001268.jpg',	// 图片的名称
     'height': 427,	// 图片的高度
     'width': 640,	// 图片的宽度
     'id': 1268	// 图片的 id
 },
 ...
],

'annotations': [
 {
     'segmentation': [[192.81,
         247.09,
         ...
         219.03,
         249.06]],  # if you have mask labels
     'area': 1035.749,	// w * h得到，目标的面积
     'iscrowd': 0,
     'image_id': 1268,	// 所属图片的 id
     'bbox': [192.81, 224.8, 74.73, 33.43],	// (x, y, w, h) (x, y) 为标注框的左上角坐标
     'category_id': 16,	// 所属的类别 id ,从 1 开始
     'id': 42986	// 目标的 id
 },
 ...
],

'categories': [
 {'id': 0, 'name': 'car'},
]

Kitti数据集标注格式

Kitti数据集使用的是txt文件进行标注，如下图：

下面引用一段kitti官方给出的说明，这段说明中解释了标注文件中每个字段的含义，我下面列出对于2D目标检测有用的字段：

• type: 目标的类别，共有8个类别，(‘Car’, ‘Van’, ‘Truck’,‘Pedestrian’, ‘Person_sitting’, ‘Cyclist’, ‘Tram’, ‘Misc’ or ‘DontCare’)
• bbox: 检测框参数，**[left, top, right, bottom]**的格式

#Values    Name      Description
----------------------------------------------------------------------------
   1    type         Describes the type of object: 'Car', 'Van', 'Truck',
                     'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram',
                     'Misc' or 'DontCare'
   1    truncated    Float from 0 (non-truncated) to 1 (truncated), where
                     truncated refers to the object leaving image boundaries
   1    occluded     Integer (0,1,2,3) indicating occlusion state:
                     0 = fully visible, 1 = partly occluded
                     2 = largely occluded, 3 = unknown
   1    alpha        Observation angle of object, ranging [-pi..pi]
   4    bbox         2D bounding box of object in the image (0-based index):
                     contains left, top, right, bottom pixel coordinates
   3    dimensions   3D object dimensions: height, width, length (in meters)
   3    location     3D object location x,y,z in camera coordinates (in meters)
   1    rotation_y   Rotation ry around Y-axis in camera coordinates [-pi..pi]
   1    score        Only for results: Float, indicating confidence in
                     detection, needed for p/r curves, higher is better.

步骤和代码

合并类别

kitti数据集中有很多个类别，在这里我们只设置三个类别(Car,Pedestrian,Cyclist)，我们可以把Van, Truck, Tram这三个类别合并入Car这个类别，把Person_sitting这个类别合并到Pedestrian这个类别，忽略DontCare和Misc类别。这里直接使用了这位博主中的modify_annotations_txt.py。

需要注意的是，该脚本文件会直接在kitti数据集上进行修改，因此，建议改之前备份kitti数据集的标签。

modify_annotations_txt.py代码：

# modify_annotations_txt.py
# 将Car’，’Cyclist’，’Pedestrian’
# 将 ‘Van’, ‘Truck’, ‘Tram’ 合并到 ‘Car’ 类别中去
# 将 ‘Person_sitting’ 合并到 ‘Pedestrian’ 类别中去
# ‘Misc’ 和 ‘Dontcare’ 这两类直接忽略

import glob
import string

txt_list = glob.glob('./label/training/label_2/*.txt') # 存储Labels文件夹所有txt文件路径

# 查看类别集合
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)

将图片划分成训练集和验证集

将kitti数据集下载下来并解压，kitti目标检测数据集一共包含7481张训练图片和7518张测试图片，本文只使用训练集数据，因此，需要将训练图片划分成训练集和验证集，我这里使用的比例是9:1。

首先将训练集图片和标注文件按照以下格式进行整理，将图片放入image_2文件夹，将标注文件放入label_2文件夹中。

data_dir
├── image_2
├──────0000001.png
├──────0000002.png
├──────0000003.png
└── label_2
├──────0000001.txt
├──────0000002.txt
└──────0000003.txt

接着运行split_datasets.py文件，其中的

• dest_dir变量指定的是划分好的数据集的存储位置。
• ratio变量用来调整训练集和验证集的比例（0.9 = 训练集的数量 / 全部数据数量）。

当脚本运行完，我们就可以得到划分好的数据集，目录如下。train_labels和val_labels中存储的分别是训练集和验证集的的标签文件，train2017和val2017存储的分别是训练集的图片和验证集的图片。

dest_dir
├── labels
│   ├── train_labels
│   └── val_labels
├── train2017
└── val2017

split_datasets.py代码如下：

import os
import random
import shutil

'''
将数据按照一定比例划分成训练集和验证集，目录结构：

data_dir
├── image_2
├──────0000001.png
├──────0000002.png
├──────0000003.png
└── label_2
├──────0000001.txt
├──────0000002.txt
└──────0000003.txt

最终得到的目录结构：
dest_dir
├── annotations
├── labels
│   ├── train_labels
│   └── val_labels
├── train2017
└── val2017


'''
# 训练集占的比例
ratio = 0.9

# 设置随机种子以确保可重复性
random.seed(42)

# 指定数据集路径和训练/验证集路径
data_dir = "/data2/2022/kyb/datasets/Kitti/training/"
dest_dir = '/data2/2022/kyb/datasets/kitti_coco/'

train_img_dir = os.path.join(dest_dir, 'train')
train_label_dir = os.path.join(dest_dir, 'labels/train_labels')
val_img_dir = os.path.join(dest_dir, 'val')
val_label_dir = os.path.join(dest_dir, 'labels/val_labels')


# 创建训练/验证集文件夹
os.makedirs(train_img_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_img_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)

# 获取该类别的所有图像
img_path = os.path.join(data_dir + 'image_2')
all_images = os.listdir(img_path)
num_images = len(all_images)

label_path = os.path.join(data_dir + 'label_2')

# 打乱顺序
random.shuffle(all_images)

# 计算分割点
split_index = int(ratio * num_images)

# 将前90%的图像复制到训练集文件夹
for image_name in all_images[:split_index]:
    src_path = os.path.join(img_path, image_name)
    dst_path = os.path.join(train_img_dir, image_name)
    shutil.copyfile(src_path, dst_path)

    src_label_path = os.path.join(label_path, image_name[:-4] + '.txt')
    dst_label_path = os.path.join(train_label_dir, image_name[:-4] + '.txt')
    shutil.copyfile(src_label_path, dst_label_path)

# 将后10%的图像复制到验证集文件夹
for image_name in all_images[split_index:]:
    src_path = os.path.join(img_path, image_name)
    dst_path = os.path.join(val_img_dir, image_name)
    shutil.copyfile(src_path, dst_path)

    src_label_path = os.path.join(label_path, image_name[:-4] + '.txt')
    dst_label_path = os.path.join(val_label_dir, image_name[:-4] + '.txt')
    shutil.copyfile(src_label_path, dst_label_path)
    

print("数据集划分完成！" + "训练集图片数目: " + str(split_index) + '验证集图片数目: '+ str(num_images - split_index))

将kitti的标签文件预处理成coco的标签文件格式即json格式

本文使用kitti2coco.py文件对kitti的标签文件进行处理，在处理前要保证以下目录结构，如果你使用的是本文之前的脚本处理的话，那么自然就是这个目录结构，如果不是的话，建议先整理成以下结构，或者自己根据代码做出调整。

data_root
├── labels
│   ├── train_labels
│   └── val_labels
├── train2017
└── val2017

kitti2coco.py脚本的主要功能就是先读取所有的训练集label，然后整理成coco格式的标注，最后将其写入json文件中。coco与kitti标注文件的格式与含义在上文中已经进行说明。唯一一点需要注意的是将kitti的bbox转换成coco的bbox的逻辑，kitti的bbox格式是(x1,y1,x2,y2)，即GT左上角点的坐标和右下角点的坐标。coco的bbox的格式，是(x,y,w,h)，其中(x,y)是左上角点的坐标。知道了这些变量的含义后，相信你肯定能理解代码为什么这么写了。

kitti2coco.py代码如下，其中data_root是数据集的根目录：

import os
import json
import argparse
import cv2


def kitti2coco(label_dir, img_dir, output_dir, suffix):
    # Create COCO annotation structure
    coco = {}
    coco['images'] = []
    coco['annotations'] = []
    coco['categories'] = []

    # Add categories
    categories = [
        {'id': 1, 'name': 'Car'},
        {'id': 2, 'name': 'Pedestrian'},
        {'id': 3, 'name': 'Cyclist'}
    ]
    coco['categories'] = categories

    # Add images and annotations
    image_id = 0
    annotation_id = 0

    for file in os.listdir(label_dir):
        if file.endswith('.txt'):
            image_path = os.path.join(img_dir, file[:-4] + '.png')
            
            # 读取图片的高宽
            img_file = cv2.imread(image_path)
            img_height, img_width = img_file.shape[0],img_file.shape[1]
            
            image = {
                'id': image_id,
                'file_name': file[:-4] + '.png',
                'height': img_height, # KITTI image height
                'width': img_width # KITTI image width
            }
            coco['images'].append(image)

            with open(os.path.join(label_dir, file), 'r') as f:
                lines = f.readlines()
                for line in lines:
                    line = line.strip().split(' ')
                    category_id = 1 if line[0] == 'Car' else 2 if line[0] == 'Pedestrian' else 3
                    bbox = [float(coord) for coord in line[4:8]]

                    x1, y1, x2, y2 = bbox
                    bbox_width = x2 - x1
                    bbox_height = y2 - y1


                    annotation = {
                        'id': annotation_id,
                        'image_id': image_id,
                        'category_id': category_id,
                        'bbox': [x1, y1, bbox_width,bbox_height],
                        'area': bbox_height*bbox_width,
                        'iscrowd': 0
                    }
                    coco['annotations'].append(annotation)
                    annotation_id += 1

            image_id += 1

    # Write COCO annotation to file
    with open(os.path.join(output_dir, 'instances_'+ suffix + '2017' +'.json'), 'w') as f:
        json.dump(coco, f)

if __name__ == '__main__':

    '''

    目录结构：
    data_root
    ├── annotations
    ├── labels
    │   ├── train_labels
    │   └── val_labels
    ├── train2017
    └── val2017

    '''
    data_root = '/data2/2022/kyb/datasets/kitti_coco/'

    # 输出路径
    outputs_path = os.path.join(data_root, 'annotations')
    os.makedirs(outputs_path, exist_ok=True)

    
    train_img_dir = os.path.join(data_root, 'train2017')    # 训练集图片的路径
    train_label_dir = os.path.join(data_root, 'labels/train_labels')    # 训练集label的路径
    kitti2coco(train_label_dir, train_img_dir, outputs_path, 'train')   # 转换训练集标注格式

    val_img_dir = os.path.join(data_root, 'val2017')    # 验证集图片的路径
    val_label_dir = os.path.join(data_root, 'labels/val_labels')    # 验证集label的路径
    kitti2coco(val_label_dir, val_img_dir, outputs_path, 'val') # 转换验证集标注格式

    print('格式转换完成！')

修改MMDetection中的类别信息和网络的输出维度

经过以上合并kitti的类别标签、划分训练集验证集、生成coco格式的json标签三步的处理。我们就成功将kitti数据集转换成coco格式了。其它数据集也是同理的，关键是弄懂两种数据集格式之间的区别。

下载MMDetection以及环境的教程官方已经写的很清楚了可以参考这里。本文就不赘述了。因为kitti数据集的类别标签个数与coco不同，因此我们要修改代码中的类别标签。

1、修改类别标签 coco.py 文件中的标签

mmdetection/mmdet/datasets/coco.py文件定义了读取COCO数据集的方式，其中classes中定义了数据集的类别，我们需要将其改成kitti数据集的类别，如下图：

METAINFO中的内容为：

METAINFO = {
    'classes':
    ('Car', 'Pedestrian', 'Cyclist'),
    'palette':
    [(220, 20, 60), (119, 11, 32), (0, 0, 142)]
}

2、修改class_names.py中的类别

mmdetection/mmdet/evaluation/functional/class_names.py中定义了评估时使用的coco数据集的标签，我们也将其改成kitti的标签。

def coco_classes() -> list:
    """Class names of COCO."""
    return [
        'Car','Pedestrian','Cyclist'
    ]

3、在coco_detection.py代码中修改数据集的根路径（这里也可以使用软连接，有兴趣的同学可以bd一下，也很好用）

mmdetection/configs/base/datasets/coco_detection.py文件中定义了coco数据集训练时的一些配置，其中data_root代表数据集的根路径，我们将其修改成我们刚才处理好的coco格式的kitti数据集的根路径。

data_root = 'data/kitti_coco/'	# 这里请修改成你自己的数据集的路径

4、修改atss的配置文件，将其类别输出维度修改成数据集类别个数

mmdetection/configs/atss/atss_r50_fpn_1x_coco.py，由于该模型原本使用的是coco数据集，coco数据集有80个类别，因此最终的分类头输出的是维度是80，但是我们的kitti数据集只有3个类别，因此要将其输出维度修改成3.

至此，我们就已经修改完毕啦。接下来运行MMDetection的训练脚本就可以开始训练了，CUDA_VISIBLE_DEVICES是选择使用第几张显卡进行训练，我这里选择的是1号显卡，如果只有一块显卡可以不用写这个参数。最终的训练记录会保存在work-dir指定的目录下：

CUDA_VISIBLE_DEVICES=1 python tools/train.py configs/atss/atss_r50_fpn_1x_coco.py --work-dir atts_outputs/

总结

由于更换了数据集，在COCO上使用的学习率可能过大导致训练不能收敛（loss不降低），或者直接跑飞（loss出现nan）的情况。如果出现这种情况，我们调低学习率即可，我是用的学习率是0.004.也是在mmdetection/configs/atss/atss_r50_fpn_1x_coco.py配置文件中修改.