PaddleDetection-MaskRcnn总体结构解析

2021SC@SDUSC

Mask Rcnn整体结构

对应PaddleDetection的流程图

modeling/architecture/mask_rcnn.py源码解析

在yaml的配置：/configs/_base_/models/mask_rcnn_r50_fpn.yml

'''
# Model Achitecture
# MaskRCNNYOLOV3结构，包括了以下主要组
MaskRCNN:
  # model anchor info flow
  anchor: Anchor
  proposal: Proposal
  mask: Mask
  # model feat info flow
  backbone: ResNet #主干网络类名
  neck: FPN #neck FPN类名
  rpn_head: RPNHead #Head成员类名
  bbox_head: BBoxHead #Head成员类名
  mask_head: MaskHead #Head成员类名
  # post process
  bbox_post_process: BBoxPostProcess #BBox后处理
  mask_post_process: MaskPostProcess #Mask后处理
'''

相关库引用与定义 

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import paddle
from ppdet.core.workspace import register
from .meta_arch import BaseArch

__all__ = ['MaskRCNN']

接下来是具体代码分析

@register
class MaskRCNN(BaseArch):
    __category__ = 'architecture'
    __inject__ = [
        'anchor',              #anchor成员类
        'proposal',            #proposal成员类
        'mask',                #mask成员类
        'backbone',            #backbone成员类
        'neck',                #neck成员类
        'rpn_head',            #rpn_head成员类
        'bbox_head',           #bbox_head成员类
        'mask_head',           #mask_head成员类
        'bbox_post_process',   #bbox_post_process成员类
        'mask_post_process',   #mask_post_process成员类
    ]

    def __init__(self,
                 anchor,
                 proposal,
                 mask,
                 backbone,
                 rpn_head,
                 bbox_head,
                 mask_head,
                 bbox_post_process,
                 mask_post_process,
                 neck=None):
        super(MaskRCNN, self).__init__()
        self.anchor = anchor
        self.proposal = proposal
        self.mask = mask
        self.backbone = backbone
        self.neck = neck
        self.rpn_head = rpn_head
        self.bbox_head = bbox_head
        self.mask_head = mask_head
        self.bbox_post_process = bbox_post_process
        self.mask_post_process = mask_post_process
}

 MaskRCNN类的定义，以及初始化函数。相关注释在代码中。

def model_arch(self):
        # Backbone 主干网络
        body_feats = self.backbone(self.inputs)
        spatial_scale = 1. / 16

        # Neck 颈
        if self.neck is not None:
            body_feats, spatial_scale = self.neck(body_feats)

        # RPN  区域选取网络
        # rpn_head returns two list: rpn_feat, rpn_head_out 
        # each element in rpn_feats contains rpn feature on each level,
        # and the length is 1 when the neck is not applied.
        # each element in rpn_head_out contains (rpn_rois_score, rpn_rois_delta)
        rpn_feat, self.rpn_head_out = self.rpn_head(self.inputs, body_feats)

        # Anchor 锚框
        # anchor_out returns a list,
        # each element contains (anchor, anchor_var)
        self.anchor_out = self.anchor(rpn_feat)

        # Proposal RoI  提供感兴趣区域
        # compute targets here when training
        rois = self.proposal(self.inputs, self.rpn_head_out, self.anchor_out)
        # BBox Head BBox头
        bbox_feat, self.bbox_head_out, self.bbox_head_feat_func = self.bbox_head(
            body_feats, rois, spatial_scale)

        rois_has_mask_int32 = None
        if self.inputs['mode'] == 'infer':#推理过程中BBox特征又bbox头获取
            bbox_pred, bboxes = self.bbox_head.get_prediction(
                self.bbox_head_out, rois)
            # Refine bbox by the output from bbox_head at test stage
            self.bboxes = self.bbox_post_process(bbox_pred, bboxes,
                                                 self.inputs['im_shape'],
                                                 self.inputs['scale_factor'])
        else:
            # Proposal RoI for Mask branch 感兴趣区域从掩膜特征获取
            # bboxes update at training stage only
            bbox_targets = self.proposal.get_targets()[0]
            self.bboxes, rois_has_mask_int32 = self.mask(self.inputs, rois,
                                                         bbox_targets)

        # Mask Head  掩膜头
        self.mask_head_out = self.mask_head(
            self.inputs, body_feats, self.bboxes, bbox_feat,
            rois_has_mask_int32, spatial_scale, self.bbox_head_feat_func)

def get_loss(self, ):
        loss = {}

        # RPN loss 计算候选区域损失
        rpn_loss_inputs = self.anchor.generate_loss_inputs(
            self.inputs, self.rpn_head_out, self.anchor_out)
        loss_rpn = self.rpn_head.get_loss(rpn_loss_inputs)
        loss.update(loss_rpn)

        # BBox loss 计算bbox损失
        bbox_targets = self.proposal.get_targets()
        loss_bbox = self.bbox_head.get_loss([self.bbox_head_out], bbox_targets)
        loss.update(loss_bbox)

        # Mask loss 计算掩膜损失
        mask_targets = self.mask.get_targets()
        loss_mask = self.mask_head.get_loss(self.mask_head_out, mask_targets)
        loss.update(loss_mask)
        #损失求和，更新loss
        total_loss = paddle.add_n(list(loss.values()))
        loss.update({'loss': total_loss})
        return loss

def get_pred(self):#预测推理结果（框、数量、掩膜）
        bbox, bbox_num = self.bboxes
        output = {
            'bbox': bbox,
            'bbox_num': bbox_num,
            'mask': self.mask_head_out
        }
        return output