wt_ocr_and_docker/StyleText/arch/decoder.py

# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import paddle
import paddle.nn as nn

from arch.base_module import SNConv, SNConvTranspose, ResBlock


class Decoder(nn.Layer):
    def __init__(self, name, encode_dim, out_channels, use_bias, norm_layer,
                 act, act_attr, conv_block_dropout, conv_block_num,
                 conv_block_dilation, out_conv_act, out_conv_act_attr):
        super(Decoder, self).__init__()
        conv_blocks = []
        for i in range(conv_block_num):
            conv_blocks.append(
                ResBlock(
                    name="{}_conv_block_{}".format(name, i),
                    channels=encode_dim * 8,
                    norm_layer=norm_layer,
                    use_dropout=conv_block_dropout,
                    use_dilation=conv_block_dilation,
                    use_bias=use_bias))
        self.conv_blocks = nn.Sequential(*conv_blocks)
        self._up1 = SNConvTranspose(
            name=name + "_up1",
            in_channels=encode_dim * 8,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up2 = SNConvTranspose(
            name=name + "_up2",
            in_channels=encode_dim * 4,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up3 = SNConvTranspose(
            name=name + "_up3",
            in_channels=encode_dim * 2,
            out_channels=encode_dim,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._pad2d = paddle.nn.Pad2D([1, 1, 1, 1], mode="replicate")
        self._out_conv = SNConv(
            name=name + "_out_conv",
            in_channels=encode_dim,
            out_channels=out_channels,
            kernel_size=3,
            use_bias=use_bias,
            norm_layer=None,
            act=out_conv_act,
            act_attr=out_conv_act_attr)

    def forward(self, x):
        if isinstance(x, (list, tuple)):
            x = paddle.concat(x, axis=1)
        output_dict = dict()
        output_dict["conv_blocks"] = self.conv_blocks.forward(x)
        output_dict["up1"] = self._up1.forward(output_dict["conv_blocks"])
        output_dict["up2"] = self._up2.forward(output_dict["up1"])
        output_dict["up3"] = self._up3.forward(output_dict["up2"])
        output_dict["pad2d"] = self._pad2d.forward(output_dict["up3"])
        output_dict["out_conv"] = self._out_conv.forward(output_dict["pad2d"])
        return output_dict


class DecoderUnet(nn.Layer):
    def __init__(self, name, encode_dim, out_channels, use_bias, norm_layer,
                 act, act_attr, conv_block_dropout, conv_block_num,
                 conv_block_dilation, out_conv_act, out_conv_act_attr):
        super(DecoderUnet, self).__init__()
        conv_blocks = []
        for i in range(conv_block_num):
            conv_blocks.append(
                ResBlock(
                    name="{}_conv_block_{}".format(name, i),
                    channels=encode_dim * 8,
                    norm_layer=norm_layer,
                    use_dropout=conv_block_dropout,
                    use_dilation=conv_block_dilation,
                    use_bias=use_bias))
        self._conv_blocks = nn.Sequential(*conv_blocks)
        self._up1 = SNConvTranspose(
            name=name + "_up1",
            in_channels=encode_dim * 8,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up2 = SNConvTranspose(
            name=name + "_up2",
            in_channels=encode_dim * 8,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up3 = SNConvTranspose(
            name=name + "_up3",
            in_channels=encode_dim * 4,
            out_channels=encode_dim,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._pad2d = paddle.nn.Pad2D([1, 1, 1, 1], mode="replicate")
        self._out_conv = SNConv(
            name=name + "_out_conv",
            in_channels=encode_dim,
            out_channels=out_channels,
            kernel_size=3,
            use_bias=use_bias,
            norm_layer=None,
            act=out_conv_act,
            act_attr=out_conv_act_attr)

    def forward(self, x, y, feature2, feature1):
        output_dict = dict()
        output_dict["conv_blocks"] = self._conv_blocks(
            paddle.concat(
                (x, y), axis=1))
        output_dict["up1"] = self._up1.forward(output_dict["conv_blocks"])
        output_dict["up2"] = self._up2.forward(
            paddle.concat(
                (output_dict["up1"], feature2), axis=1))
        output_dict["up3"] = self._up3.forward(
            paddle.concat(
                (output_dict["up2"], feature1), axis=1))
        output_dict["pad2d"] = self._pad2d.forward(output_dict["up3"])
        output_dict["out_conv"] = self._out_conv.forward(output_dict["pad2d"])
        return output_dict


class SingleDecoder(nn.Layer):
    def __init__(self, name, encode_dim, out_channels, use_bias, norm_layer,
                 act, act_attr, conv_block_dropout, conv_block_num,
                 conv_block_dilation, out_conv_act, out_conv_act_attr):
        super(SingleDecoder, self).__init__()
        conv_blocks = []
        for i in range(conv_block_num):
            conv_blocks.append(
                ResBlock(
                    name="{}_conv_block_{}".format(name, i),
                    channels=encode_dim * 4,
                    norm_layer=norm_layer,
                    use_dropout=conv_block_dropout,
                    use_dilation=conv_block_dilation,
                    use_bias=use_bias))
        self._conv_blocks = nn.Sequential(*conv_blocks)
        self._up1 = SNConvTranspose(
            name=name + "_up1",
            in_channels=encode_dim * 4,
            out_channels=encode_dim * 4,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up2 = SNConvTranspose(
            name=name + "_up2",
            in_channels=encode_dim * 8,
            out_channels=encode_dim * 2,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._up3 = SNConvTranspose(
            name=name + "_up3",
            in_channels=encode_dim * 4,
            out_channels=encode_dim,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1,
            use_bias=use_bias,
            norm_layer=norm_layer,
            act=act,
            act_attr=act_attr)
        self._pad2d = paddle.nn.Pad2D([1, 1, 1, 1], mode="replicate")
        self._out_conv = SNConv(
            name=name + "_out_conv",
            in_channels=encode_dim,
            out_channels=out_channels,
            kernel_size=3,
            use_bias=use_bias,
            norm_layer=None,
            act=out_conv_act,
            act_attr=out_conv_act_attr)

    def forward(self, x, feature2, feature1):
        output_dict = dict()
        output_dict["conv_blocks"] = self._conv_blocks.forward(x)
        output_dict["up1"] = self._up1.forward(output_dict["conv_blocks"])
        output_dict["up2"] = self._up2.forward(
            paddle.concat(
                (output_dict["up1"], feature2), axis=1))
        output_dict["up3"] = self._up3.forward(
            paddle.concat(
                (output_dict["up2"], feature1), axis=1))
        output_dict["pad2d"] = self._pad2d.forward(output_dict["up3"])
        output_dict["out_conv"] = self._out_conv.forward(output_dict["pad2d"])
        return output_dict