Detail-recovery Image Deraining via Context Aggregation Networks 网络结构分析

《Detail-recovery Image Deraining via Context Aggregation Networks》这篇文章选自CVPR2020，针对图像去雨过程中的细节损失恢复问题提出了可拆卸细节补充模块，接下来重点分析一下网络结构和实现代码。

这个一个两分支并行结构，两条并行分支分别是提取雨水条纹的RRN和提取图像细节的DRN。
输入的雨图传入到RRN得到输出雨条纹RAIN STREAKS，同时将雨图传入到DRN得到输出的细节补充Detail Repair Feature，再把雨图减去雨条纹，加上细节补充就得到了去雨修复图像。
有必要说明的一点是，两个分支模块是独立的，可拆卸的。

接下来我们分析一下RRN和DRN模块的结构。
DRN：
DRN模块和RRN模块相似，是由一层3×3Conv、PRelu，16层RRB，一层3×3Conv、BN，一层残差结构，一层3×3Conv构成。
其中，RRB结构图如下：

RRB由一层Conv、BN、PRelu，一层Conv、BN，一层SE模块，一层残差结构构成。
在SE模块中，输入图像经过一层全局平均池化，两层全连接层（参数分别是Relu和Sigmoid）得到channel权重，再和输入图像相乘得到最后的输出。

RRN：
RRN模块由一层3×3Conv、PRelu，16层SDCAB，一层3×3Conv、BN，一层残差结构，一层3×3Conv构成。
其中，SDCAB结构图如下：

SDCAB由一层DCCL，一层BN、PRelu，一层DCCL，一层BN，一层残差结构构成。
DCCL是由三个拥有不同大小感受野的并行卷积模块组成，其感受野的大小分别是1×1,3×3,5×5，输出cat在一起，再经过一层1×1Conv就得到最终的输出。

上述网络实现代码如下：

ef get_the_end_model(input_channel_num=3, feature_dim=64, resunit_num=16):

    def _back_net(inputs):
        def _residual_block(inputs):
            x = _empty_block(inputs)
            x = BatchNormalization()(x)
            x = PReLU(shared_axes=[1, 2])(x)
            x = _empty_block(x)
            x = BatchNormalization()(x)
            m = Add()([x, inputs])
            return m

        def _empty_block(inputs):
            x1 = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs)
            x2 = Conv2D(feature_dim, (3, 3), dilation_rate=3, padding="same", kernel_initializer="he_normal")(inputs)
            x3 = Conv2D(feature_dim, (3, 3), dilation_rate=5, padding="same", kernel_initializer="he_normal")(inputs)
            x = concatenate([x1, x2, x3], axis=-1)
            x_out = Conv2D(feature_dim, (1, 1), padding="same", kernel_initializer="he_normal")(x)
            return x_out

        x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs)  # 输出64层
        x = PReLU(shared_axes=[1, 2])(x)
        x0 = x

        for i in range(resunit_num):
            x = _residual_block(x)  # 16层SDCAB

        x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x)
        x = BatchNormalization()(x)
        x = Add()([x, x0])  # 全局残差
        x = Conv2D(input_channel_num, (3, 3), padding="same", kernel_initializer="he_normal")(x)  # 输出3层
        return x

    def _rain_net(inputs):
        def _residual_block(inputs, number):
            x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs)  # 3*3卷积，64层输出
            x = BatchNormalization()(x)
            x = PReLU(shared_axes=[1, 2])(x)
            x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x)  # 3*3卷积，64层输出
            x = BatchNormalization()(x)
            filters = 64  # SE模块
            se_shape = (1, 1, filters)
            se = GlobalAveragePooling2D()(x)
            se = Reshape(se_shape)(se)  # 形状变换，se-->se_shape
            se = Dense(number, activation="relu",
                       kernel_initializer="he_normal",use_bias=False)(se)  # FC1+Relu,(转换为稠密张量)
            se = Dense(filters, activation="hard_sigmoid",
                       kernel_initializer="he_normal", use_bias=False)(se)  # FC2+Sigmod
            x = multiply([x, se])  # 权重值相乘
            m = Add()([x, inputs])  # 全局残差
            return m

        x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(inputs)  # 3*3卷积，64层输出
        x = PReLU(shared_axes=[1, 2])(x)
        x0 = x

        for i in range(resunit_num):  # 16层RRB
            x = _residual_block(x, 4)

        x = Conv2D(feature_dim, (3, 3), padding="same", kernel_initializer="he_normal")(x)
        x = BatchNormalization()(x)
        x = Add()([x, x0])  # 全局残差
        x = Conv2D(input_channel_num, (3, 3), padding="same", kernel_initializer="he_normal")(x)  # 3*3卷积，3层输出
        return x

    inputs = Input(shape=(None, None, input_channel_num), name='Rain_image')  # 0*0*3
    Rain = _rain_net(inputs)  # Rain Residual Network

    out1 = Subtract()([inputs, Rain])  # 去雨后的图像
    Back_in = Add()([out1, inputs])
    Back = _back_net(Back_in)  # 输入的是原始的雨图和去雨后的图像的叠加，猜测这样做得到的细节更好
    out = Add()([out1, Back])  # 加上提取到的细节

    model = Model(inputs=inputs, outputs=[out1, out])  # 定义模型的输入输出
    return model