最近想跑一下旷世开源的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