MaskRCNN源码阅读

MaskRCNN源码阅读

前言

这两天读了下MaskRCNN的源码，在github上找的Mask R-CNN for object detection and instance segmentation on Keras and TensorFlow。

之前看过FasterRCNN的源码，不过当时没有做记录。

总的来说，MaskRCNN算是集大成者。其在FasterRCNN的框架上修改而成，backbone由经典的VGG变成了ResNet+FPN(Feature Pyramid Network)，RPN基本上没变，Roi pooling相应的升级成RoiAlign（可以更好的解决mask预测问题），预测头网络在原先的分类预测和回归预测的基础上又添加了一个新分支，用来预测mask。第一个将深度学习应用在语义分割上的应该是FCN(Fully Convolutional Networks)吧。

关于模型的代码实现主要在model.py文件中。

定义模型网络结果的代码主要在MaskRCNN类中的build方法中。推荐先看这个方法，有个概况之后，再细看各个子网络的实现。

输入定义

分别定义了在training和inference模式下的输入

Bottom-up

应该就是ResNet网络

        # Build the shared convolutional layers.
        # Bottom-up Layers
        # Returns a list of the last layers of each stage, 5 in total.
        # Don't create the thead (stage 5), so we pick the 4th item in the list.
        if callable(config.BACKBONE):
            _, C2, C3, C4, C5 = config.BACKBONE(input_image, stage5=True,
                                                train_bn=config.TRAIN_BN)
        else:
            _, C2, C3, C4, C5 = resnet_graph(input_image, config.BACKBONE,
                                             stage5=True, train_bn=config.TRAIN_BN)

Top-down

以P4为例。先将P5上采样，得到与P4长宽一样的特征图，然后将C4进行卷积（横向连接），最后将两者相加（插个题外话，其实，有些论文中的网络是相乘或者其他的操作）。最后在将P4进行次卷积，据说是为了目的是消除上采样的混叠效应（aliasing effect）。

        P4 = KL.Add(name="fpn_p4add")([
            KL.UpSampling2D(size=(2, 2), name="fpn_p5upsampled")(P5),
            KL.Conv2D(config.TOP_DOWN_PYRAMID_SIZE, (1, 1), name='fpn_c4p4')(C4)])

        P4 = KL.Conv2D(config.TOP_DOWN_PYRAMID_SIZE, (3, 3), padding="SAME", name="fpn_p4")(P4)

注意的是，用于RPN的特征图和用于头部网络的特征图并不完全一致。

RPN

FasterRCNN的RPN只需要在接受一个特征图，而如今要接受多尺度的特征图。在得到RPN在不同尺度上输出结果（即在不同特征图的每个点上的每个anchor上的处理结果）后还需要再整理一下，如将不同的尺度上的输出结果中的bbox信息聚合在一块。

        # RPN Model
        rpn = build_rpn_model(config.RPN_ANCHOR_STRIDE,
                              len(config.RPN_ANCHOR_RATIOS), config.TOP_DOWN_PYRAMID_SIZE)
        # Loop through pyramid layers
        layer_outputs = []  # list of lists
        for p in rpn_feature_maps:
            layer_outputs.append(rpn([p]))
        # Concatenate layer outputs
        # Convert from list of lists of level outputs to list of lists
        # of outputs across levels.
        # e.g. [[a1, b1, c1], [a2, b2, c2]] => [[a1, a2], [b1, b2], [c1, c2]]
        output_names = ["rpn_class_logits", "rpn_class", "rpn_bbox"]
        outputs = list(zip(*layer_outputs))
        outputs = [KL.Concatenate(axis=1, name=n)(list(o))
                   for o, n in zip(outputs, output_names)]

        rpn_class_logits, rpn_class, rpn_bbox = outputs

其中，不同bbox的坐标是使用的归一化的坐标，所以在不同特征图上的预测结果是可以混合在一块的。

ProposalLayer

RPN网络输出那么多信息，需要进一步处理，得到有效的proposal。ProposalLayer根据anchor score和NMS进行选择。顺便调整下bbox位置及大小。

        # Generate proposals
        # Proposals are [batch, N, (y1, x1, y2, x2)] in normalized coordinates
        # and zero padded.
        proposal_count = config.POST_NMS_ROIS_TRAINING if mode == "training"\
            else config.POST_NMS_ROIS_INFERENCE
        rpn_rois = ProposalLayer(
            proposal_count=proposal_count,
            nms_threshold=config.RPN_NMS_THRESHOLD,
            name="ROI",
            config=config)([rpn_class, rpn_bbox, anchors])

在最初的时候，一直想不通ProposalLayer层是怎么工作的，因为之前只知道卷积层和pool层，直到看了FasterRCNN的源代码后才知道Layer还可以这么搞。MaskRCNN的ProposalLayer和FasterRCNN的差不多。

DetectionTargetLayer

产生用于训练的样本。大致过程是先去掉数据集中过于拥挤的目标框，计算proposal和gt_boxes(真实的目标框)的重叠情况（矩阵形式），正负样本取样，文中的正负样本比例为2:1。给每个GT boxes和GT masks分配正样本（正ROI），和其他的后续处理。

            # Generate detection targets
            # Subsamples proposals and generates target outputs for training
            # Note that proposal class IDs, gt_boxes, and gt_masks are zero
            # padded. Equally, returned rois and targets are zero padded.
            rois, target_class_ids, target_bbox, target_mask =\
                DetectionTargetLayer(config, name="proposal_targets")([
                    target_rois, input_gt_class_ids, gt_boxes, input_gt_masks])

PyramidROIAlign

ROI Pooling的升级版。原理可以在资料中找到。我原以为实现起来会很复杂，实际上并不是这样的。代码中先选择要在哪个层次的特征图上处理后，使用的是采样点为1的ROIalign（可以使用tf.image.crop_and_resize，该函数支持双线性插值。将最初最初的大自然图像想象成一个2D连续信号，图像经过离散采样后得到了数字图像，记录了每个整数坐标上的信息，整数坐标之间的信息通过插值算法拟合。将这个思路应用在特征图上）

            # Crop and Resize
            # From Mask R-CNN paper: "We sample four regular locations, so
            # that we can evaluate either max or average pooling. In fact,
            # interpolating only a single value at each bin center (without
            # pooling) is nearly as effective."
            #
            # Here we use the simplified approach of a single value per bin,
            # which is how it's done in tf.crop_and_resize()
            # Result: [batch * num_boxes, pool_height, pool_width, channels]
            pooled.append(tf.image.crop_and_resize(
                feature_maps[i], level_boxes, box_indices, self.pool_shape,
                method="bilinear"))

fpn_classifier_graph

和FasterRCNN大同小异，除了采用了ROIAlign之外。

build_fpn_mask_graph

除了采用ROIAlign外，基本上也都是常规操作吧。可以参考FCN网络理解。

DetectionLayer

在inference模式时，模型先进行fpn_classifier_graph处理，然后将其输出输入到DetectionLayer处理后，将其输出作为代替training模式下的rois输入到build_fpn_mask_graph中。

            # Generate detection targets
            # Subsamples proposals and generates target outputs for training
            # Note that proposal class IDs, gt_boxes, and gt_masks are zero
            # padded. Equally, returned rois and targets are zero padded.
            rois, target_class_ids, target_bbox, target_mask =\
                DetectionTargetLayer(config, name="proposal_targets")([
                    target_rois, input_gt_class_ids, gt_boxes, input_gt_masks])

主要处理流程是

1. 确定ROI的类别
2. 修正bbox
3. 过滤背景和低置信度的boxes
4. 应用NMS
5. 保留top_k目标

training和inference

在ProposalLayer层之前，两者只有输入不同，之后两者的区别主要在于

• training模式：DetectionTargetLayer -> NetworkHeads(fpn_classifier_graph + build_fpn_mask_graph) -> Losses
• inference模式：NetworkHeads(fpn_classifier_graph ±> DetectionLayer -> build_fpn_mask_graph)

参考资料

在阅读源码的时候，参考这下面的资料，推荐先阅读资料在看源码。
令人拍案称奇的Mask RCNN
MASK RCNN 源码阅读(UPDATE)

Hellcatzm/Mask_RCNN: Mask R-CNN on Keras and TensorFlow，个人注释版