first commit

ppocr/modeling/transforms/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+__all__ = ["build_transform"]
+
+
+def build_transform(config):
+    from .tps import TPS
+    from .stn import STN_ON
+    from .tsrn import TSRN
+    from .tbsrn import TBSRN
+    from .gaspin_transformer import GA_SPIN_Transformer as GA_SPIN
+
+    support_dict = ["TPS", "STN_ON", "GA_SPIN", "TSRN", "TBSRN"]
+
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "transform only support {}".format(support_dict)
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/modeling/transforms/gaspin_transformer.py

@@ -0,0 +1,319 @@
+# copyright (c) 2022 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import numpy as np
+import functools
+from .tps import GridGenerator
+
+"""This code is refer from:
+https://github.com/hikopensource/DAVAR-Lab-OCR/davarocr/davar_rcg/models/transformations/gaspin_transformation.py
+"""
+
+
+class SP_TransformerNetwork(nn.Layer):
+    """
+    Sturture-Preserving Transformation (SPT) as Equa. (2) in Ref. [1]
+    Ref: [1] SPIN: Structure-Preserving Inner Offset Network for Scene Text Recognition. AAAI-2021.
+    """
+
+    def __init__(self, nc=1, default_type=5):
+        """Based on SPIN
+        Args:
+            nc (int): number of input channels (usually in 1 or 3)
+            default_type (int): the complexity of transformation intensities (by default set to 6 as the paper)
+        """
+        super(SP_TransformerNetwork, self).__init__()
+        self.power_list = self.cal_K(default_type)
+        self.sigmoid = nn.Sigmoid()
+        self.bn = nn.InstanceNorm2D(nc)
+
+    def cal_K(self, k=5):
+        """
+
+        Args:
+            k (int): the complexity of transformation intensities (by default set to 6 as the paper)
+
+        Returns:
+            List: the normalized intensity of each pixel in [0,1], denoted as \beta [1x(2K+1)]
+
+        """
+        from math import log
+
+        x = []
+        if k != 0:
+            for i in range(1, k + 1):
+                lower = round(
+                    log(1 - (0.5 / (k + 1)) * i) / log((0.5 / (k + 1)) * i), 2
+                )
+                upper = round(1 / lower, 2)
+                x.append(lower)
+                x.append(upper)
+        x.append(1.00)
+        return x
+
+    def forward(self, batch_I, weights, offsets, lambda_color=None):
+        """
+
+        Args:
+            batch_I (Tensor): batch of input images [batch_size x nc x I_height x I_width]
+            weights:
+            offsets: the predicted offset by AIN, a scalar
+            lambda_color: the learnable update gate \alpha in Equa. (5) as
+                          g(x) = (1 - \alpha) \odot x + \alpha \odot x_{offsets}
+
+        Returns:
+            Tensor: transformed images by SPN as Equa. (4) in Ref. [1]
+                        [batch_size x I_channel_num x I_r_height x I_r_width]
+
+        """
+        batch_I = (batch_I + 1) * 0.5
+        if offsets is not None:
+            batch_I = batch_I * (1 - lambda_color) + offsets * lambda_color
+        batch_weight_params = paddle.unsqueeze(paddle.unsqueeze(weights, -1), -1)
+        batch_I_power = paddle.stack([batch_I.pow(p) for p in self.power_list], axis=1)
+
+        batch_weight_sum = paddle.sum(batch_I_power * batch_weight_params, axis=1)
+        batch_weight_sum = self.bn(batch_weight_sum)
+        batch_weight_sum = self.sigmoid(batch_weight_sum)
+        batch_weight_sum = batch_weight_sum * 2 - 1
+        return batch_weight_sum
+
+
+class GA_SPIN_Transformer(nn.Layer):
+    """
+    Geometric-Absorbed SPIN Transformation (GA-SPIN) proposed in Ref. [1]
+
+
+    Ref: [1] SPIN: Structure-Preserving Inner Offset Network for Scene Text Recognition. AAAI-2021.
+    """
+
+    def __init__(
+        self,
+        in_channels=1,
+        I_r_size=(32, 100),
+        offsets=False,
+        norm_type="BN",
+        default_type=6,
+        loc_lr=1,
+        stn=True,
+    ):
+        """
+        Args:
+            in_channels (int): channel of input features,
+                                set it to 1 if the grayscale images and 3 if RGB input
+            I_r_size (tuple): size of rectified images (used in STN transformations)
+            offsets (bool): set it to False if use SPN w.o. AIN,
+                            and set it to True if use SPIN (both with SPN and AIN)
+            norm_type (str): the normalization type of the module,
+                            set it to 'BN' by default, 'IN' optionally
+            default_type (int): the K chromatic space,
+                                set it to 3/5/6 depend on the complexity of transformation intensities
+            loc_lr (float): learning rate of location network
+            stn (bool): whether to use stn.
+
+        """
+        super(GA_SPIN_Transformer, self).__init__()
+        self.nc = in_channels
+        self.spt = True
+        self.offsets = offsets
+        self.stn = stn  # set to True in GA-SPIN, while set it to False in SPIN
+        self.I_r_size = I_r_size
+        self.out_channels = in_channels
+        if norm_type == "BN":
+            norm_layer = functools.partial(nn.BatchNorm2D, use_global_stats=True)
+        elif norm_type == "IN":
+            norm_layer = functools.partial(
+                nn.InstanceNorm2D, weight_attr=False, use_global_stats=False
+            )
+        else:
+            raise NotImplementedError(
+                "normalization layer [%s] is not found" % norm_type
+            )
+
+        if self.spt:
+            self.sp_net = SP_TransformerNetwork(in_channels, default_type)
+            self.spt_convnet = nn.Sequential(
+                # 32*100
+                nn.Conv2D(in_channels, 32, 3, 1, 1, bias_attr=False),
+                norm_layer(32),
+                nn.ReLU(),
+                nn.MaxPool2D(kernel_size=2, stride=2),
+                # 16*50
+                nn.Conv2D(32, 64, 3, 1, 1, bias_attr=False),
+                norm_layer(64),
+                nn.ReLU(),
+                nn.MaxPool2D(kernel_size=2, stride=2),
+                # 8*25
+                nn.Conv2D(64, 128, 3, 1, 1, bias_attr=False),
+                norm_layer(128),
+                nn.ReLU(),
+                nn.MaxPool2D(kernel_size=2, stride=2),
+                # 4*12
+            )
+            self.stucture_fc1 = nn.Sequential(
+                nn.Conv2D(128, 256, 3, 1, 1, bias_attr=False),
+                norm_layer(256),
+                nn.ReLU(),
+                nn.MaxPool2D(kernel_size=2, stride=2),
+                nn.Conv2D(256, 256, 3, 1, 1, bias_attr=False),
+                norm_layer(256),
+                nn.ReLU(),  # 2*6
+                nn.MaxPool2D(kernel_size=2, stride=2),
+                nn.Conv2D(256, 512, 3, 1, 1, bias_attr=False),
+                norm_layer(512),
+                nn.ReLU(),  # 1*3
+                nn.AdaptiveAvgPool2D(1),
+                nn.Flatten(1, -1),  # batch_size x 512
+                nn.Linear(512, 256, weight_attr=nn.initializer.Normal(0.001)),
+                nn.BatchNorm1D(256),
+                nn.ReLU(),
+            )
+            self.out_weight = 2 * default_type + 1
+            self.spt_length = 2 * default_type + 1
+            if offsets:
+                self.out_weight += 1
+            if self.stn:
+                self.F = 20
+                self.out_weight += self.F * 2
+                self.GridGenerator = GridGenerator(self.F * 2, self.F)
+
+            # self.out_weight*=nc
+            # Init structure_fc2 in LocalizationNetwork
+            initial_bias = self.init_spin(default_type * 2)
+            initial_bias = initial_bias.reshape(-1)
+            param_attr = ParamAttr(
+                learning_rate=loc_lr,
+                initializer=nn.initializer.Assign(np.zeros([256, self.out_weight])),
+            )
+            bias_attr = ParamAttr(
+                learning_rate=loc_lr, initializer=nn.initializer.Assign(initial_bias)
+            )
+            self.stucture_fc2 = nn.Linear(
+                256, self.out_weight, weight_attr=param_attr, bias_attr=bias_attr
+            )
+            self.sigmoid = nn.Sigmoid()
+
+            if offsets:
+                self.offset_fc1 = nn.Sequential(
+                    nn.Conv2D(128, 16, 3, 1, 1, bias_attr=False),
+                    norm_layer(16),
+                    nn.ReLU(),
+                )
+                self.offset_fc2 = nn.Conv2D(16, in_channels, 3, 1, 1)
+                self.pool = nn.MaxPool2D(2, 2)
+
+    def init_spin(self, nz):
+        """
+        Args:
+            nz (int): number of paired \betas exponents, which means the value of K x 2
+
+        """
+        init_id = [0.00] * nz + [5.00]
+        if self.offsets:
+            init_id += [-5.00]
+            # init_id *=3
+        init = np.array(init_id)
+
+        if self.stn:
+            F = self.F
+            ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
+            ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2))
+            ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2))
+            ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+            ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+            initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+            initial_bias = initial_bias.reshape(-1)
+            init = np.concatenate([init, initial_bias], axis=0)
+        return init
+
+    def forward(self, x, return_weight=False):
+        """
+        Args:
+            x (Tensor): input image batch
+            return_weight (bool): set to False by default,
+                                  if set to True return the predicted offsets of AIN, denoted as x_{offsets}
+
+        Returns:
+            Tensor: rectified image [batch_size x I_channel_num x I_height x I_width], the same as the input size
+        """
+
+        if self.spt:
+            feat = self.spt_convnet(x)
+            fc1 = self.stucture_fc1(feat)
+            sp_weight_fusion = self.stucture_fc2(fc1)
+            sp_weight_fusion = sp_weight_fusion.reshape(
+                [x.shape[0], self.out_weight, 1]
+            )
+            if self.offsets:  # SPIN w. AIN
+                lambda_color = sp_weight_fusion[:, self.spt_length, 0]
+                lambda_color = (
+                    self.sigmoid(lambda_color).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+                )
+                sp_weight = sp_weight_fusion[:, : self.spt_length, :]
+                offsets = self.pool(self.offset_fc2(self.offset_fc1(feat)))
+
+                assert offsets.shape[2] == 2  # 2
+                assert offsets.shape[3] == 6  # 16
+                offsets = self.sigmoid(offsets)  # v12
+
+                if return_weight:
+                    return offsets
+                offsets = nn.functional.upsample(
+                    offsets, size=(x.shape[2], x.shape[3]), mode="bilinear"
+                )
+
+                if self.stn:
+                    batch_C_prime = sp_weight_fusion[
+                        :, (self.spt_length + 1) :, :
+                    ].reshape([x.shape[0], self.F, 2])
+                    build_P_prime = self.GridGenerator(batch_C_prime, self.I_r_size)
+                    build_P_prime_reshape = build_P_prime.reshape(
+                        [build_P_prime.shape[0], self.I_r_size[0], self.I_r_size[1], 2]
+                    )
+
+            else:  # SPIN w.o. AIN
+                sp_weight = sp_weight_fusion[:, : self.spt_length, :]
+                lambda_color, offsets = None, None
+
+                if self.stn:
+                    batch_C_prime = sp_weight_fusion[:, self.spt_length :, :].reshape(
+                        [x.shape[0], self.F, 2]
+                    )
+                    build_P_prime = self.GridGenerator(batch_C_prime, self.I_r_size)
+                    build_P_prime_reshape = build_P_prime.reshape(
+                        [build_P_prime.shape[0], self.I_r_size[0], self.I_r_size[1], 2]
+                    )
+
+            x = self.sp_net(x, sp_weight, offsets, lambda_color)
+            if self.stn:
+                is_fp16 = False
+                if build_P_prime_reshape.dtype != paddle.float32:
+                    data_type = build_P_prime_reshape.dtype
+                    x = x.cast(paddle.float32)
+                    build_P_prime_reshape = build_P_prime_reshape.cast(paddle.float32)
+                    is_fp16 = True
+                x = F.grid_sample(
+                    x=x, grid=build_P_prime_reshape, padding_mode="border"
+                )
+                if is_fp16:
+                    x = x.cast(data_type)
+        return x

ppocr/modeling/transforms/stn.py

@@ -0,0 +1,147 @@
+# 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.
+"""
+This code is refer from:
+https://github.com/ayumiymk/aster.pytorch/blob/master/lib/models/stn_head.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import numpy as np
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+
+
+def conv3x3_block(in_channels, out_channels, stride=1):
+    n = 3 * 3 * out_channels
+    w = math.sqrt(2.0 / n)
+    conv_layer = nn.Conv2D(
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        weight_attr=nn.initializer.Normal(mean=0.0, std=w),
+        bias_attr=nn.initializer.Constant(0),
+    )
+    block = nn.Sequential(conv_layer, nn.BatchNorm2D(out_channels), nn.ReLU())
+    return block
+
+
+class STN(nn.Layer):
+    def __init__(self, in_channels, num_ctrlpoints, activation="none"):
+        super(STN, self).__init__()
+        self.in_channels = in_channels
+        self.num_ctrlpoints = num_ctrlpoints
+        self.activation = activation
+        self.stn_convnet = nn.Sequential(
+            conv3x3_block(in_channels, 32),  # 32x64
+            nn.MaxPool2D(kernel_size=2, stride=2),
+            conv3x3_block(32, 64),  # 16x32
+            nn.MaxPool2D(kernel_size=2, stride=2),
+            conv3x3_block(64, 128),  # 8*16
+            nn.MaxPool2D(kernel_size=2, stride=2),
+            conv3x3_block(128, 256),  # 4*8
+            nn.MaxPool2D(kernel_size=2, stride=2),
+            conv3x3_block(256, 256),  # 2*4,
+            nn.MaxPool2D(kernel_size=2, stride=2),
+            conv3x3_block(256, 256),
+        )  # 1*2
+        self.stn_fc1 = nn.Sequential(
+            nn.Linear(
+                2 * 256,
+                512,
+                weight_attr=nn.initializer.Normal(0, 0.001),
+                bias_attr=nn.initializer.Constant(0),
+            ),
+            nn.BatchNorm1D(512),
+            nn.ReLU(),
+        )
+        fc2_bias = self.init_stn()
+        self.stn_fc2 = nn.Linear(
+            512,
+            num_ctrlpoints * 2,
+            weight_attr=nn.initializer.Constant(0.0),
+            bias_attr=nn.initializer.Assign(fc2_bias),
+        )
+
+    def init_stn(self):
+        margin = 0.01
+        sampling_num_per_side = int(self.num_ctrlpoints / 2)
+        ctrl_pts_x = np.linspace(margin, 1.0 - margin, sampling_num_per_side)
+        ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
+        ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1 - margin)
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0).astype(
+            np.float32
+        )
+        if self.activation == "none":
+            pass
+        elif self.activation == "sigmoid":
+            ctrl_points = -np.log(1.0 / ctrl_points - 1.0)
+        ctrl_points = paddle.to_tensor(ctrl_points)
+        fc2_bias = paddle.reshape(
+            ctrl_points, shape=[ctrl_points.shape[0] * ctrl_points.shape[1]]
+        )
+        return fc2_bias
+
+    def forward(self, x):
+        x = self.stn_convnet(x)
+        batch_size, _, h, w = x.shape
+        x = paddle.reshape(x, shape=(batch_size, -1))
+        img_feat = self.stn_fc1(x)
+        x = self.stn_fc2(0.1 * img_feat)
+        if self.activation == "sigmoid":
+            x = F.sigmoid(x)
+        x = paddle.reshape(x, shape=[-1, self.num_ctrlpoints, 2])
+        return img_feat, x
+
+
+class STN_ON(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        tps_inputsize,
+        tps_outputsize,
+        num_control_points,
+        tps_margins,
+        stn_activation,
+    ):
+        super(STN_ON, self).__init__()
+        self.tps = TPSSpatialTransformer(
+            output_image_size=tuple(tps_outputsize),
+            num_control_points=num_control_points,
+            margins=tuple(tps_margins),
+        )
+        self.stn_head = STN(
+            in_channels=in_channels,
+            num_ctrlpoints=num_control_points,
+            activation=stn_activation,
+        )
+        self.tps_inputsize = tps_inputsize
+        self.out_channels = in_channels
+
+    def forward(self, image):
+        stn_input = paddle.nn.functional.interpolate(
+            image, self.tps_inputsize, mode="bilinear", align_corners=True
+        )
+        stn_img_feat, ctrl_points = self.stn_head(stn_input)
+        x, _ = self.tps(image, ctrl_points)
+        return x

ppocr/modeling/transforms/tbsrn.py

@@ -0,0 +1,298 @@
+# copyright (c) 2022 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.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/scene-text-telescope/model/tbsrn.py
+"""
+
+import math
+import warnings
+import numpy as np
+import paddle
+from paddle import nn
+import string
+
+warnings.filterwarnings("ignore")
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+from .stn import STN as STNHead
+from .tsrn import GruBlock, mish, UpsampleBLock
+from ppocr.modeling.heads.sr_rensnet_transformer import (
+    Transformer,
+    LayerNorm,
+    PositionwiseFeedForward,
+    MultiHeadedAttention,
+)
+
+
+def positionalencoding2d(d_model, height, width):
+    """
+    :param d_model: dimension of the model
+    :param height: height of the positions
+    :param width: width of the positions
+    :return: d_model*height*width position matrix
+    """
+    if d_model % 4 != 0:
+        raise ValueError(
+            "Cannot use sin/cos positional encoding with "
+            "odd dimension (got dim={:d})".format(d_model)
+        )
+    pe = paddle.zeros([d_model, height, width])
+    # Each dimension use half of d_model
+    d_model = int(d_model / 2)
+    div_term = paddle.exp(
+        paddle.arange(0.0, d_model, 2, dtype="int64") * -(math.log(10000.0) / d_model)
+    )
+    pos_w = paddle.arange(0.0, width, dtype="float32").unsqueeze(1)
+    pos_h = paddle.arange(0.0, height, dtype="float32").unsqueeze(1)
+
+    pe[0:d_model:2, :, :] = (
+        paddle.sin(pos_w * div_term).transpose([1, 0]).unsqueeze(1).tile([1, height, 1])
+    )
+    pe[1:d_model:2, :, :] = (
+        paddle.cos(pos_w * div_term).transpose([1, 0]).unsqueeze(1).tile([1, height, 1])
+    )
+    pe[d_model::2, :, :] = (
+        paddle.sin(pos_h * div_term).transpose([1, 0]).unsqueeze(2).tile([1, 1, width])
+    )
+    pe[d_model + 1 :: 2, :, :] = (
+        paddle.cos(pos_h * div_term).transpose([1, 0]).unsqueeze(2).tile([1, 1, width])
+    )
+
+    return pe
+
+
+class FeatureEnhancer(nn.Layer):
+    def __init__(self):
+        super(FeatureEnhancer, self).__init__()
+
+        self.multihead = MultiHeadedAttention(h=4, d_model=128, dropout=0.1)
+        self.mul_layernorm1 = LayerNorm(features=128)
+
+        self.pff = PositionwiseFeedForward(128, 128)
+        self.mul_layernorm3 = LayerNorm(features=128)
+
+        self.linear = nn.Linear(128, 64)
+
+    def forward(self, conv_feature):
+        """
+        text : (batch, seq_len, embedding_size)
+        global_info: (batch, embedding_size, 1, 1)
+        conv_feature: (batch, channel, H, W)
+        """
+        batch = conv_feature.shape[0]
+        position2d = (
+            positionalencoding2d(64, 16, 64)
+            .cast("float32")
+            .unsqueeze(0)
+            .reshape([1, 64, 1024])
+        )
+        position2d = position2d.tile([batch, 1, 1])
+        conv_feature = paddle.concat(
+            [conv_feature, position2d], 1
+        )  # batch, 128(64+64), 32, 128
+        result = conv_feature.transpose([0, 2, 1])
+        origin_result = result
+        result = self.mul_layernorm1(
+            origin_result + self.multihead(result, result, result, mask=None)[0]
+        )
+        origin_result = result
+        result = self.mul_layernorm3(origin_result + self.pff(result))
+        result = self.linear(result)
+        return result.transpose([0, 2, 1])
+
+
+def str_filt(str_, voc_type):
+    alpha_dict = {
+        "digit": string.digits,
+        "lower": string.digits + string.ascii_lowercase,
+        "upper": string.digits + string.ascii_letters,
+        "all": string.digits + string.ascii_letters + string.punctuation,
+    }
+    if voc_type == "lower":
+        str_ = str_.lower()
+    for char in str_:
+        if char not in alpha_dict[voc_type]:
+            str_ = str_.replace(char, "")
+    str_ = str_.lower()
+    return str_
+
+
+class TBSRN(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        scale_factor=2,
+        width=128,
+        height=32,
+        STN=True,
+        srb_nums=5,
+        mask=False,
+        hidden_units=32,
+        infer_mode=False,
+    ):
+        super(TBSRN, self).__init__()
+        in_planes = 3
+        if mask:
+            in_planes = 4
+        assert math.log(scale_factor, 2) % 1 == 0
+        upsample_block_num = int(math.log(scale_factor, 2))
+        self.block1 = nn.Sequential(
+            nn.Conv2D(in_planes, 2 * hidden_units, kernel_size=9, padding=4),
+            nn.PReLU(),
+            # nn.ReLU()
+        )
+        self.srb_nums = srb_nums
+        for i in range(srb_nums):
+            setattr(self, "block%d" % (i + 2), RecurrentResidualBlock(2 * hidden_units))
+
+        setattr(
+            self,
+            "block%d" % (srb_nums + 2),
+            nn.Sequential(
+                nn.Conv2D(2 * hidden_units, 2 * hidden_units, kernel_size=3, padding=1),
+                nn.BatchNorm2D(2 * hidden_units),
+            ),
+        )
+
+        # self.non_local = NonLocalBlock2D(64, 64)
+        block_ = [UpsampleBLock(2 * hidden_units, 2) for _ in range(upsample_block_num)]
+        block_.append(nn.Conv2D(2 * hidden_units, in_planes, kernel_size=9, padding=4))
+        setattr(self, "block%d" % (srb_nums + 3), nn.Sequential(*block_))
+        self.tps_inputsize = [height // scale_factor, width // scale_factor]
+        tps_outputsize = [height // scale_factor, width // scale_factor]
+        num_control_points = 20
+        tps_margins = [0.05, 0.05]
+        self.stn = STN
+        self.out_channels = in_channels
+        if self.stn:
+            self.tps = TPSSpatialTransformer(
+                output_image_size=tuple(tps_outputsize),
+                num_control_points=num_control_points,
+                margins=tuple(tps_margins),
+            )
+
+            self.stn_head = STNHead(
+                in_channels=in_planes,
+                num_ctrlpoints=num_control_points,
+                activation="none",
+            )
+        self.infer_mode = infer_mode
+
+        self.english_alphabet = (
+            "-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+        )
+        self.english_dict = {}
+        for index in range(len(self.english_alphabet)):
+            self.english_dict[self.english_alphabet[index]] = index
+        transformer = Transformer(alphabet="-0123456789abcdefghijklmnopqrstuvwxyz")
+        self.transformer = transformer
+        for param in self.transformer.parameters():
+            param.trainable = False
+
+    def label_encoder(self, label):
+        batch = len(label)
+
+        length = [len(i) for i in label]
+        length_tensor = paddle.to_tensor(length, dtype="int64")
+
+        max_length = max(length)
+        input_tensor = np.zeros((batch, max_length))
+        for i in range(batch):
+            for j in range(length[i] - 1):
+                input_tensor[i][j + 1] = self.english_dict[label[i][j]]
+
+        text_gt = []
+        for i in label:
+            for j in i:
+                text_gt.append(self.english_dict[j])
+        text_gt = paddle.to_tensor(text_gt, dtype="int64")
+
+        input_tensor = paddle.to_tensor(input_tensor, dtype="int64")
+        return length_tensor, input_tensor, text_gt
+
+    def forward(self, x):
+        output = {}
+        if self.infer_mode:
+            output["lr_img"] = x
+            y = x
+        else:
+            output["lr_img"] = x[0]
+            output["hr_img"] = x[1]
+            y = x[0]
+        if self.stn and self.training:
+            _, ctrl_points_x = self.stn_head(y)
+            y, _ = self.tps(y, ctrl_points_x)
+        block = {"1": self.block1(y)}
+        for i in range(self.srb_nums + 1):
+            block[str(i + 2)] = getattr(self, "block%d" % (i + 2))(block[str(i + 1)])
+
+        block[str(self.srb_nums + 3)] = getattr(self, "block%d" % (self.srb_nums + 3))(
+            (block["1"] + block[str(self.srb_nums + 2)])
+        )
+
+        sr_img = paddle.tanh(block[str(self.srb_nums + 3)])
+        output["sr_img"] = sr_img
+
+        if self.training:
+            hr_img = x[1]
+
+            # add transformer
+            label = [str_filt(i, "lower") + "-" for i in x[2]]
+            length_tensor, input_tensor, text_gt = self.label_encoder(label)
+            hr_pred, word_attention_map_gt, hr_correct_list = self.transformer(
+                hr_img, length_tensor, input_tensor
+            )
+            sr_pred, word_attention_map_pred, sr_correct_list = self.transformer(
+                sr_img, length_tensor, input_tensor
+            )
+            output["hr_img"] = hr_img
+            output["hr_pred"] = hr_pred
+            output["text_gt"] = text_gt
+            output["word_attention_map_gt"] = word_attention_map_gt
+            output["sr_pred"] = sr_pred
+            output["word_attention_map_pred"] = word_attention_map_pred
+
+        return output
+
+
+class RecurrentResidualBlock(nn.Layer):
+    def __init__(self, channels):
+        super(RecurrentResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2D(channels, channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2D(channels)
+        self.gru1 = GruBlock(channels, channels)
+        # self.prelu = nn.ReLU()
+        self.prelu = mish()
+        self.conv2 = nn.Conv2D(channels, channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2D(channels)
+        self.gru2 = GruBlock(channels, channels)
+        self.feature_enhancer = FeatureEnhancer()
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                paddle.nn.initializer.XavierUniform(p)
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = self.bn1(residual)
+        residual = self.prelu(residual)
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+
+        size = residual.shape
+        residual = residual.reshape([size[0], size[1], -1])
+        residual = self.feature_enhancer(residual)
+        residual = residual.reshape([size[0], size[1], size[2], size[3]])
+        return x + residual

ppocr/modeling/transforms/tps.py

@@ -0,0 +1,321 @@
+# 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.
+"""
+This code is refer from:
+https://github.com/clovaai/deep-text-recognition-benchmark/blob/master/modules/transformation.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import numpy as np
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+        bn_name = "bn_" + name
+        self.bn = nn.BatchNorm(
+            out_channels,
+            act=act,
+            param_attr=ParamAttr(name=bn_name + "_scale"),
+            bias_attr=ParamAttr(bn_name + "_offset"),
+            moving_mean_name=bn_name + "_mean",
+            moving_variance_name=bn_name + "_variance",
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class LocalizationNetwork(nn.Layer):
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(LocalizationNetwork, self).__init__()
+        self.F = num_fiducial
+        F = num_fiducial
+        if model_name == "large":
+            num_filters_list = [64, 128, 256, 512]
+            fc_dim = 256
+        else:
+            num_filters_list = [16, 32, 64, 128]
+            fc_dim = 64
+
+        self.block_list = []
+        for fno in range(0, len(num_filters_list)):
+            num_filters = num_filters_list[fno]
+            name = "loc_conv%d" % fno
+            conv = self.add_sublayer(
+                name,
+                ConvBNLayer(
+                    in_channels=in_channels,
+                    out_channels=num_filters,
+                    kernel_size=3,
+                    act="relu",
+                    name=name,
+                ),
+            )
+            self.block_list.append(conv)
+            if fno == len(num_filters_list) - 1:
+                pool = nn.AdaptiveAvgPool2D(1)
+            else:
+                pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+            in_channels = num_filters
+            self.block_list.append(pool)
+        name = "loc_fc1"
+        stdv = 1.0 / math.sqrt(num_filters_list[-1] * 1.0)
+        self.fc1 = nn.Linear(
+            in_channels,
+            fc_dim,
+            weight_attr=ParamAttr(
+                learning_rate=loc_lr,
+                name=name + "_w",
+                initializer=nn.initializer.Uniform(-stdv, stdv),
+            ),
+            bias_attr=ParamAttr(name=name + ".b_0"),
+            name=name,
+        )
+
+        # Init fc2 in LocalizationNetwork
+        initial_bias = self.get_initial_fiducials()
+        initial_bias = initial_bias.reshape(-1)
+        name = "loc_fc2"
+        param_attr = ParamAttr(
+            learning_rate=loc_lr,
+            initializer=nn.initializer.Assign(np.zeros([fc_dim, F * 2])),
+            name=name + "_w",
+        )
+        bias_attr = ParamAttr(
+            learning_rate=loc_lr,
+            initializer=nn.initializer.Assign(initial_bias),
+            name=name + "_b",
+        )
+        self.fc2 = nn.Linear(
+            fc_dim, F * 2, weight_attr=param_attr, bias_attr=bias_attr, name=name
+        )
+        self.out_channels = F * 2
+
+    def forward(self, x):
+        """
+        Estimating parameters of geometric transformation
+        Args:
+            image: input
+        Return:
+            batch_C_prime: the matrix of the geometric transformation
+        """
+        B = x.shape[0]
+        i = 0
+        for block in self.block_list:
+            x = block(x)
+        x = x.squeeze(axis=2).squeeze(axis=2)
+        x = self.fc1(x)
+
+        x = F.relu(x)
+        x = self.fc2(x)
+        x = x.reshape(shape=[-1, self.F, 2])
+        return x
+
+    def get_initial_fiducials(self):
+        """see RARE paper Fig. 6 (a)"""
+        F = self.F
+        ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
+        ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2))
+        ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2))
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+        return initial_bias
+
+
+class GridGenerator(nn.Layer):
+    def __init__(self, in_channels, num_fiducial):
+        super(GridGenerator, self).__init__()
+        self.eps = 1e-6
+        self.F = num_fiducial
+
+        name = "ex_fc"
+        initializer = nn.initializer.Constant(value=0.0)
+        param_attr = ParamAttr(
+            learning_rate=0.0, initializer=initializer, name=name + "_w"
+        )
+        bias_attr = ParamAttr(
+            learning_rate=0.0, initializer=initializer, name=name + "_b"
+        )
+        self.fc = nn.Linear(
+            in_channels, 6, weight_attr=param_attr, bias_attr=bias_attr, name=name
+        )
+
+    def forward(self, batch_C_prime, I_r_size):
+        """
+        Generate the grid for the grid_sampler.
+        Args:
+            batch_C_prime: the matrix of the geometric transformation
+            I_r_size: the shape of the input image
+        Return:
+            batch_P_prime: the grid for the grid_sampler
+        """
+        C = self.build_C_paddle()
+        P = self.build_P_paddle(I_r_size)
+
+        inv_delta_C_tensor = self.build_inv_delta_C_paddle(C).astype("float32")
+        P_hat_tensor = self.build_P_hat_paddle(C, paddle.to_tensor(P)).astype("float32")
+
+        inv_delta_C_tensor.stop_gradient = True
+        P_hat_tensor.stop_gradient = True
+
+        batch_C_ex_part_tensor = self.get_expand_tensor(batch_C_prime)
+
+        batch_C_ex_part_tensor.stop_gradient = True
+
+        batch_C_prime_with_zeros = paddle.concat(
+            [batch_C_prime, batch_C_ex_part_tensor], axis=1
+        )
+        batch_T = paddle.matmul(inv_delta_C_tensor, batch_C_prime_with_zeros)
+        batch_P_prime = paddle.matmul(P_hat_tensor, batch_T)
+        return batch_P_prime
+
+    def build_C_paddle(self):
+        """Return coordinates of fiducial points in I_r; C"""
+        F = self.F
+        ctrl_pts_x = paddle.linspace(-1.0, 1.0, int(F / 2), dtype="float64")
+        ctrl_pts_y_top = -1 * paddle.ones([int(F / 2)], dtype="float64")
+        ctrl_pts_y_bottom = paddle.ones([int(F / 2)], dtype="float64")
+        ctrl_pts_top = paddle.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = paddle.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        C = paddle.concat([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+        return C  # F x 2
+
+    def build_P_paddle(self, I_r_size):
+        I_r_height, I_r_width = I_r_size
+        I_r_grid_x = (
+            paddle.arange(-I_r_width, I_r_width, 2, dtype="float64") + 1.0
+        ) / paddle.to_tensor(np.array([I_r_width])).astype("float64")
+
+        I_r_grid_y = (
+            paddle.arange(-I_r_height, I_r_height, 2, dtype="float64") + 1.0
+        ) / paddle.to_tensor(np.array([I_r_height])).astype("float64")
+
+        # P: self.I_r_width x self.I_r_height x 2
+        P = paddle.stack(paddle.meshgrid(I_r_grid_x, I_r_grid_y), axis=2)
+        P = paddle.transpose(P, perm=[1, 0, 2])
+        # n (= self.I_r_width x self.I_r_height) x 2
+        return P.reshape([-1, 2])
+
+    def build_inv_delta_C_paddle(self, C):
+        """Return inv_delta_C which is needed to calculate T"""
+        F = self.F
+        hat_eye = paddle.eye(F, dtype="float64")  # F x F
+        hat_C = (
+            paddle.norm(C.reshape([1, F, 2]) - C.reshape([F, 1, 2]), axis=2) + hat_eye
+        )
+        hat_C = (hat_C**2) * paddle.log(hat_C)
+        delta_C = paddle.concat(  # F+3 x F+3
+            [
+                paddle.concat(
+                    [paddle.ones((F, 1), dtype="float64"), C, hat_C], axis=1
+                ),  # F x F+3
+                paddle.concat(
+                    [
+                        paddle.zeros((2, 3), dtype="float64"),
+                        paddle.transpose(C, perm=[1, 0]),
+                    ],
+                    axis=1,
+                ),  # 2 x F+3
+                paddle.concat(
+                    [
+                        paddle.zeros((1, 3), dtype="float64"),
+                        paddle.ones((1, F), dtype="float64"),
+                    ],
+                    axis=1,
+                ),  # 1 x F+3
+            ],
+            axis=0,
+        )
+        inv_delta_C = paddle.inverse(delta_C)
+        return inv_delta_C  # F+3 x F+3
+
+    def build_P_hat_paddle(self, C, P):
+        F = self.F
+        eps = self.eps
+        n = P.shape[0]  # n (= self.I_r_width x self.I_r_height)
+        # P_tile: n x 2 -> n x 1 x 2 -> n x F x 2
+        P_tile = paddle.tile(paddle.unsqueeze(P, axis=1), (1, F, 1))
+        C_tile = paddle.unsqueeze(C, axis=0)  # 1 x F x 2
+        P_diff = P_tile - C_tile  # n x F x 2
+        # rbf_norm: n x F
+        rbf_norm = paddle.norm(P_diff, p=2, axis=2, keepdim=False)
+
+        # rbf: n x F
+        rbf = paddle.multiply(paddle.square(rbf_norm), paddle.log(rbf_norm + eps))
+        P_hat = paddle.concat([paddle.ones((n, 1), dtype="float64"), P, rbf], axis=1)
+        return P_hat  # n x F+3
+
+    def get_expand_tensor(self, batch_C_prime):
+        B, H, C = batch_C_prime.shape
+        batch_C_prime = batch_C_prime.reshape([B, H * C])
+        batch_C_ex_part_tensor = self.fc(batch_C_prime)
+        batch_C_ex_part_tensor = batch_C_ex_part_tensor.reshape([-1, 3, 2])
+        return batch_C_ex_part_tensor
+
+
+class TPS(nn.Layer):
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(TPS, self).__init__()
+        self.loc_net = LocalizationNetwork(
+            in_channels, num_fiducial, loc_lr, model_name
+        )
+        self.grid_generator = GridGenerator(self.loc_net.out_channels, num_fiducial)
+        self.out_channels = in_channels
+
+    def forward(self, image):
+        image.stop_gradient = False
+        batch_C_prime = self.loc_net(image)
+        batch_P_prime = self.grid_generator(batch_C_prime, image.shape[2:])
+        batch_P_prime = batch_P_prime.reshape([-1, image.shape[2], image.shape[3], 2])
+        is_fp16 = False
+        if batch_P_prime.dtype != paddle.float32:
+            data_type = batch_P_prime.dtype
+            image = image.cast(paddle.float32)
+            batch_P_prime = batch_P_prime.cast(paddle.float32)
+            is_fp16 = True
+        batch_I_r = F.grid_sample(x=image, grid=batch_P_prime)
+        if is_fp16:
+            batch_I_r = batch_I_r.cast(data_type)
+
+        return batch_I_r

ppocr/modeling/transforms/tps_spatial_transformer.py

@@ -0,0 +1,170 @@
+# 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.
+"""
+This code is refer from:
+https://github.com/ayumiymk/aster.pytorch/blob/master/lib/models/tps_spatial_transformer.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import numpy as np
+import itertools
+
+
+def grid_sample(input, grid, canvas=None):
+    input.stop_gradient = False
+
+    is_fp16 = False
+    if grid.dtype != paddle.float32:
+        data_type = grid.dtype
+        input = input.cast(paddle.float32)
+        grid = grid.cast(paddle.float32)
+        is_fp16 = True
+    output = F.grid_sample(input, grid)
+    if is_fp16:
+        output = output.cast(data_type)
+        grid = grid.cast(data_type)
+
+    if canvas is None:
+        return output
+    else:
+        input_mask = paddle.ones(shape=input.shape)
+        if is_fp16:
+            input_mask = input_mask.cast(paddle.float32)
+            grid = grid.cast(paddle.float32)
+        output_mask = F.grid_sample(input_mask, grid)
+        if is_fp16:
+            output_mask = output_mask.cast(data_type)
+        padded_output = output * output_mask + canvas * (1 - output_mask)
+        return padded_output
+
+
+# phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+def compute_partial_repr(input_points, control_points):
+    N = input_points.shape[0]
+    M = control_points.shape[0]
+    pairwise_diff = paddle.reshape(input_points, shape=[N, 1, 2]) - paddle.reshape(
+        control_points, shape=[1, M, 2]
+    )
+    # original implementation, very slow
+    # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+    pairwise_diff_square = pairwise_diff * pairwise_diff
+    pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
+    repr_matrix = 0.5 * pairwise_dist * paddle.log(pairwise_dist)
+    # fix numerical error for 0 * log(0), substitute all nan with 0
+    mask = np.array(repr_matrix != repr_matrix)
+    repr_matrix[mask] = 0
+    return repr_matrix
+
+
+# output_ctrl_pts are specified, according to our task.
+def build_output_control_points(num_control_points, margins):
+    margin_x, margin_y = margins
+    num_ctrl_pts_per_side = num_control_points // 2
+    ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+    ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+    ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+    ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+    ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+    output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+    output_ctrl_pts = paddle.to_tensor(output_ctrl_pts_arr)
+    return output_ctrl_pts
+
+
+class TPSSpatialTransformer(nn.Layer):
+    def __init__(self, output_image_size=None, num_control_points=None, margins=None):
+        super(TPSSpatialTransformer, self).__init__()
+        self.output_image_size = output_image_size
+        self.num_control_points = num_control_points
+        self.margins = margins
+
+        self.target_height, self.target_width = output_image_size
+        target_control_points = build_output_control_points(num_control_points, margins)
+        N = num_control_points
+
+        # create padded kernel matrix
+        forward_kernel = paddle.zeros(shape=[N + 3, N + 3])
+        target_control_partial_repr = compute_partial_repr(
+            target_control_points, target_control_points
+        )
+        target_control_partial_repr = paddle.cast(
+            target_control_partial_repr, forward_kernel.dtype
+        )
+        forward_kernel[:N, :N] = target_control_partial_repr
+        forward_kernel[:N, -3] = 1
+        forward_kernel[-3, :N] = 1
+        target_control_points = paddle.cast(target_control_points, forward_kernel.dtype)
+        forward_kernel[:N, -2:] = target_control_points
+        forward_kernel[-2:, :N] = paddle.transpose(target_control_points, perm=[1, 0])
+        # compute inverse matrix
+        inverse_kernel = paddle.inverse(forward_kernel)
+
+        # create target coordinate matrix
+        HW = self.target_height * self.target_width
+        target_coordinate = list(
+            itertools.product(range(self.target_height), range(self.target_width))
+        )
+        target_coordinate = paddle.to_tensor(target_coordinate)  # HW x 2
+        Y, X = paddle.split(target_coordinate, target_coordinate.shape[1], axis=1)
+        Y = Y / (self.target_height - 1)
+        X = X / (self.target_width - 1)
+        target_coordinate = paddle.concat(
+            [X, Y], axis=1
+        )  # convert from (y, x) to (x, y)
+        target_coordinate_partial_repr = compute_partial_repr(
+            target_coordinate, target_control_points
+        )
+        target_coordinate_repr = paddle.concat(
+            [
+                target_coordinate_partial_repr,
+                paddle.ones(shape=[HW, 1]),
+                target_coordinate,
+            ],
+            axis=1,
+        )
+
+        # register precomputed matrices
+        self.inverse_kernel = inverse_kernel
+        self.padding_matrix = paddle.zeros(shape=[3, 2])
+        self.target_coordinate_repr = target_coordinate_repr
+        self.target_control_points = target_control_points
+
+    def forward(self, input, source_control_points):
+        assert source_control_points.ndimension() == 3
+        assert source_control_points.shape[1] == self.num_control_points
+        assert source_control_points.shape[2] == 2
+        batch_size = source_control_points.shape[0]
+
+        padding_matrix = paddle.expand(self.padding_matrix, shape=[batch_size, 3, 2])
+        Y = paddle.concat(
+            [source_control_points.astype(padding_matrix.dtype), padding_matrix], 1
+        )
+        mapping_matrix = paddle.matmul(self.inverse_kernel, Y)
+        source_coordinate = paddle.matmul(self.target_coordinate_repr, mapping_matrix)
+
+        grid = paddle.reshape(
+            source_coordinate, shape=[-1, self.target_height, self.target_width, 2]
+        )
+        grid = paddle.clip(
+            grid, 0, 1
+        )  # the source_control_points may be out of [0, 1].
+        # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+        grid = 2.0 * grid - 1.0
+        output_maps = grid_sample(input, grid, canvas=None)
+        return output_maps, source_coordinate

ppocr/modeling/transforms/tsrn.py

@@ -0,0 +1,215 @@
+# 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.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/text-gestalt/model/tsrn.py
+"""
+
+import math
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from collections import OrderedDict
+import sys
+import numpy as np
+import warnings
+import math, copy
+import cv2
+
+warnings.filterwarnings("ignore")
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+from .stn import STN as STN_model
+from ppocr.modeling.heads.sr_rensnet_transformer import Transformer
+
+
+class TSRN(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        scale_factor=2,
+        width=128,
+        height=32,
+        STN=False,
+        srb_nums=5,
+        mask=False,
+        hidden_units=32,
+        infer_mode=False,
+        **kwargs,
+    ):
+        super(TSRN, self).__init__()
+        in_planes = 3
+        if mask:
+            in_planes = 4
+        assert math.log(scale_factor, 2) % 1 == 0
+        upsample_block_num = int(math.log(scale_factor, 2))
+        self.block1 = nn.Sequential(
+            nn.Conv2D(in_planes, 2 * hidden_units, kernel_size=9, padding=4), nn.PReLU()
+        )
+        self.srb_nums = srb_nums
+        for i in range(srb_nums):
+            setattr(self, "block%d" % (i + 2), RecurrentResidualBlock(2 * hidden_units))
+
+        setattr(
+            self,
+            "block%d" % (srb_nums + 2),
+            nn.Sequential(
+                nn.Conv2D(2 * hidden_units, 2 * hidden_units, kernel_size=3, padding=1),
+                nn.BatchNorm2D(2 * hidden_units),
+            ),
+        )
+
+        block_ = [UpsampleBLock(2 * hidden_units, 2) for _ in range(upsample_block_num)]
+        block_.append(nn.Conv2D(2 * hidden_units, in_planes, kernel_size=9, padding=4))
+        setattr(self, "block%d" % (srb_nums + 3), nn.Sequential(*block_))
+        self.tps_inputsize = [height // scale_factor, width // scale_factor]
+        tps_outputsize = [height // scale_factor, width // scale_factor]
+        num_control_points = 20
+        tps_margins = [0.05, 0.05]
+        self.stn = STN
+        if self.stn:
+            self.tps = TPSSpatialTransformer(
+                output_image_size=tuple(tps_outputsize),
+                num_control_points=num_control_points,
+                margins=tuple(tps_margins),
+            )
+
+            self.stn_head = STN_model(
+                in_channels=in_planes,
+                num_ctrlpoints=num_control_points,
+                activation="none",
+            )
+        self.out_channels = in_channels
+
+        self.r34_transformer = Transformer()
+        for param in self.r34_transformer.parameters():
+            param.trainable = False
+        self.infer_mode = infer_mode
+
+    def forward(self, x):
+        output = {}
+        if self.infer_mode:
+            output["lr_img"] = x
+            y = x
+        else:
+            output["lr_img"] = x[0]
+            output["hr_img"] = x[1]
+            y = x[0]
+        if self.stn and self.training:
+            _, ctrl_points_x = self.stn_head(y)
+            y, _ = self.tps(y, ctrl_points_x)
+        block = {"1": self.block1(y)}
+        for i in range(self.srb_nums + 1):
+            block[str(i + 2)] = getattr(self, "block%d" % (i + 2))(block[str(i + 1)])
+
+        block[str(self.srb_nums + 3)] = getattr(self, "block%d" % (self.srb_nums + 3))(
+            (block["1"] + block[str(self.srb_nums + 2)])
+        )
+
+        sr_img = paddle.tanh(block[str(self.srb_nums + 3)])
+
+        output["sr_img"] = sr_img
+
+        if self.training:
+            hr_img = x[1]
+            length = x[2]
+            input_tensor = x[3]
+
+            # add transformer
+            sr_pred, word_attention_map_pred, _ = self.r34_transformer(
+                sr_img, length, input_tensor
+            )
+
+            hr_pred, word_attention_map_gt, _ = self.r34_transformer(
+                hr_img, length, input_tensor
+            )
+
+            output["hr_img"] = hr_img
+            output["hr_pred"] = hr_pred
+            output["word_attention_map_gt"] = word_attention_map_gt
+            output["sr_pred"] = sr_pred
+            output["word_attention_map_pred"] = word_attention_map_pred
+
+        return output
+
+
+class RecurrentResidualBlock(nn.Layer):
+    def __init__(self, channels):
+        super(RecurrentResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2D(channels, channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2D(channels)
+        self.gru1 = GruBlock(channels, channels)
+        self.prelu = mish()
+        self.conv2 = nn.Conv2D(channels, channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2D(channels)
+        self.gru2 = GruBlock(channels, channels)
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = self.bn1(residual)
+        residual = self.prelu(residual)
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+        residual = self.gru1(residual.transpose([0, 1, 3, 2])).transpose([0, 1, 3, 2])
+
+        return self.gru2(x + residual)
+
+
+class UpsampleBLock(nn.Layer):
+    def __init__(self, in_channels, up_scale):
+        super(UpsampleBLock, self).__init__()
+        self.conv = nn.Conv2D(
+            in_channels, in_channels * up_scale**2, kernel_size=3, padding=1
+        )
+
+        self.pixel_shuffle = nn.PixelShuffle(up_scale)
+        self.prelu = mish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.pixel_shuffle(x)
+        x = self.prelu(x)
+        return x
+
+
+class mish(nn.Layer):
+    def __init__(
+        self,
+    ):
+        super(mish, self).__init__()
+        self.activated = True
+
+    def forward(self, x):
+        if self.activated:
+            x = x * (paddle.tanh(F.softplus(x)))
+        return x
+
+
+class GruBlock(nn.Layer):
+    def __init__(self, in_channels, out_channels):
+        super(GruBlock, self).__init__()
+        assert out_channels % 2 == 0
+        self.conv1 = nn.Conv2D(in_channels, out_channels, kernel_size=1, padding=0)
+        self.gru = nn.GRU(out_channels, out_channels // 2, direction="bidirectional")
+
+    def forward(self, x):
+        # x: b, c, w, h
+        x = self.conv1(x)
+        x = x.transpose([0, 2, 3, 1])  # b, w, h, c
+        batch_size, w, h, c = x.shape
+        x = x.reshape([-1, h, c])  # b*w, h, c
+        x, _ = self.gru(x)
+        x = x.reshape([-1, w, h, c])
+        x = x.transpose([0, 3, 1, 2])
+        return x