目标检测与分割之MaskRCNN代码结构流程全面梳理+总结_mask r-cnn分割分支在那个文件

学习目标检测不久，花一个多月的空闲时间啃完了maskRCNN论文和tensorflow+keras开源代码，感觉进入了一个新的境界，在此做一个比较详细的流程梳理和总结，同时也希望对其他学习目标检测和语义分割的同学有点帮助。
本文主要针对代码结构进行梳理，对于原理还不熟悉的同学可以参考这几篇博文：
格灵深瞳DeepGlint：干货 | 目标检测入门，看这篇就够了（已更完）
白裳：一文读懂Faster RCNN
stone：令人拍案称奇的Mask RCNN
另外由于博主比较懒，借用了以下系列博文的部分框图，侵删：
BINGO Hong：MASK_RCNN代码详解(1)-Basebone部分

开始前必须膜拜一下华人学者的骄傲——何恺明（K.He）大神。MaskRCNN虽然模型比较复杂，不太适合应用于实时性要求较高的场景，但其检测精度从2017至今仍未被超越，更值得一提的是，模型的BackBone(Resnet101)以及前身FasterRCNN都出自他手，更要命的是，比我大不了几岁%￥#&^*
现在进入正题。本文共分三部分，分别是：
第一部分 数据输入
第二部分 模型介绍
第三部分 Train模式和Inference模式流程总结

第一部分 数据输入
data_generator无限循环读取数据，每次读取一个batch，包括三个真值：input_gt_class_ids, input_gt_boxes, input_gt_masks，两个输入数据：input_image, input_image_meta , 两个rpn标签：input_rpn_match, input_rpn_bbox，每batch_size次打包输入模型。

一、generate_pyramid_anchors 在BackBone每一层特征图上生成anchors。
对于2-6每层特征图的每个像素点，分别生成大小为RPN_ANCHOR_SCALES (32, 64, 128, 256, 512)，尺寸比例为RPN_ANCHOR_RATIOS [0.5, 1, 2]三种大小的anchors，总数为3*（256²⁺¹²⁸2+64²⁺³²2+16^2）=261888个。

二、load_image_gt生成所需格式的数据和真值
（包括image，image_meta, gt_class_ids, gt_boxes, gt_masks）

1. load_image读取图片，得到image
2. load_mask 根据标注文件，生成对应于每个标注目标的二值化的gt_masks和class_ids
3. utils.resize_image根据IMAGE_MIN_DIM(800)，IMAGE_MIN_SCALE(0)， IMAGE_MAX_DIM(1024)， IMAGE_RESIZE_MODE("square“)对图片进行缩放，生成window, scale, padding, crop等
4. 根据scale, padding, crop调整mask大小
5. 从mask中找出左下、右上坐标，作为gt_boxes
6. 找到感兴趣的类别active_class_ids
7. 如果使用mini_mask，则只保留bbox内部的数据
8. 得到所有image_meta数据：

meta = np.array(
        [image_id] +                  # size=1
        list(original_image_shape) +     # size=3
        list(image_shape) +            # size=3
        list(window) +               # size=4 (y1, x1, y2, x2) 
        [scale] +                     # size=1
        list(active_class_ids)           # size=num_classes)
1
2
3
4
5
6
7

1. 真值数量控制在最多MAX_GT_INSTANCES(100个)

三、build_rpn_targets 生成rpn层的训练标签，与head网络分开训练
输入：image.shape, anchors, gt_class_ids, gt_boxes, config
输出：rpn_match, rpn_bbox

1. 按gt_class_id < 0 筛选出包含多个对象的crowd_boxes，去掉这些boxes得到新的gt_boxes，算出anchors和gt_boxes的IOU矩阵
2. .找到每个anchors与所有gt_box的最大IOU，最大IOU< 0.3且与crowd_boxes最大IOU<0.001的为负样本。
3. 找到与每个gt_box的IOU最大的anchors，设为正样本，与每个gt_box的最大IOU>0.7的也设为正样本。保证每个gt_box都有对应的anchors
4. 平衡正负样本数量，保证正负样本1:1，总数为RPN_TRAIN_ANCHORS_PER_IMAGE（256个），正样本不够用负样本补齐。
5. 对于每一个正样本，将IOU最大的gt_box作为对应的真值，算出两个box的变换尺度因子。

  rpn_bbox[ix] = [
            (gt_center_y - a_center_y) / a_h,
            (gt_center_x - a_center_x) / a_w,
            np.log(gt_h / a_h),
            np.log(gt_w / a_w),
1
2
3
4
5

6. 返回样本标签rpn_match（261888 维，分为-1,0,1）和正样本尺度变换因子rpn_bbox（256维）

第二部分 模型部分

一、BackBone网络
作为底层的特征提取网络
1.深度残差网络ResNet101：
保证在堆叠网络的过程中，网络至少不会因为继续堆叠而产生退化


2. 特征金字塔FPN：
同时利用低层特征图的空间信息和高层特征图的语义信息

二、RPN网络
build_rpn_model生成RPN网络，得到所有anchors的分类、回归信息
输入FPN每层特征图，分别计算，再将结果拼接在一起
以P2特征图为例：(batch256256*256)

1. 先经过 (512个33256卷积核)生成shared层
2. . shared层分为两个分支，一个分类一个回归。
3. 分类分支经过6个11512卷积核，变为batch2562566，reshape得到rpn_class_logits（batch2562563*2）然后经过softmax得到rpn_probs.
4. 回归分支经过12个11512卷积核，变为batch25625612, reshape得到rpn_bbox（batch2562563*4）
5. 得到一个特征图的三个结果后，把所有特征图合并，然后用outputs = list(zip(*layer_outputs))将三个结果放在外层，特征图放在内层，最后得到rpn_class_logits, rpn_class, rpn_bbox（也可以直接输入ROIs，不用RPN）
   

三、ProposalLayer层
根据RPN网络输出的分类、回归信息，经过NMS得到最终的ROIs
输入： [rpn_class, rpn_bbox, anchors]
输出：rpn_rois

1. 先按照scores排序，找到得分最高的PRE_NMS_LIMIT （6000个）anchors。
2. 根据delta得到box的坐标
3. 把box坐标crop到0-1范围内
4. 执行NMS，阈值为RPN_NMS_THRESHOLD（0.7），train模式保留POST_NMS_ROIS_TRAINING（2000个），inference模式保留POST_NMS_ROIS_INFERENCE（1000个），不够的padding 0
5. 得到proposals,也就是rois。

四、DetectionTargetLayer层（train模式特有）
生成head网络的训练标签
输入：target_rois, input_gt_class_ids, gt_boxes, input_gt_masks
输出：rois, target_class_ids, target_bbox, target_mask

1. 去掉输入中的全0数据（padding得到）
2. 按gt_class_id < 0 筛选出包含多个对象的crowd_boxes，去掉这些boxes得到新的gt_boxes，算出proposals和gt_boxes的IOU矩阵(这步与rpn层相同)
3. 找到每个proposal与所有gt_boxes的最大的IOU，如果>0.5,设为正样本，小于0.5且与crowd_boxes 的最大IOU小于0.001的设为负样本。（与rpn层不同）
4. 随机挑选正样本，数量=TRAIN_ROIS_PER_IMAGE(200)*config.ROI_POSITIVE_RATIO(0.33), 随机挑选负样本，数量 = 正样本数量/ratio – 正样本数量。
5. 找到每个proposal对应的最大IOU的索引，对应关系存在roi_gt_box_assignment中，即每个proposal对应哪个gt_box
6. 找到这些gt_boxes,算出delta
7. 找到对应的gt_masks，用tf.image.crop_and_resize根据proposal的大小抠出来，得到masks，注意不是直接用gt_masks
8. 将正负样本合在一起，得到rois, 若总数不够TRAIN_ROIS_PER_IMAGE，padding 0补齐。

五、Head网络
对筛选出来的ROIs进行分类、回归、mask分割操作

1.fpn_classifier_graph分类和回归分支
首先经过PyramidROIAlign，采用双线性插值将每个ROI精确地转换为相同的大小，再输入分类回归网络，得到每个ROI的分类和bbox回归结果
输入：rois, mrcnn_feature_maps, input_image_meta
输出：mrcnn_class_logits, mrcnn_class, mrcnn_bbox
其中train模式输入的是DetectionTargetLayer层筛选的用于训练的roi, 数量是TRAIN_ROIS_PER_IMAGE（200个）
Inference模式输入的是ProposalLayer层输出的全部的用来预测的ROI，数量是POST_NMS_ROIS_INFERENCE（1000个）

1）PyramidROIAlign层
输入：[rois, image_meta] + feature_maps
输出：pooled
a.根据roi的w和h与图片大小的比例算出应该对应特征图哪个层：roi_level = log2_graph(tf.sqrt(h * w) / (224.0 / tf.sqrt(image_area))), 并把level缩放到2-5之间
b.找到每一个level的box：level_boxes和在所有box中的索引：box_to_level（第几个batch,第几个），和box在batch中的索引box_indices
c.将level_boxes和box_indices做梯度阻断（和RPN不通）
d.在这个level的特征图中根据level_boxes, box_indices抠图，并采用双线性插值

 pooled.append(tf.image.crop_and_resize(
                feature_maps[i], level_boxes, box_indices, self.pool_shape,
                method="bilinear")) 
1
2
3

e.循环每个level,将抠的图合在一起，得到pooled
f.把pooled按batch排序(方法是先把box_to_level按batch排序，得到序号，把pooled重排)，然后转成[batch, num, 7, 7, 256]大小，7为pool_size

2）FPN_Classifier网络
对于ROIAlign层输出的[batch, num, 7, 7, 256]大小的数据，用KL.TimeDistributed实现在第一维（num）将每个batch的每个ROI拆分,，并行执行各项操作
a. mrcnn_class_conv1: 1024 * 7 * 7，归一化，relu
b. mrcnn_class_conv2: 1024 * 1 * 1，归一化，relu
c. 去掉多余的维度，从(batch, 200, 1, 1, 1024)降到(batch, 200, 1024)
d. 分两个分支，一个经过num_classes维全连接，softmax，得到mrcnn_class_logits和mrcnn_class
e. 一个经过num_classes * 4维全连接，reshape，得到mrcnn_bbox_fc和mrcnn_bbox

2.fpn_mask_graph mask分割分支
每个ROI同样先经过PyramidROIAlign，然后输入MASK分割网络，得到mask结果
输入：rois, mrcnn_feature_maps, input_image_meta,
输出：mrcnn_mask
其中train模式输入的是DetectionTargetLayer层筛选的用于训练的roi, 数量是TRAIN_ROIS_PER_IMAGE（200个）
Inference模式输入的是DetectionLayer层输出的NMS后的检测结果，数量是实际数量，最多DETECTION_MAX_INSTANCES（100个）
1) PyramidROIAlign层
和分类和回归分支不同的只有pool_size,此处为14
2) FPN_mask网络
a. mrcnn_mask_conv1,2,3,4: 256 * 3 * 3，归一化，relu
b. mrcnn_mask_deconv: 256 * 2 * 2, 步长2，relu, 2倍上采样
c. num_classes个1*1卷积得到mrcnn_mask

六、LOSS部分
RPN损失
1. rpn_class_loss
输入：input_rpn_match（真值）, rpn_class_logits（检测值）
将input_rpn_match中-1,1的标签变成0,1，并去掉原来标签为0的（非正非负），和相应的检测值rpn_class_logits求平均交叉熵：

loss = K.sparse_categorical_crossentropy(target=anchor_class,output=rpn_class_logits，from_logits=True)
1

2. rpn_bbox_loss
输入：input_rpn_bbox （真值）, input_rpn_match, rpn_bbox（检测值）
只找input_rpn_match为1的anchors，并且需要将真值的每个batch裁出和input_rpn_match为1的anchors一样的数量（取前count个），然后算平均smooth_l1_loss:

loss = smooth_l1_loss(target_bbox, rpn_bbox)
1

MRCNN损失
3. mrcnn_class_loss
输入：target_class_ids, mrcnn_class_logits, active_class_ids
求mrcnn_class_logits和target_class_ids的稀疏交叉熵，并找到检测值对应logit最大的类，去除不需要的类。

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target_class_ids, logits=pred_class_logits)
1

4. mrcnn_bbox_loss
输入：target_bbox, target_class_ids, mrcnn_bbox
target_bbox（每个roi只有一个）找到对应target_class_ids大于0的，mrcnn_bbox（每个class_id预测了一个）找到target_class_ids>0且与真值class_id一致的那个bbox，算平均smooth_l1_loss:

loss = K.switch(tf.size(target_bbox) > 0, smooth_l1_loss(y_true=target_bbox, y_pred=pred_bbox), tf.constant(0.0))
1

5. mrcnn_mask_loss
输入：target_mask, target_class_ids, mrcnn_mask
真值为0,1，检测值为0-1，与bbox类似，选择对应真值target_class_ids>0 且class_id与真值一致的mask计算二值化交叉熵：

loss = K.switch(tf.size(y_true) > 0, K.binary_crossentropy(target=y_true, output=y_pred), tf.constant(0.0))
1

七、DetectionLayer层（inference模式特有）
将分类和回归结果应用于ROI，并执行NMS和置信度过滤，得到最终的检测结果
输入：rpn_rois, mrcnn_class, mrcnn_bbox, input_image_meta
输出：detections(包含refined_rois, class_ids, class_scores)
a. 得到每个ROI预测概率最大的类别class_ids
b. 找到每个ROI预测的class_id对应的class_scores, deltas
c. 将delta应用于ROI，得到refined_rois
d. 按照image_meta中window的大小裁剪到window内部。
e. 筛选出class_id大于0且scores > DETECTION_MIN_CONFIDENCE(0.9)的ROI
f. 对于每个class_id，做NMS,阈值为DETECTION_NMS_THRESHOLD(0.3)，保留个数是DETECTION_MAX_INSTANCES(100)，不够的补-1
g. 对NMS后的ROI去掉值为-1的，找到对应的class_scores，然后排序，取class_scores最大的DETECTION_MAX_INSTANCES个refined_rois，不够的padding 0，与class_ids和class_scores合并为detections输出。

第三部分 Train模式和Inference模式流程总结
一、Train模式流程总结
.通过数据生成器不断读取数据和真值，转换为需要的格式

1. 将输入数据对应于特征金字塔每一层的规模，生成anchors(261888个)，并根据与真值的IOU筛选正负样本(二分类，比例1:1，共256个，正样本不够用负样本补)，生成RPN网络训练标签(包括正负分类标签、bbox回归deltas)
2. 初始化模型，模式为’train’
3. 输入数据经过BackBone和RPN层后，得到每个anchors的分类、回归信息，在此根据上一步筛选的样本和计算的RPN训练标签，计算RPN分类、回归损失，训练RPN网络
4. RPN网络输出的anchors分类、回归信息输入ProposalLayer，先根据得分排序，取前PRE_NMS_LIMIT(6000)个，再经过NMS变为POST_NMS_ROIS_TRAINING (2000)个proposals
5. Proposals经过DetectionTargetLayer，根据与真值的IOU筛选正负样本（只考虑正负，比例为1:2，总数为TRAIN_ROIS_PER_IMAGE(200)个，正样本不够用pad0补)，生成head网络训练标签(包括多分类标签、bbox回归deltas、masks)
6. 训练用的200个ROI分别通过FPN分类/回归分支和FPN mask两个分支，每个分支都首先进行PyramidROIAlign，将每个ROI用公式对应到最优特征图进行抠图， 得到[batch, num, 7, 7, 256]大小，然后分别计算class_logits, bbox, mask
7. 根据步骤5的标签和步骤6的检测值，计算MRCNN分类、回归、mask损失，并训练head网络，注意由于bbox和mask每个类别都预测了一个，只计算与标签类别一致的类别的损失。
8. 两次训练的区别是RPN网络是全图一次卷积，head网络是200个ROI并行计算。

二、Inference模式流程总结

1. 建立模型，加载权重，模式为’inference’
2. 读取数据，转换为需要的格式
3. 按照与训练网络同样的方式生成anchors(261888个)
4. 将输入数据和anchors输入模型进行预测，得到detections和mrcnn_mask：
   A. 数据输入模型后经过BackBone, RPN网络, ProposalLayer层，得到ROIS，数量为POST_NMS_ROIS_INFERENCE（1000）个。
   B. ROIS先输入fpn_classifier_graph，得到分类、回归信息
   C. 然后经过DetectionLayer，将分类和回归结果应用于ROI，并执行NMS和置信度过滤，得到最终的boxes以及分类和回归信息，数量为最多DETECTION_MAX_INSTANCES（100）个，多了取分数最高的100个，不够padding 0
   D. 将boxes输入fpn_mask_graph，得到mask
5. 将检测结果进行后处理，包括把mask进行二值化处理、mask和boxes按照图片原始信息进行缩放、筛选等，最后输出最终结果。

关于MaskRCNN的介绍告一段落，大家如果有疑问可以留言，尽量回答，共同进步。同时由于作者比较跨界，所以其他任何问题也欢迎留言~~~