YOLOX实战之自建数据集进行训练_yolox制作自己数据集

一、配置环境
ubuntu20.04
rtx 1080ti
cuda 11.2
在配置环境之初由于人习惯，首先创建了一个conda虚拟环境，避免和我之前的环境产生冲突：

conda create -n yolox python=3.6    # 创建环境
source activate yolox             # 激活环境
1
2

创建好环境之后就要对依赖库进行配置

git clone git@github.com:Megvii-BaseDetection/YOLOX.git  #拉取yolox官方源码
cd YOLOX 
pip3 install -U pip 
pip3 install -r requirements.txt              # 在安装requirements.txt依赖库时首先将torch注释掉
python3 setup.py develop  # or  安装yolox
1
2
3
4
5

因为默认requirements.txt安装的是安装torch最新版，由于我在跑代码时遇到了一些问题，所以我自己装了torch1.8版本
新的更新：
由于近期任务加重，采用了分布式训练的，但是cudnn报错了：
RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED
经过排查是pytorch的版本与cuda\dudnn版本不兼容产生的问题，所以我重新更新了pytorch版本：

conda install pytorch==1.10.0 torchvision==0.11.1 torchaudio==0.10.0 cudatoolkit=11.3 -c pytorch
1

此指令为以前版本

~~pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html~~ 
1

然后安装nvidia混合精度库apex：

git clone https://github.com/NVIDIA/apex
cd apex
pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
1
2
3

安装pycocotools：

pip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

1
2

如果安装失败的话可以用以下命令：

 pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI
1

在分布式训练时调用pycocotools画图是出现了如下错误：

ImportError: Cannot load backend ‘TkAgg‘ which requires the ‘tk‘ interactive framework
1

后来发现是matplotlib版本过高导致的，卸载你的环境中matplotlib高版本，下载3.2.1版本，解决报错

pip uninstall matplotlib
pip install matplotlib==3.2.1
1
2

二、demo测试
在github中将权重文件下载下来，下载完成后直接放到yolox的根目录下：

运行如下指令：

python tools/demo.py image -f exps/example/custom/yolox_s.py -c ./latest_ckpt.pth --path assets/dog.jpg --conf 0.3 --nms 0.65 --tsize 640 --save_result --device gpu
1

其中：

image是推理模式，如果是输入视频则为video
-f 模型配置文件
-c 为权重文件地址
–path是测试的图片路径
–conf 置信度阈值
–nms nms的iou阈值
–tsize 测试图片大小
–save_result 是否保存推理结果
注意一定要将模型配置文件中的num_classes进行调整

运行成功示例如下（注意：他这里的demo代码出错了也不会报错，直接跳出终止，所以如果你运行了发现没有如下的运行结果，可能是程序终止了，需要自己排查问题。

三、数据准备和配置修改
官方提供的程序中只支持voc和coco两种格式的数据训练，而我的是yolo格式的，所以需要对格式进行转换，转换代码如下：

from xml.dom.minidom import Document
import os
import cv2


# def makexml(txtPath, xmlPath, picPath):  # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径
def makexml(picPath, txtPath, xmlPath):  # txt所在文件夹路径，xml文件保存路径，图片所在文件夹路径
    """此函数用于将yolo格式txt标注文件转换为voc格式xml标注文件
    """
    dic = {'0': "0",  # 创建字典用来对类型进行转换
           '1': "1",  # 此处的字典要与自己的classes.txt文件中的类对应，且顺序要一致
           }
    files = os.listdir(txtPath)
    print(files)
    for i, name in enumerate(files):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        txtFile = open(txtPath + name)
#        print(txtFile)
        txtList = txtFile.readlines()
#        print(txtList)
        img = cv2.imread(picPath + name[0:-4] + ".jpg")
        print(name[0:-4])
        Pheight, Pwidth, Pdepth = img.shape

        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束

        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(name[0:-4] + ".jpg")
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束

        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束

        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束

        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束

        annotation.appendChild(size)  # size标签结束

        for j in txtList:
            oneline = j.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(dic[oneline[0]])
 #           print(namecontent)
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束

            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束

            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束

            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束

            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束

            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束

            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束

            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束

            object.appendChild(bndbox)  # bndbox标签结束

            annotation.appendChild(object)  # object标签结束

        f = open(xmlPath + name[0:-4] + ".xml", 'w')
        xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        f.close()


if __name__ == "__main__":
    picPath = "model/YOLOX/datasets/VOC/VOCdevkit/VOC2007/JPEGImages/"  # 图片所在文件夹路径，后面的/一定要带上
    txtPath = "model/YOLOX/datasets/VOC/VOCdevkit/VOC2007/labels/lables/"  # txt所在文件夹路径，后面的/一定要带上
    xmlPath = "model/YOLOX/datasets/VOC/VOCdevkit/VOC2007/Annotations/"  # xml文件保存路径，后面的/一定要带上
    makexml(picPath, txtPath, xmlPath)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

以上代码只需要依照自身情况对dic、picPath、txtPath、xmlPath进行更改即可转换。以上是yolo格式转voc格式。下面是yolo格式转coco格式：

"""
YOLO 格式的数据集转化为 COCO 格式的数据集
--root_path 输入根路径
"""

import os
import cv2
import json
from tqdm import tqdm
import argparse
import glob

parser = argparse.ArgumentParser("ROOT SETTING")
parser.add_argument('--root_path', type=str, default='coco', help="root path of images and labels")
arg = parser.parse_args()

# 默认划分比例为 8:1:1。 第一个划分点在8/10处，第二个在9/10。
VAL_SPLIT_POINT = 4 / 5
TEST_SPLIT_POINT = 9 / 10

root_path = arg.root_path
print(root_path)

# 原始标签路径
originLabelsDir = os.path.join(root_path, 'labels/*/*.txt')
# 原始标签对应的图片路径
originImagesDir = os.path.join(root_path, 'images/*/*.jpg')
# dataset用于保存所有数据的图片信息和标注信息
train_dataset = {'categories': [], 'annotations': [], 'images': []}
val_dataset = {'categories': [], 'annotations': [], 'images': []}
test_dataset = {'categories': [], 'annotations': [], 'images': []}

# 打开类别标签
with open(os.path.join(root_path, 'classes.txt')) as f:
    classes = f.read().strip().split()

# 建立类别标签和数字id的对应关系
for i, cls in enumerate(classes, 1):
    train_dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'fish'})
    val_dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'fish'})
    test_dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'fish'})

# 读取images文件夹的图片名称
indexes = glob.glob(originImagesDir)
print(len(indexes))
# ---------------接着将，以上数据转换为COCO所需要的格式---------------
for k, index in enumerate(tqdm(indexes)):

    txtFile = index.replace('images', 'labels').replace('jpg', 'txt')
    # 用opencv读取图片，得到图像的宽和高
    im = cv2.imread(index)
    H, W, _ = im.shape

    # 切换dataset的引用对象，从而划分数据集
    if k + 1 > round(len(indexes) * VAL_SPLIT_POINT):
        if k + 1 > round(len(indexes) * TEST_SPLIT_POINT):
            dataset = test_dataset
        else:
            dataset = val_dataset
    else:
        dataset = train_dataset

    # 添加图像的信息到dataset中

    if (os.path.exists(txtFile)):
        with open(txtFile, 'r') as fr:
            dataset['images'].append({'file_name': index.replace("\\", "/"),
                                      'id': k,
                                      'width': W,
                                      'height': H})
            labelList = fr.readlines()
            for label in labelList:
                label = label.strip().split()
                x = float(label[1])
                y = float(label[2])
                w = float(label[3])
                h = float(label[4])
                # convert x,y,w,h to x1,y1,x2,y2
                # imagePath = os.path.join(originImagesDir,
                #                            txtFile.replace('txt', 'jpg'))
                image = cv2.imread(index)
                x1 = (x - w / 2) * W
                y1 = (y - h / 2) * H
                x2 = (x + w / 2) * W
                y2 = (y + h / 2) * H
                x1 = int(x1)
                y1 = int(y1)
                x2 = int(x2)
                y2 = int(y2)
                # 为了与coco标签方式对，标签序号从1开始计算
                cls_id = int(label[0]) + 1
                width = max(0, x2 - x1)
                height = max(0, y2 - y1)
                dataset['annotations'].append({
                    'area': width * height,
                    'bbox': [x1, y1, width, height],
                    'category_id': int(cls_id),
                    'id': i,
                    'image_id': k,
                    'iscrowd': 0,
                    # mask, 矩形是从左上角点按顺时针的四个顶点
                    'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
                })
            # print(dataset)
            # break
    else:
        continue

# 保存结果的文件夹
folder = os.path.join(root_path, 'annotations')
if not os.path.exists(folder):
    os.makedirs(folder)
for phase in ['train', 'val', 'test']:
    json_name = os.path.join(root_path, 'annotations/{}.json'.format(phase))
    with open(json_name, 'w', encoding="utf-8") as f:
        if phase == 'train':
            json.dump(train_dataset, f, ensure_ascii=False, indent=1)
        if phase == 'val':
            json.dump(val_dataset, f, ensure_ascii=False, indent=1)
        if phase == 'test':
            json.dump(test_dataset, f, ensure_ascii=False, indent=1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

在运行代码之前数据集根目录下放置classes.txt文件。只需要指定–root_path即可进行转换。
转换完成后我们需要对数据集进行分割，分割数据集代码如下：

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
#xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
#数据集的划分，地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

分割完成后我们需要将处理好的数据集按照如下目录进行配置：

├── data #手动创建data、VOCdevkit、VOC2007、Annotations、JPEGImages、ImageSets、Main这些文件夹
│   ├── VOCdevkit
│   │   ├── VOC2007
│   │   │   ├── Annotations #把test.txt、trainval.txt对应的xml文件放在这
│   │   │   ├── JPEGImages #把test.txt、trainval.txt对应的图片放在这
│   │   │   ├── ImageSets
│   │   │   │   ├── Main
│   │   │   │   │   ├── test.txt 
│   │   │   │   │   ├── trainval.txt

1
2
3
4
5
6
7
8
9
10

对yolox/data/datasets/voc_classes.py进行修改，修改成你训练数据的类别（保险起见，最好每一个类别后都加上 ,）：

VOC_CLASSES = (
"panda",
"tiger",
)
1
2
3
4

接着修改yolox/exp/yolox_base.py

将self.num_classes修改为自己的类别数
self.num_classes = 20 (我的数据集是 20)

还可以修改 self.inputsize, self.random_size 改变训练尺寸大小

还可以修改 self.test_size 改变测试的尺寸大小

1
2
3
4
5
6
7
8

注意在训练voc数据集时一定要按照如下格式进行配置,然后我们修改exps/example/yolox_voc/yolox_voc_s.py文件三个地方，把你的类别和文件路径改过来。

下列一定要修改，否则会在训练时出现ap为零的情况

对yolox/data/datasets/voc.py 下的 _do_python_eval 方法进行修改

annopath = os.path.join(rootpath, "Annotations", "{:s}.xml")

修改为

annopath = os.path.join(rootpath, "Annotations", "{}.xml")

1
2
3
4
5
6

修改完过后直接运行如下命令开始训练：

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 1 -b 16 --fp16 -o -c yolox_s.pth

1
2

-d 使用多少张显卡训练
-b 批次大小
–fp16 是否开启半精度训练
查看模型训练结果：
使用命令tensorboard --logdir=event-dir(event路径)
直接命令行或者在event当前路径新建.bat文件，编辑tensorboard --logdir=event-dir，保存双击运行

四、模型测试
1、修改一下demo.py，导入VOC_classes
2、在olox/data/datasets/init.py加入VOC类别

# demo.py 
#开头需要 引入voc类别
from yolox.data.datasets import VOC_CLASSES
predictor = Predictor(model, exp, VOC_CLASSES, trt_file, decoder, args.device)# Predictor 默认类别是COCO_CLASSES
#   yolox/data/datasets/__init__.py 里需要也加入 voc类别
from .voc_classes import VOC_CLASSES



1
2
3
4
5
6
7
8
9

后运行demo.py进行预测

python tools/demo.py image -n yolox-s -c yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device [cpu/gpu]
1

五、模型转换

 python tools/export_onnx.py --output-name yolox_s.onnx -f  exps/default/yolox_s.py -c YOLOX_outputs/yolox_voc_s/best_ckpt.pth

1
2

六、如果在训练中中断运行如下指令即可

python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py -d 2 -b 8 --resume -c /YOLOX_outputs//latest_ckpt.pth

1
2

七、训练COCO数据集

首先是制作自己的COCO数据集：
数据集格式转换见博客：目标检测各种数据集进行互转
转换完成后将数据集文件排布如下：

其中train2017、test2017、val2017为三个子集的图片，annotations为标注信息，具体内容如下：

将制作好的数据集打包成一个包之后，为了我们训练的方便，我们往往采用软连接的形式对数据集进行使用。具体指令如下：

ln -s /path/to/your/COCO ./datasets/COCO
1

使用以上链接时需要切换到工程跟=根目录下方可使用
将数据集制作好之后，我们将工程中的yolox/data/datasets中的coco_classes.py的类别列表定义成自己列表。并且将yolox/exp中的yolox_base。