极市平台
极市导读
Yolo系列因为其灵活性，一直是目标检测热门算法。无奈用它训练自己的数据集有些不好用，于是有空就搞了一下，训练自己的数据集。 >>加入极市CV技术交流群，走在计算机视觉的最前沿
代码：https://github.com/Megvii-BaseDetection/YOLOX
论文：https://arxiv.org/abs/2107.08430
Yolo系列因为其灵活性，一直是目标检测热门算法。无奈用它训练自己的数据集有些不好用，于是有空就搞了一下，训练自己的数据集。
1.安装YOLOX
git clone git@github.com:Megvii-BaseDetection/YOLOX.git
cd YOLOX
pip3 install -U pip && pip3 install -r requirements.txt
pip3 install -v -e . # or python3 setup.py develop
pip3 install cython; pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
2.下载预训练权重
https://github.com/Megvii-BaseDetection/YOLOX/blob/main/exps/default/yolox_s.py
3.准备自己的Voc数据集
-----datasets
------VOCdevkit
------DATA_NAME # 你自己存储数据集的文件夹名称
------JPEGImages
------000000000000000.jpg
------Annotations
------000000000000000.xml
------ImageSets
-------Main
------trainval.txt
------test.txt
4.配置文件编辑（config.yaml）
CLASSES:
- person # 数据集的标签，本教程只检测人
CLASSES_NUM: 1 # 待检测的类别个数
SUB_NAME: 'custom' # 上一步中的DATA_NAME
5.修改yolox文件，适配自己的数据集
5.1
首先在exps/example/yolox_voc/yolox__voc_s.py文件最前面写入下面的代码，主要是采用yaml解析config.yaml获得SUB_NAME
import sys
sys.path.insert(1,"../../")
# parseYaml库是自己编写的用于解析yaml
import parseYaml
cfg = parseYaml.get_config("./config.yaml")
DATA_NAME = cfg.SUB_NAME
注：parseYaml脚本如下：
import yaml
import os
from easydict import EasyDict as edict
class YamlParser(edict):
""" This is yaml parser based on EasyDict.
"""
def __init__(self, cfg_dict=None, config_file=None):
if cfg_dict is None:
cfg_dict = {}
if config_file is not None:
assert(os.path.isfile(config_file))
with open(config_file, 'r') as fo:
cfg_dict.update(yaml.load(fo.read(),Loader=yaml.FullLoader))
super(YamlParser, self).__init__(cfg_dict)
def merge_from_file(self, config_file):
with open(config_file, 'r') as fo:
self.update(yaml.load(fo.read()))
def merge_from_dict(self, config_dict):
self.update(config_dict)
def get_config(config_file=None):
return YamlParser(config_file=config_file)
5.2 修改voc_classes.py
cfg = parseYaml.get_config("./config.yaml")
if cfg.CUSTOM:
VOC_CLASSES = cfg.CLASSES
else:
VOC_CLASSES = (
"person",
"aeroplane",
"bicycle",
"bird",
"boat",
"bus",
"bottle",
"car",
"cat",
"chair",
"cow",
"diningtable",
"dog",
"horse",
"motorbike",
"pottedplant",
"sheep",
"sofa",
"train",
"tvmonitor",
)
5.3
修改Exp类的_init__方法，主要是采用yaml解析获得CLASS__NUM
def __init__(self):
super(Exp, self).__init__()
self.num_classes = cfg.CLASSES_NUM # 获得检测的类别个数
self.depth = 0.33
self.width = 0.50
self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
5.4 修改数据加载过程
dataset = VOCDetection(
data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
# image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
image_sets=[(DATA_NAME, 'trainval')], # 适配自己的数据集名称
img_size=self.input_size,
preproc=TrainTransform(
rgb_means=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_labels=50,
),
custom=True, # 新增custom参数
)
5.5
根据5.3中的custom参数，修改voc.py中的VOCDetection的_init_方法
class VOCDetection(Dataset):
def __init__(
self,
data_dir,
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
img_size=(416, 416),
preproc=None,
target_transform=AnnotationTransform(),
dataset_name="VOC0712",
custom = True # 新增
):
super().__init__(img_size)
self.root = data_dir
self.image_set = image_sets
self.img_size = img_size
self.preproc = preproc
self.target_transform = target_transform
self.name = dataset_name
self._annopath = os.path.join("%s", "Annotations", "%s.xml")
self._imgpath = os.path.join("%s", "JPEGImages", "%s.jpg")
self._classes = VOC_CLASSES
self.ids = list()
self.custom = custom
if self.custom: # 处理自己的数据集
self.base_dir,self.custom_name = image_sets[0] # DATA_NAME
rootpath = os.path.join(self.root, self.base_dir)
for line in open(
os.path.join(rootpath, "ImageSets", "Main", self.custom_name + ".txt")
):
self.ids.append((rootpath, line.strip()))
else: # 处理默认的Voc数据集
for (year, name) in image_sets:
self._year = year
rootpath = os.path.join(self.root, "VOC" + year)
for line in open(
os.path.join(rootpath, "ImageSets", "Main", name + ".txt")
):
self.ids.append((rootpath, line.strip()))
5.6 修改get_eval_loader方法
valdataset = VOCDetection(
data_dir=os.path.join(get_yolox_datadir(), "VOCdevkit"),
# image_sets=[('2007', 'test')],
image_sets=[(DATA_NAME, 'test')],
img_size=self.test_size,
preproc=ValTransform(
rgb_means=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
),
custom=True,
)
6.执行训练
python tools/train.py -f exps/example/yolox_voc/yolox_voc_s.py
-expn TEST
-d 4
-b 64
--fp16
-o
-c weights/yolox_s.pth
7.执行推理验证
python tools/demo.py image/video/webcam
-f exps/example/yolox_voc/yolox_voc_s.py
-c YOLOX_outputs/yolox_voc_s/best_ckpt.pth.tar
--path img/1.jpg
--conf 0.25
--nms 0.45
--tsize 640
--save_result
--device gpu
# if choose webcam
--camid 0/"rtsp:"
如果觉得有用，就请分享到朋友圈吧！
公众号后台回复“transformer”获取最新Transformer综述论文下载～
# CV技术社群邀请函 #
备注：姓名-学校/公司-研究方向-城市（如：小极-北大-目标检测-深圳）
即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群
每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~
