MMDetection 目标检测使用

OpenMMLab 是香港中文大学-商汤科技联合实验室 MMLab 自 2018 年 10 月开源的一个计算机视觉领域的 AI 算法框架。其包含了众多计算机视觉多个方向的算法框架，本篇介绍 MMDetection 库，运行服务器 Ubuntu 18.04。

MMDetection 的安装需要在一些基础库之上进行，如 pytorch，mmcv 等。假设显卡驱动、cuda、cudnn 等均已安装配置好，接下来就是安装 python 和一些包。我们这里采用 miniconda 来进行 python 安装和环境配置。

# 安装 miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
chmod +x miniconda.sh
bash ./miniconda.sh -b -p /opt/miniconda

# 配置 miniconda 环境
echo "export MINICONDA_HOME=/opt/miniconda" >> ~/.bashrc
echo "export PATH=$MINICONDA_HOME/bin:$PATH" >> ~/.bashrc
source ~/.bashrc

# 安装虚拟环境
conda create -n mm python=3.9
conda activate mm

# 安装 pytorch，假设这里的 cuda 版本是 11.1
pip install torch==1.10.0+cu111 torchvision==0.11.0+cu111 torchaudio==0.10.0 -f https://download.pytorch.org/whl/torch_stable.html

# 安装 mmcv，注意这里的 mmcv 索引链接地址中需要根据上面 cuda 和 torch 版本进行下载安装
pip install mmcv -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.10.0/index.html

# 安装 MMDetection
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection
pip install -e .

打开 Python 命令提示符，输入如下内容，如果正常显示，说明 mmdetection 安装成功：

import mmdet
print(mmdet.__version__)

使用 MMDetection 进行图片分类

安装成功后，在克隆的仓库里有源代码，里面包含了很多算法的配置信息。我们这里以 faster-rcnn 为例，进行图片目标检测应用。

首先需要下载预训练好的模型，访问网址：mmdetection-model_zoo，选择 faster-rcnn:

cd mmdetection
mkdir checkpoints
wget https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_1x_coco/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth -P checkpoints

然后，进行图片目标检测（假设下面代码运行目录是 mmdetection）：

# 导入依赖包
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
# 指定配置文件
config = 'configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
checkpoint = 'checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
# 指定设备，这里使用第一块GPU
device = "cuda:0"
# 初始化模型
model = init_detector(config, checkpoint, device=device)
# 模型推断
img = 'demo/demo.jpg'
result = inference_detector(model, img)
# 把结果画在图片上
show_result_pyplot(model, img, result, score_thr=0.9)

rpn 区域提议

首先需要下载预训练好的模型，访问网址：mmdetection-model_zoo，选择 rpn:

cd mmdetection
wget https://download.openmmlab.com/mmdetection/v2.0/rpn/rpn_r50_fpn_1x_coco/rpn_r50_fpn_1x_coco_20200218-5525fa2e.pth -P checkpoints

然后，进行图片目标检测（假设下面代码运行目录是 mmdetection）：

# 导入依赖包
from mmdet.apis import inference_detector, init_detector, show_result_pyplot
# 指定配置文件
config = 'configs/rpn/rpn_r50_fpn_1x_coco.py'
checkpoint = 'checkpoints/rpn_r50_fpn_1x_coco_20200218-5525fa2e.pth'
# 指定设备，这里使用第一块GPU
device = "cuda:0"
# 初始化模型
model = init_detector(config, checkpoint, device=device)
# 模型推断
img = 'demo/demo.jpg'
rpn_result = inference_detector(model, img)
# 把结果画在图片上
model.show_result(img, rpn_result, top_k=100)

使用 mmdetection 模型微调

现在我们模型框架、预训练模型都有了，只需要使用自己的数据集进行再次训练，使得模型能够更加准确的在我们自己数据上进行预测。

对于个人数据集，有三种组织数据集方法来结合 mmdetection 进行模型微调：

1. 把数据集转化为 CoCo 组织形式；
2. 使用 mmdetection 提供的数据类型 CustomDataset，把数据写成 pkl 保存本地使用；
3. 继承 mmdetection 提供的数据类型 CustomDataset，撰写自己的数据类型，省去 pkl，节省本地空间增加运行速度。

首先，下载数据集，这里我们采用 kitti_tiny 数据集:

wget wget https://download.openmmlab.com/mmdetection/data/kitti_tiny.zip -P data

其次，将下载的数据保存并解压到目录下：

cd data
unzip -q kitti_tiny.zip

最后，根据不同的数据组织形式进行配置文件生产和模型训练。

COCO 格式

请自行组织。

CustomDataset 格式

把数据转化为中间 pkl 格式：

import os.path as osp
import numpy as np
import mmcv

def convert_kitti_to_middle(ann_file, out_file, img_prefix):
    CLASSES = ('Car', 'Pedestrian', 'Cyclist')
    # 类别反查表
    cat2label = {k: i for i, k in enumerate(CLASSES)}
    # 图像列表
    image_list = mmcv.list_from_file(ann_file)
    # 所有图像和标注的信息存储在一个列表中
    data_infos = []
    # convert annotations to middle format
    for image_id in image_list:
        filename = f'{img_prefix}/{image_id}.jpeg'
        image = mmcv.imread(filename)
        height, width = image.shape[:2]

# 单张图像的信息存储在字典中
        data_info = dict(filename=f'{image_id}.jpeg', width=width, height=height)
        # 读取标注信息，并处理
        label_prefix = img_prefix.replace('image_2', 'label_2')
        lines = mmcv.list_from_file(osp.join(label_prefix, f'{image_id}.txt'))
        content = [line.strip().split(' ') for line in lines]
        bbox_names = [x[0] for x in content]
        bboxes = [[float(info) for info in x[4:8]] for x in content]

gt_bboxes = []
        gt_labels = []
        gt_bboxes_ignore = []
        gt_labels_ignore = []

# filter 'DontCare'
        for bbox_name, bbox in zip(bbox_names, bboxes):
            if bbox_name in cat2label:
                gt_labels.append(cat2label[bbox_name])
                gt_bboxes.append(bbox)
            else:
                gt_labels_ignore.append(-1)
                gt_bboxes_ignore.append(bbox)
        # 将标注信息（坐标和标签）转换为 ndarray
        data_anno = dict(
            bboxes = np.array(gt_bboxes, dtype=np.float32).reshape(-1, 4),
            labels = np.array(gt_labels, dtype=np.longlong),
            bboxes_ignore = np.array(gt_bboxes_ignore, dtype=np.float32).reshape(-1, 4),
            labels_ignore = np.array(gt_labels_ignore, dtype=np.longlong)
        )
        data_info.update(ann=data_anno)
        # 所有图像和标注的信息存储在一个列表中
        data_infos.append(data_info)
    mmcv.dump(data_infos, out_file)


convert_kitti_to_middle('data/kitti_tiny/train.txt', 'data/kitti_tiny/train_middle.pkl', 'data/kitti_tiny/training/image_2')
convert_kitti_to_middle('data/kitti_tiny/val.txt', 'data/kitti_tiny/val_middle.pkl', 'data/kitti_tiny/training/image_2')

生产配置文件：

from mmcv import Config
from mmdet.apis import set_random_seed

# 获取基本配置文件参数
cfg = Config.fromfile("configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py")

cfg.device='cuda'

# 修改数据类型以及文件路径
classes = ('Car', 'Pedestrain', 'Cyclist')
cfg.dataset_type = 'CustomDataset'
cfg.data_root = 'data/kitti_tiny'
cfg.classes = classes

dtype = 'CustomDataset'
droot = 'data/kitti_tiny'

cfg.data.test.type = dtype
cfg.data.test.data_root = droot
cfg.data.test.ann_file = 'train_middle.pkl'
cfg.data.test.img_prefix = 'training/image_2'
cfg.data.test.classes = classes

cfg.data.train.type = dtype
cfg.data.train.data_root = droot
cfg.data.train.ann_file = 'train_middle.pkl'
cfg.data.train.img_prefix = 'training/image_2'
cfg.data.train.classes = classes

cfg.data.val.type = dtype
cfg.data.val.data_root = droot
cfg.data.val.ann_file = 'val_middle.pkl'
cfg.data.val.img_prefix = 'training/image_2'
cfg.data.val.classes = classes

# 修改 bbox_head 中的类别数
cfg.model.roi_head.bbox_head.num_classes = 3
# 使用预训练好的 faster_rcnn 模型用于 fine tuning
cfg.load_from = 'checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco-5324cff8.pth'
# 设置工作目录用于存放日志和临时文件
cfg.work_dir = 'work_dir'

if not os.path.exists(cfg.work_dir):
    os.path.makedirs(cfg.work_dir)

# 原本的学习率是在8个GPU卡基础上训练设置的，现在单卡需要除以8
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# 由于是自定义数据集，需要修改评价方法
cfg.evaluation.metric = 'mAP'
# 设置evaluation间隔减少运行时间
cfg.evaluation.interval = 12
# 设置存档点间隔减少存储空间的消耗
cfg.checkpoint_config.interval = 12

# 固定随机种子是的结果可复现
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# 打印所有的配置参数
print(f"Config: \n {cfg.pretty_text}")

cfg.dump(F'{cfg.work_dir}/customformat.py')

训练模型：

import mmcv
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector
import os.path as osp
from mmcv import Config

cfg = Config.fromfile("work_dir/customformat.py")

# 构建数据集
datasets = [build_dataset(cfg.data.train)]

# 构建检测模型
model = build_detector(cfg.model)

# 添加类别文字属性提高可视化效果
model.CLASSES = datasets[0].CLASSES

# 创建工作目录并测试训练模型
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
train_detector(model, datasets, cfg, distributed=False, validate=True)

自定义 KittiTinyDataset 格式

import os.path as osp
import mmcv
import numpy as np
from mmdet.datasets.builder import DATASETS
from mmdet.datasets.custom import CustomDataset

from mmcv import Config
from mmdet.apis import set_random_seed


# 注册表的装饰器，这样 builder 根据配置文件构建时
# 就可以识别 kittiTinyDataset 数据集类型
# 继承 CustomDataset 并重载 load_annotations 函数
@DATASETS.register_module()
class KittiTinyDataset(CustomDataset):
    CLASSES = ('Car', 'Pedestrain', 'Cyclist')

def load_annotations(self, ann_file):
        cat2label = {k: i for i, k in enumerate(self.CLASSES)}
        # load image list from file
        image_list = mmcv.list_from_file(self.ann_file)

data_infos = []
        # convert annotations to middle format
        for image_id in image_list:
            filename = f"{self.img_prefix}/{image_id}.jpeg"
            image = mmcv.imread(filename)
            height, width = image.shape[:2]

data_info = dict(filename=f'{image_id}.jpeg', width=width, height=height)
            # load_annotation
            label_prefix = self.img_prefix.replace('image_2', 'label_2')
            lines = mmcv.list_from_file(osp.join(label_prefix, f'{image_id}.txt'))

content = [line.strip().split(' ') for line in lines]
            bbox_names = [x[0] for x in content]
            bboxes = [[float(info) for info in x[4:8]] for x in content]

gt_bboxes = []
            gt_labels = []
            gt_bboxes_ignore = []
            gt_labels_ignore = []

# filter 'DontCare'
            for bbox_name, bbox in zip(bbox_names, bboxes):
                if bbox_name in cat2label:
                    gt_labels.append(cat2label[bbox_name])
                    gt_bboxes.append(bbox)
                else:
                    gt_labels_ignore.append(-1)
                    gt_bboxes_ignore.append(bbox)
            data_anno = dict(
                bboxes=np.array(gt_bboxes, dtype=np.float32).reshape(-1, 4),
                labels=np.array(gt_labels, dtype=np.longlong),
                bboxes_ignore=np.array(gt_bboxes_ignore, dtype=np.float32).reshape(-1, 4),
                labels_ignore=np.array(gt_labels_ignore, dtype=np.longlong)
            )
            data_info.update(ann=data_anno)
            data_infos.append(data_info)

return data_infos

# 获取基本配置文件参数
cfg = Config.fromfile("configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py")

cfg.device='cuda'

# 修改数据类型以及文件路径
classes = ('Car', 'Pedestrain', 'Cyclist')
cfg.dataset_type = 'KittiTinyDataset'
cfg.data_root = 'data/kitti_tiny'
cfg.classes = classes

dtype = 'KittiTinyDataset'
droot = 'data/kitti_tiny'

cfg.data.test.type = dtype
cfg.data.test.data_root = droot
cfg.data.test.ann_file = 'train.txt'
cfg.data.test.img_prefix = 'training/image_2'
cfg.data.test.classes = classes

cfg.data.train.type = dtype
cfg.data.train.data_root = droot
cfg.data.train.ann_file = 'train.txt'
cfg.data.train.img_prefix = 'training/image_2'
cfg.data.train.classes = classes

cfg.data.val.type = dtype
cfg.data.val.data_root = droot
cfg.data.val.ann_file = 'val.txt'
cfg.data.val.img_prefix = 'training/image_2'
cfg.data.val.classes = classes

# 修改 bbox_head 中的类别数
cfg.model.roi_head.bbox_head.num_classes = 3
# 使用预训练好的 faster_rcnn 模型用于 fine tuning
cfg.load_from = 'checkpoints/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco-5324cff8.pth'
# 设置工作目录用于存放日志和临时文件
cfg.work_dir = 'work_dir_2'

if not os.path.exists(cfg.work_dir):
    os.path.makedirs(cfg.work_dir)

# 原本的学习率是在8个GPU卡基础上训练设置的，现在单卡需要除以8
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

# 由于是自定义数据集，需要修改评价方法
cfg.evaluation.metric = 'mAP'
# 设置evaluation间隔减少运行时间
cfg.evaluation.interval = 12
# 设置存档点间隔减少存储空间的消耗
cfg.checkpoint_config.interval = 12

# 固定随机种子是的结果可复现
cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

# 打印所有的配置参数
print(f"Config: \n {cfg.pretty_text}")

cfg.dump(F'{cfg.work_dir}/customKittiFormat.py')

import mmcv
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector
import os.path as osp
from mmcv import Config

# 构建数据集
datasets = [build_dataset(cfg.data.train)]

# 构建检测模型
model = build_detector(cfg.model)

# 添加类别文字属性提高可视化效果
model.CLASSES = datasets[0].CLASSES

# 创建工作目录并测试训练模型
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
train_detector(model, datasets, cfg, distributed=False, validate=True)