# 先从 resnet 抽取四个不同阶段的特征图 C2-C5。

_, C2, C3, C4, C5 =

resnet_graph(input_image, config.BACKBONE,stage5=True, train_bn=config.TRAIN_BN)

# Top-down Layers 构建自上而下的网络结构

# 从 C5开始处理，先卷积来转换特征图尺寸

P5 = KL.Conv2D(256, (1, 1), name='fpn_c5p5')(C5)

# 上采样之后的P5和卷积之后的 C4像素相加得到 P4，后续的过程就类似了

P4 = KL.Add(name="fpn_p4add")([

            KL.UpSampling2D(size=(2, 2), name="fpn_p5upsampled")(P5),

            KL.Conv2D(256, (1, 1),name='fpn_c4p4')(C4)])

P3 = KL.Add(name="fpn_p3add")([

            KL.UpSampling2D(size=(2, 2), name="fpn_p4upsampled")(P4),

            KL.Conv2D(256, (1, 1), name='fpn_c3p3')(C3)])

P2 = KL.Add(name="fpn_p2add")([

            KL.UpSampling2D(size=(2, 2),name="fpn_p3upsampled")(P3),

            KL.Conv2D(256, (1, 1), name='fpn_c2p2')(C2)])

# P2-P5最后又做了一次3*3的卷积，作用是消除上采样带来的混叠效应

# Attach 3x3 conv to all P layers to get the final feature maps.

P2 = KL.Conv2D(256, (3, 3), padding="SAME", name="fpn_p2")(P2)

P3 = KL.Conv2D(256, (3, 3), padding="SAME",name="fpn_p3")(P3)

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

P5 = KL.Conv2D(256, (3, 3), padding="SAME",name="fpn_p5")(P5)

# P6 is used for the 5th anchor scale in RPN. Generated by

# subsampling from P5 with stride of 2.

P6 = KL.MaxPooling2D(pool_size=(1, 1), strides=2,name="fpn_p6")(P5)

# 注意 P6是用在 RPN 目标区域提取网络里面的，而不是用在 FPN 网络

# Note that P6 is used in RPN, but not in the classifier heads.

rpn_feature_maps = [P2, P3, P4, P5, P6]

# 最后得到了5个融合了不同层级特征的特征图列表；

# Assign each ROI to a level in the pyramid based on the ROI area.

# 这里的 boxes 是 ROI 的框，用来计算得到每个 ROI 框的面积

y1, x1, y2, x2 = tf.split(boxes, 4, axis=2)

h = y2 - y1

w = x2 - x1

# Use shape of first image. Images in a batch must have the same size.

# 这里得到原图的尺寸，计算原图的面积

image_shape = parse_image_meta_graph(image_meta)['image_shape'][0]

# Equation 1 in the Feature Pyramid Networks paper. Account for

# the fact that our coordinates are normalized here.

# e.g. a 224x224 ROI (in pixels) maps to P4

# 原图面积

image_area = tf.cast(image_shape[0] * image_shape[1], tf.float32)

# 分两步计算每个 ROI 框需要在哪个层的特征图中进行 pooling

roi_level = log2_graph(tf.sqrt(h * w) / (224.0 / tf.sqrt(image_area)))

roi_level = tf.minimum(5, tf.maximum(

    2, 4 + tf.cast(tf.round(roi_level), tf.int32)))

# 代码里面的计算替换为以下计算方式：

roi_level = min(5, max(2, 4 + log2(sqrt(w * h) / ( 224 / sqrt(image_area)) ) ) )

for i, level in enumerate(range(2, 6)):

            # 先找出需要在第 level 层计算ROI

            ix = tf.where(tf.equal(roi_level, level))

            level_boxes = tf.gather_nd(boxes, ix)

            # Box indicies for crop_and_resize.

            box_indices = tf.cast(ix[:, 0], tf.int32)

            # Keep track of which box is mapped to which level

            box_to_level.append(ix)

            # Stop gradient propogation to ROI proposals

            level_boxes = tf.stop_gradient(level_boxes)

            box_indices = tf.stop_gradient(box_indices)

            # 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]

            # 使用 tf.image.crop_and_resize 进行 ROI pooling

            pooled.append(tf.image.crop_and_resize(

                feature_maps[i], level_boxes, box_indices, self.pool_shape,

                method="bilinear"))

# ROI Pooling

# Shape: [batch, num_boxes, pool_height, pool_width, channels]

# 得到经过 ROI pooling 之后的特征列表

x = PyramidROIAlign([pool_size, pool_size],

                    name="roi_align_classifier")([rois, image_meta] + feature_maps)

# 将上面得到的特征列表送入 2 个1024通道数的卷积层以及 2 个 rulu 激活层

# Two 1024 FC layers (implemented with Conv2D for consistency)

x = KL.TimeDistributed(KL.Conv2D(1024, (pool_size, pool_size), padding="valid"),

                       name="mrcnn_class_conv1")(x)

x = KL.TimeDistributed(BatchNorm(), name='mrcnn_class_bn1')(x, training=train_bn)

x = KL.Activation('relu')(x)

x = KL.TimeDistributed(KL.Conv2D(1024, (1, 1)),

                       name="mrcnn_class_conv2")(x)

x = KL.TimeDistributed(BatchNorm(), name='mrcnn_class_bn2')(x, training=train_bn)

x = KL.Activation('relu')(x)

shared = KL.Lambda(lambda x: K.squeeze(K.squeeze(x, 3), 2),

                   name="pool_squeeze")(x)

# 分类层

# Classifier head

mrcnn_class_logits = KL.TimeDistributed(KL.Dense(num_classes),

                                        name='mrcnn_class_logits')(shared)

mrcnn_probs = KL.TimeDistributed(KL.Activation("softmax"),

                                 name="mrcnn_class")(mrcnn_class_logits)

# BBOX 的位置偏移回归层

# BBox head

# [batch, boxes, num_classes * (dy, dx, log(dh), log(dw))]

x = KL.TimeDistributed(KL.Dense(num_classes * 4, activation='linear'),

                       name='mrcnn_bbox_fc')(shared)

# Reshape to [batch, boxes, num_classes, (dy, dx, log(dh), log(dw))]

s = K.int_shape(x)

mrcnn_bbox = KL.Reshape((s[1], num_classes, 4), name="mrcnn_bbox")(x)