Mask_YOLO_Caffe:基于Mask-RCNN和YOLO检测的工程实现(附源代码)_github

本文由知乎作者Leon Wu授权转载,不得擅自二次转载。

本文将YOLO和maskRCNN进行了结合,用yolo进行人体检测的同时可以得到instance的mask和keypoints。

首先奉上github代码, 有训练好的demo可以尝试,希望可以值得大家一玩。

​https://github.com/leon-liangwu/MaskYolo_Caffe​

下面是一些简单的介绍:

YOLO最初版本是由Github中关于YOLO的Caffe版本实现有很多,最初yolo出现由Redmon开发的darknet深度学习框。这个框架用纯C语言实现了视觉深度学习中常用的layers,依赖很少。但是如果想要用其他的backbone比如resenet,mobilenet这些网络的话,很难去实现。所以很多同学将YOLO移植到Caffe中,方便使用Caffe中的imagenet模型进行finetune。

实现的功能:Caffe-MaskYolo

  • Demos for object detection, mask segmentation and keypoints recognition
  • YOLO v2 (RegionLossLayer) and v3 (YoloLossLayer) are supported
  • Instance Mask segmentation with Yolo
  • Keypoints Recognition with yolo
  • training data preparation and training

一些samples

Mask_YOLO_Caffe:基于Mask-RCNN和YOLO检测的工程实现(附源代码)_3d_02

Mask_YOLO_Caffe:基于Mask-RCNN和YOLO检测的工程实现(附源代码)_深度学习_03

Mask_YOLO_Caffe:基于Mask-RCNN和YOLO检测的工程实现(附源代码)_github_04

具体实现

YOLO的核心其实是它的Loss layer. darknet 中对应yolo v1到 v3的进化的loss layer分别是:

detection_layer.c
region_layer.c
yolo_layer.c

所以将YOLO移植到Caffe中最重要的就是在Caffe中实现对应的层,这里我实现了V2和V3的caffe 对应的层:

MaskYolo_Caffe/caffe-maskyolo/src/caffe/layers/region_loss_layer.cpp
MaskYolo_Caffe/caffe-maskyolo/src/caffe/layers/yolo_loss_layer.cpp

其次,生成lmdb和数据层的实现:

在github中有详细地介绍如何生成训练使用的lmdb,不详细介绍,大概的方式就是将mask和keypoints的ground truth存在lmdb的image channel中,数据层进行解析,可以看一下数据层的形式:

name: "MobileNet"
layer {
name: "data"
type: "KpsBoxData"
top: "data"
top: "label"
include {
phase: TRAIN
}
transform_param {
scale: 0.0078125
mean_value: 128.0
mean_value: 128.0
mean_value: 128.0
resize_w: 320
resize_h: 224
jitter: 0.2
}
data_param {
source: "/data/Machine_Learning/lmdb/mask_kps_lmdb"
batch_size: 1
backend: LMDB
min_size_and: 0.08
min_size_or: 0.001
}
}

上面这两步已经实现了YOLO的caffe版本,可以进行物体检测了,下面是我如何和MaskRCNN进行结合。实际上关键点就是如何将YOLO的检测结果输出给roi_pooling层。因此我写了一个decode_rois_layer 用来将yolo的检测结果作为proposals,并输入给roi pooling 层。

Demo & Train

# compile the pythonapi of cocoapi
cd ROOT_MaskYolo/lib/cocoapi/PythonAPI
make -j

# use the following command to generate lmdb which contains mask and keypoints information
cd ROOT_MaskYolo
python scripts/createdata_mask_kps.py --coco_dir=/path/to/coco --lmdb_dir=/path/to/lmdb

# the training for mask consists of 2 steps
cd ./models/mb_body_mask

# 1. freeze the weights of detection network, only update the roi mask part
sh train_maskyolo_step1.sh

# 2. update all the network with finetuning the model of step1
sh train_maskyolo_step2.sh

Reference

You Only Look Once: Unified, Real-Time Object :detection https://arxiv.org/abs/1506.00264

YOLO9000: Better, Faster, Stronger :https://arxiv.org/abs/1612.00824

YOLOv3: An Incremental Improvement

Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks:https://arxiv.org/abs/1506.00149

Mask R-CNN:https://arxiv.org/abs/1703.00687