【目标检测】YOLOX训练王者荣耀数据集
前言
最近想跑一下旷世开源的YOLOX，本想着属于YOLO系列，代码大致都和YOLOv5差不多，没想到代码整体差异还是挺大的，跑通的过程中踩了不少坑，这篇就来记录一下整个流程。
官方源码：https://github.com/Megvii-BaseDetection/YOLOX
项目整体框架
项目框架和YOLOv6比较相似，训练、测试、检测三个核心函数都封装在tools文件夹中。
另外项目里有个dataset文件夹用来专门存放数据集，对于源码不熟悉建议先把数据集放在这个文件夹下。
数据预处理
官方的tutorials包含了CoCo和VOC两个经典数据集的训练方式，本篇主要采用的是VOC数据集的训练方式。
首先，需要将自己的数据集按照VOC的格式进行放置，注意文件名不要变化。
之后，在Main文件夹下新建下面两个txt文件。
运行下面的脚本[2]，按照9：1的比例划分训练验证集和测试集。
import os
import random
trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = r'C:\Users\xy\Desktop\Work\YOLOX\datasets\VOC\VOCdevkit\VOC2007\Annotations'
txtsavepath = r'C:\Users\xy\Desktop\Work\YOLOX\datasets\VOC\VOCdevkit\VOC2007\ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftest = open(r'C:\Users\xy\Desktop\Work\YOLOX\datasets\VOC\VOCdevkit\VOC2007\ImageSets\Main\test.txt', 'w')
ftrain = open(r'C:\Users\xy\Desktop\Work\YOLOX\datasets\VOC\VOCdevkit\VOC2007\ImageSets\Main\trainval.txt', 'w')
for i in list:
name = total_xml[i][:-4] + '\n'
if i in trainval:
ftest.write(name)
else:
ftrain.write(name)
ftrain.close()
ftest.close()
修改配置文件
首先修改yolox/data/dataloading.py文件：
再修改exps/example/yolox_voc/yolox_voc_s.py文件：
修改数据集类别个数：
修改训练验证集路径：
修改测试集路径：
然后，修改yolox/data/datasets/voc_classes.py：
这里修改成自己数据集的类别：
接着，修改yolox/evaluators/voc_eval.py ：
该文件用于读取xml标签信息
然后，修改yolox\exp\yolox_base.py
这里包含了很多训练的超参数(很奇怪开发者为什么不按照惯例把它写在train.py中)
其中，我主要修改这几个量：
- self.max_epoch = 300：指定训练epoch
- self.print_interval = 10：每经过10个batch打印一次
- self.eval_interval = 1：每训练一次进行验证
- self.data_num_workers = 0 ：只使用主线程
最后，修改主程序tools/train.py
主要修改这三个量：
下面也可以修改"-c", "--ckpt"来指定官方提供的预训练模型，不指定则从0开始训练。
运行train.py就可以成功开始训练了。
如果还遇到一些报错，可以尝试在文件根目录下运行python setup.py install来更新一些依赖。
检测
训练完之后，保存的模型会存放在tools/YOLOX_outputs/yolox_voc_s/best_ckpt.pth
注意检测前还需要修改类别标签，YOLOX和YOLOv6一样，类别标签并不保存在模型之中，而是在检测时再进行映射，猜测这么做可能是可以减小模型体积，加快推理速度。
在yolox/data/datasets/__init__.py中添加from .voc_classes import VOC_CLASSES
然后在demo.py中，将所有COCO_CLASSES替换成VOC_CLASSES
检测图片
python tools/demo.py --save_result -f exps/example/yolox_voc/yolox_voc_s.py -c tools/YOLOX_outputs/yolox_voc_s/best_ckpt.pth --device gpu --tsize 640 --path assets/dog.jpg
检测视频
python tools/demo.py --save_result -f exps/example/yolox_voc/yolox_voc_s.py -c tools/YOLOX_outputs/yolox_voc_s/best_ckpt.pth --device gpu --tsize 640 --demo video --path C:\Users\xy\Desktop\sucai\wangzhe\zhuge.mp4
拓展：添加帧率显示
对于视频，如果需要显示帧率，可以进行如下操作：
在visual函数中添加：
def visual(self, output, img_info, cls_conf=0.35):
ratio = img_info["ratio"]
img = img_info["raw_img"]
if output is None:
return img
output = output.cpu()
bboxes = output[:, 0:4]
# preprocessing: resize
bboxes /= ratio
cls = output[:, 6]
scores = output[:, 4] * output[:, 5]
vis_res = vis(img, bboxes, scores, cls, cls_conf, self.cls_names)
# 添加帧率检测
cv2.putText(img, "FPS:{:.1f}".format(1. / (time.time() - self.tt)), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2,
(225, 221, 156), 4)
self.tt = time.time()
return vis_res
在Predictor中添加self.tt = time.time()
class Predictor(object):
def __init__(
self,
model,
exp,
cls_names=VOC_CLASSES,
trt_file=None,
decoder=None,
device="cpu",
fp16=False,
legacy=False,
):
self.model = model
self.cls_names = cls_names
self.decoder = decoder
self.num_classes = exp.num_classes
self.confthre = exp.test_conf
self.nmsthre = exp.nmsthre
self.test_size = exp.test_size
self.device = device
self.fp16 = fp16
self.preproc = ValTransform(legacy=legacy)
self.tt = time.time()
结果展示
我使用网络上搜集到的王者荣耀数据集进行训练，再检测自己的一局对局视频，和YOLOv5/YOLOv7/YOLOR进行对比，效果如下
YOLOv5/YOLOv7/YOLOX/YOLOR对比测试
B站Link：https://www.bilibili.com/video/BV1zU4y1C7Kn
数据集从网络搜集：https://aistudio.baidu.com/aistudio/datasetdetail/165546
总结
从各算法的对比检测效果来看，YOLOX的帧率是最快的，但是这里面也需要考虑到YOLOX未渲染中文带来的影响。不过，从视觉上看，YOLOX的检测精度不如其它算法，这也可能是数据集较少导致的。
YOLOX的学习成本比YOLOv5要高，里面需要修改的参数太多，可能是企业更多考虑了技术的版权问题或是单纯地在秀技，总之，写成这样，对于初学者真不友好。
References
[1]https://zhuanlan.zhihu.com/p/460677014
[2]https://blog.csdn.net/weixin_44123583/article/details/124194355