mmdetection 用自定义 coco 数据集做目标检测

1.1 mmdetction 安装

1.1.1 系统环境需求

参考 mmdetection 官方文档 https://mmdetection.readthedocs.io/en/latest/INSTALL.html，系统环境需求如下：

• Linux (Windows is not officially supported)
• Python 3.5+
• PyTorch 1.1 or higher
• CUDA 9.0 or higher
• NCCL 2
• GCC 4.9 or higher
• mmcv

我的系统环境：

• CentOS 7.2
• Python 3.7
• PyTorch 1.1.0
• CUDA 10.0
• NCCL 2
• GCC 7.4.0
• mmcv

安装 mmcv 时的依赖项如下：

addict
numpy
pyyaml
six

其中，我的环境中缺少 addict 和 pyyaml，从 https://pypi.org/ 中下载源码离线安装。

附：查看深度学习软件/库/工具的命令速查表：

https://tech.amikelive.com/node-841/command-cheatsheet-checking-versions-of-installed-software-libraries-tools-for-deep-learning-on-ubuntu-16-04/

1.1.2 安装 mmdetection

官方文档的安装说明如下，适合网络环境好的条件下进行在线安装，

# Clone the mmdetection repository
git clone https://github.com/open-mmlab/mmdetection.git
cd mmdetection

# Install build requirements and then install mmdetection.
pip install -r requirements/build.txt
pip install "git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI"
pip install -v -e .  # or "python setup.py develop"

这里采用离线安装方式，
（1）检查 mmdetection 的依赖项，满足要求。

cat requirements/build.txt
# These must be installed before building mmdetection
numpy
torch>=1.1

（2）安装 cocoapi，

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py install

（3）安装 mmdetection，依赖项检查结果，terminaltables 没有，采用源码离线安装。

安装完成以后，打开 Python 进行验证，

from mmdet.apis import init_detector, inference_detector, show_result

1.2 训练自定义数据集 CatDog

1.2.1 准备数据集

创建 data 文件夹，软链接到数据集根目录，其中标记好的数据集采用 coco 数据格式，

cd mmdetection
mkdir data
export COCO_ROOT=/data1/Projects/datasets/coco
ln -s $COCO_ROOT data

在 mmdet/datasets/cat_dog.py:

from .coco import CocoDataset
from .registry import DATASETS

@DATASETS.register_module
class CatDog(CocoDataset):
    CLASSES = ('dog', 'cat')

在 mmdet/datasets/__init__.py:

from .cat_dog import CatDog

1.2.2 修改 faster_rcnn 模型配置

下载 resnet50 的预训练模型，放入 $TORCH_HOME，

export TORCH_HOME=/data1/Projects/pretrained_models
echo $TORCH_HOME

mkdir -p /data1/Projects/pretrained_models/checkpoints/
mv resnet50-19c8e357.pth /data1/Projects/pretrained_models/checkpoints/

用 faster_rcnn 做为模型进行目标检测，拷贝一份 configs/faster_rcnn_r50_fpn_1x.py 为 configs/cat_dog_faster_rcnn_r50_fpn_1x.py，在 cat_dog_faster_rcnn_r50_fpn_1x.py，加入预训练的 resnet50 加载路径，

# model settings
import os

os.environ['TORCH_HOME'] = '/data1/Projects/pretrained_models'

在 config 文件 configs/cat_dog_faster_rcnn_r50_fpn_1x.py 使用 CatDog 数据集，分类数 3 包括狗，猫和背景，

num_classes=3,

CatDog 数据集为类 coco 数据集，针对 coco 数据集修改：

• 定义数据种类，修改 mmdetection/mmdet/datasets/coco.py，把 CLASSES 的 tuple 改为自己数据集对应的种类。

CLASSES = ('dog', 'cat')

• 在 mmdetection/mmdet/core/evaluation/class_names.py 修改 coco_classes 数据集类别，这个关系到后面 test 的时候结果图中显示的类别名称。

def coco_classes():
    return [ 'dog', 'cat']

1.2.3 训练模型

使用单个 GPU 进行训练，指定 --work_dir 保存模型结果，

python tools/train.py configs/cat_dog_faster_rcnn_r50_fpn_1x.py \
--gpus 1 \
--work_dir './work_dirs/cat_dog_faster_rcnn_r50_fpn_1x'

1.2.4 测试图片

1.2.4.1 测试单张图片

import numpy as np
from mmdet.apis import init_detector, inference_detector
import mmcv
import cv2

threshold = 0.9  # confidence score
config_file = './configs/cat_dog_faster_rcnn_r50_fpn_1x.py'
checkpoint_file = './work_dirs/cat_dog_faster_rcnn_r50_fpn_1x/latest.pth'
# 通过配置文件（config file）和模型文件（checkpoint file）构建检测模型
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# 测试单张图片并展示结果
img_path = '/data1/Projects/datasets/cat_dog_single/cat.12176.jpg'
result = inference_detector(model, img_path)
bboxes = np.vstack(result)
# print(bboxes.shape)
labels = [
    np.full(bbox.shape[0], i, dtype=np.int32)
    for i, bbox in enumerate(result)
]
labels = np.concatenate(labels)
# print(labels)
img = cv2.imread(img_path)
scores = bboxes[:, -1]
inds = scores > threshold
bboxes = bboxes[inds, :]
# print(bboxes.shape)
labels = labels[inds]
# print(labels)
class_names = model.CLASSES
cat_dog_dict = {}
for label in labels:
    cat_dog_dict[class_names[label]] = cat_dog_dict.get(class_names[label], 0) + 1
print(cat_dog_dict)
for k, v in cat_dog_dict.items():
    print('{0} 有 {1} 只 {2}'.format(img_path, v, k))
for bbox in bboxes:
    left_top = (bbox[0], bbox[1])
    right_bottom = (bbox[2], bbox[3])
    cv2.rectangle(img, left_top, right_bottom, color=(0, 255, 0))
cv2.imwrite('/data1/Projects/datasets/cat_dog_single/res_cat.12176.jpg', img)

1.2.4.2 测试多张图片

from mmdet.apis import init_detector, inference_detector, show_result
import mmcv
import numpy as np
import glob
import os

config_file = 'configs/cat_dog_faster_rcnn_r50_fpn_1x.py'
checkpoint_file = 'work_dirs/cat_dog_faster_rcnn_r50_fpn_1x/latest.pth'
score_thr = 0.9

# 通过配置文件（config file）和模型文件（checkpoint file）构建检测模型
model = init_detector(config_file, checkpoint_file, device='cuda:0')

# 测试单张图片
# img = '/data1/Projects/datasets/cat_dog_single/cat.12176.jpg'
# result = inference_detector(model, img)
# show_result(img, result, model.CLASSES, score_thr=score_thr,
#             out_file='/data1/Projects/datasets/cat_dog_single/res_cat.12176.jpg')

# 测试多张图片

# imgs = glob.glob('/data1/Projects/datasets/coco/val2017/*.jpg')
imgs = glob.glob('/data1/Projects/datasets/test/*.jpg')
for i, img in enumerate(imgs):
    # 画 bounding boxes 到图片上
    # print(i, imgs[i])
    result = inference_detector(model, img)
    file_name = imgs[i].split('/')[-1]
    out_file = os.path.join('/data1/Projects/datasets/test_det', file_name)
    show_result(img, result, model.CLASSES, score_thr=score_thr, out_file=out_file)

    # 输出图片的猫和狗的数量
    img = mmcv.imread(img)
    img = img.copy()
    bbox_result = result
    bboxes = np.vstack(bbox_result)
    labels = [
        np.full(bbox.shape[0], i, dtype=np.int32)
        for i, bbox in enumerate(bbox_result)
    ]
    labels = np.concatenate(labels)

    # 根据阈值调整输出的 bboxes 和 labels
    scores = bboxes[:, -1]
    inds = scores > score_thr
    bboxes = bboxes[inds, :]
    labels = labels[inds]

    class_names = model.CLASSES
    cat_dog_dict = {}
    for label in labels:
        cat_dog_dict[class_names[label]] = cat_dog_dict.get(class_names[label], 0) + 1
    for k, v in cat_dog_dict.items():
        print('{0} 有 {1} 只 {2}'.format(imgs[i].split('/')[-1], v, k))
    print('--------------------')

测试多张图片，输出结果如下：

---

微信公众号「padluo」，分享数据科学家的自我修养，既然遇见，不如一起成长。

---