Press "Enter" to skip to content

【目标检测】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] + '
'
    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

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注