实践教程｜如何用YOLOX训练自己的数据集？

↑ 点击 蓝字  关注极市平台

作者 | JuLec@知乎（已授权）

来源 | https://zhuanlan.zhihu.com/p/402210371 

编辑 | 极市平台

极市导读

 

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:"

如果觉得有用，就请分享到朋友圈吧！

△点击卡片关注极市平台，获取 最新CV干货

公众号后台回复“transformer”获取最新Transformer综述论文下载～

极市干货

课程/比赛： 珠港澳人工智能算法大赛 ｜保姆级零基础人工智能教程

算法trick ： 目标检测比赛中的tricks集锦 ｜ 从39个kaggle竞赛中总结出来的图像分割的Tips和Tricks

技术综述： 一文弄懂各种loss function ｜ 工业图像异常检测最新研究总结（2019-2020）

# CV技术社群邀请函 #

△长按添加极市小助手

添加极市小助手微信（ID : cvmart4）

备注：姓名-学校/公司-研究方向-城市（如：小极-北大-目标检测-深圳）

即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群

每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~

觉得有用麻烦给个在看啦~