作者 | JuLec@知乎（已授权）
来源 | https://zhuanlan.zhihu.com/p/402210371
编辑 | 极市平台
导读
Yolo系列因为其灵活性，一直是目标检测热门算法。无奈用它训练自己的数据集有些不好用，于是有空就搞了一下，训练自己的数据集。
代码：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:"
本文仅做学术分享，如有侵权，请联系删文。
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
