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.pyconfigs/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('--------------------')

测试多张图片,输出结果如下:

mmdetection测试多张图片输出结果.png
bounding box.png

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

数据分析二维码.gif

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 206,602评论 6 481
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 88,442评论 2 382
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 152,878评论 0 344
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 55,306评论 1 279
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 64,330评论 5 373
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,071评论 1 285
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 38,382评论 3 400
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 37,006评论 0 259
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 43,512评论 1 300
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,965评论 2 325
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,094评论 1 333
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,732评论 4 323
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 39,283评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,286评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,512评论 1 262
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 45,536评论 2 354
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,828评论 2 345

推荐阅读更多精彩内容