first commit

ppocr/modeling/architectures/__init__.py

@@ -0,0 +1,145 @@
+# 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.
+
+import copy
+import importlib
+
+from paddle.jit import to_static
+from paddle.static import InputSpec
+
+from .base_model import BaseModel
+from .distillation_model import DistillationModel
+
+__all__ = ["build_model", "apply_to_static"]
+
+
+def build_model(config):
+    config = copy.deepcopy(config)
+    if not "name" in config:
+        arch = BaseModel(config)
+    else:
+        name = config.pop("name")
+        mod = importlib.import_module(__name__)
+        arch = getattr(mod, name)(config)
+    return arch
+
+
+def apply_to_static(model, config, logger):
+    if config["Global"].get("to_static", False) is not True:
+        return model
+    assert (
+        "d2s_train_image_shape" in config["Global"]
+    ), "d2s_train_image_shape must be assigned for static training mode..."
+    supported_list = [
+        "DB",
+        "SVTR_LCNet",
+        "TableMaster",
+        "LayoutXLM",
+        "SLANet",
+        "SVTR",
+        "SVTR_HGNet",
+        "LaTeXOCR",
+        "UniMERNet",
+        "PP-FormulaNet-S",
+        "PP-FormulaNet-L",
+    ]
+    if config["Architecture"]["algorithm"] in ["Distillation"]:
+        algo = list(config["Architecture"]["Models"].values())[0]["algorithm"]
+    else:
+        algo = config["Architecture"]["algorithm"]
+    assert (
+        algo in supported_list
+    ), f"algorithms that supports static training must in in {supported_list} but got {algo}"
+
+    specs = [
+        InputSpec([None] + config["Global"]["d2s_train_image_shape"], dtype="float32")
+    ]
+
+    if algo == "SVTR_LCNet":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None], dtype="int64"),
+                InputSpec([None], dtype="float64"),
+            ]
+        )
+    elif algo == "TableMaster":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"], 4], dtype="float32"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"], 1], dtype="float32"
+                ),
+                InputSpec([None, 6], dtype="float32"),
+            ]
+        )
+    elif algo == "LayoutXLM":
+        specs = [
+            [
+                InputSpec(shape=[None, 512], dtype="int64"),  # input_ids
+                InputSpec(shape=[None, 512, 4], dtype="int64"),  # bbox
+                InputSpec(shape=[None, 512], dtype="int64"),  # attention_mask
+                InputSpec(shape=[None, 512], dtype="int64"),  # token_type_ids
+                InputSpec(shape=[None, 3, 224, 224], dtype="float32"),  # image
+                InputSpec(shape=[None, 512], dtype="int64"),  # label
+            ]
+        ]
+    elif algo == "SLANet":
+        specs.append(
+            [
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2], dtype="int64"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2, 4], dtype="float32"
+                ),
+                InputSpec(
+                    [None, config["Global"]["max_text_length"] + 2, 1], dtype="float32"
+                ),
+                InputSpec([None], dtype="int64"),
+                InputSpec([None, 6], dtype="float64"),
+            ]
+        )
+    elif algo == "SVTR":
+        specs.append(
+            [
+                InputSpec([None, config["Global"]["max_text_length"]], dtype="int64"),
+                InputSpec([None], dtype="int64"),
+            ]
+        )
+    elif algo == "LaTeXOCR":
+        specs = [
+            [
+                InputSpec(shape=[None, 1, None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+            ]
+        ]
+    elif algo in ["UniMERNet", "PP-FormulaNet-S", "PP-FormulaNet-L"]:
+        specs = [
+            [
+                InputSpec(
+                    [None] + config["Global"]["d2s_train_image_shape"], dtype="float32"
+                ),
+                InputSpec(shape=[None, None], dtype="float32"),
+                InputSpec(shape=[None, None], dtype="float32"),
+            ]
+        ]
+    model = to_static(model, input_spec=specs)
+    logger.info("Successfully to apply @to_static with specs: {}".format(specs))
+    return model

ppocr/modeling/architectures/base_model.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2021 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.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+from ppocr.modeling.transforms import build_transform
+from ppocr.modeling.backbones import build_backbone
+from ppocr.modeling.necks import build_neck
+from ppocr.modeling.heads import build_head
+
+__all__ = ["BaseModel"]
+
+
+class BaseModel(nn.Layer):
+    def __init__(self, config):
+        """
+        the module for OCR.
+        args:
+            config (dict): the super parameters for module.
+        """
+        super(BaseModel, self).__init__()
+        in_channels = config.get("in_channels", 3)
+        model_type = config["model_type"]
+        # build transform,
+        # for rec, transform can be TPS,None
+        # for det and cls, transform should to be None,
+        # if you make model differently, you can use transform in det and cls
+        if "Transform" not in config or config["Transform"] is None:
+            self.use_transform = False
+        else:
+            self.use_transform = True
+            config["Transform"]["in_channels"] = in_channels
+            self.transform = build_transform(config["Transform"])
+            in_channels = self.transform.out_channels
+
+        # build backbone, backbone is need for del, rec and cls
+        if "Backbone" not in config or config["Backbone"] is None:
+            self.use_backbone = False
+        else:
+            self.use_backbone = True
+            config["Backbone"]["in_channels"] = in_channels
+            self.backbone = build_backbone(config["Backbone"], model_type)
+            in_channels = self.backbone.out_channels
+
+        # build neck
+        # for rec, neck can be cnn,rnn or reshape(None)
+        # for det, neck can be FPN, BIFPN and so on.
+        # for cls, neck should be none
+        if "Neck" not in config or config["Neck"] is None:
+            self.use_neck = False
+        else:
+            self.use_neck = True
+            config["Neck"]["in_channels"] = in_channels
+            self.neck = build_neck(config["Neck"])
+            in_channels = self.neck.out_channels
+
+        # # build head, head is need for det, rec and cls
+        if "Head" not in config or config["Head"] is None:
+            self.use_head = False
+        else:
+            self.use_head = True
+            config["Head"]["in_channels"] = in_channels
+            self.head = build_head(config["Head"])
+
+        self.return_all_feats = config.get("return_all_feats", False)
+
+    def forward(self, x, data=None):
+        y = dict()
+        if self.use_transform:
+            x = self.transform(x)
+        if self.use_backbone:
+            x = self.backbone(x)
+        if isinstance(x, dict):
+            y.update(x)
+        else:
+            y["backbone_out"] = x
+        final_name = "backbone_out"
+        if self.use_neck:
+            x = self.neck(x)
+            if isinstance(x, dict):
+                y.update(x)
+            else:
+                y["neck_out"] = x
+            final_name = "neck_out"
+        if self.use_head:
+            x = self.head(x, targets=data)
+            # for multi head, save ctc neck out for udml
+            if isinstance(x, dict) and "ctc_neck" in x.keys():
+                y["neck_out"] = x["ctc_neck"]
+                y["head_out"] = x
+            elif isinstance(x, dict):
+                y.update(x)
+            else:
+                y["head_out"] = x
+            final_name = "head_out"
+        if self.return_all_feats:
+            if self.training:
+                return y
+            elif isinstance(x, dict):
+                return x
+            else:
+                return {final_name: x}
+        else:
+            return x

ppocr/modeling/architectures/distillation_model.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2021 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.
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+from ppocr.modeling.transforms import build_transform
+from ppocr.modeling.backbones import build_backbone
+from ppocr.modeling.necks import build_neck
+from ppocr.modeling.heads import build_head
+from .base_model import BaseModel
+from ppocr.utils.save_load import load_pretrained_params
+
+__all__ = ["DistillationModel"]
+
+
+class DistillationModel(nn.Layer):
+    def __init__(self, config):
+        """
+        the module for OCR distillation.
+        args:
+            config (dict): the super parameters for module.
+        """
+        super().__init__()
+        self.model_list = []
+        self.model_name_list = []
+        for key in config["Models"]:
+            model_config = config["Models"][key]
+            freeze_params = False
+            pretrained = None
+            if "freeze_params" in model_config:
+                freeze_params = model_config.pop("freeze_params")
+            if "pretrained" in model_config:
+                pretrained = model_config.pop("pretrained")
+            model = BaseModel(model_config)
+            if pretrained is not None:
+                load_pretrained_params(model, pretrained)
+            if freeze_params:
+                for param in model.parameters():
+                    param.trainable = False
+            self.model_list.append(self.add_sublayer(key, model))
+            self.model_name_list.append(key)
+
+    def forward(self, x, data=None):
+        result_dict = dict()
+        for idx, model_name in enumerate(self.model_name_list):
+            result_dict[model_name] = self.model_list[idx](x, data)
+        return result_dict

ppocr/modeling/backbones/__init__.py

@@ -0,0 +1,152 @@
+# 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_backbone"]
+
+
+def build_backbone(config, model_type):
+    if model_type == "det" or model_type == "table":
+        from .det_mobilenet_v3 import MobileNetV3
+        from .det_resnet import ResNet
+        from .det_resnet_vd import ResNet_vd
+        from .det_resnet_vd_sast import ResNet_SAST
+        from .det_pp_lcnet import PPLCNet
+        from .rec_lcnetv3 import PPLCNetV3
+        from .rec_hgnet import PPHGNet_small
+        from .rec_vit import ViT
+        from .det_pp_lcnet_v2 import PPLCNetV2_base
+        from .rec_repvit import RepSVTR_det
+        from .rec_vary_vit import Vary_VIT_B
+        from .rec_pphgnetv2 import PPHGNetV2_B4
+
+        support_dict = [
+            "MobileNetV3",
+            "ResNet",
+            "ResNet_vd",
+            "ResNet_SAST",
+            "PPLCNet",
+            "PPLCNetV3",
+            "PPHGNet_small",
+            "PPLCNetV2_base",
+            "RepSVTR_det",
+            "Vary_VIT_B",
+            "PPHGNetV2_B4",
+        ]
+        if model_type == "table":
+            from .table_master_resnet import TableResNetExtra
+
+            support_dict.append("TableResNetExtra")
+    elif model_type == "rec" or model_type == "cls":
+        from .rec_mobilenet_v3 import MobileNetV3
+        from .rec_resnet_vd import ResNet
+        from .rec_resnet_fpn import ResNetFPN
+        from .rec_mv1_enhance import MobileNetV1Enhance
+        from .rec_nrtr_mtb import MTB
+        from .rec_resnet_31 import ResNet31
+        from .rec_resnet_32 import ResNet32
+        from .rec_resnet_45 import ResNet45
+        from .rec_resnet_aster import ResNet_ASTER
+        from .rec_micronet import MicroNet
+        from .rec_efficientb3_pren import EfficientNetb3_PREN
+        from .rec_svtrnet import SVTRNet
+        from .rec_vitstr import ViTSTR
+        from .rec_resnet_rfl import ResNetRFL
+        from .rec_densenet import DenseNet
+        from .rec_resnetv2 import ResNetV2
+        from .rec_hybridvit import HybridTransformer
+        from .rec_donut_swin import DonutSwinModel
+        from .rec_shallow_cnn import ShallowCNN
+        from .rec_lcnetv3 import PPLCNetV3
+        from .rec_hgnet import PPHGNet_small
+        from .rec_vit_parseq import ViTParseQ
+        from .rec_repvit import RepSVTR
+        from .rec_svtrv2 import SVTRv2
+        from .rec_vary_vit import Vary_VIT_B, Vary_VIT_B_Formula
+        from .rec_pphgnetv2 import (
+            PPHGNetV2_B4,
+            PPHGNetV2_B4_Formula,
+            PPHGNetV2_B6_Formula,
+        )
+
+        support_dict = [
+            "MobileNetV1Enhance",
+            "MobileNetV3",
+            "ResNet",
+            "ResNetFPN",
+            "MTB",
+            "ResNet31",
+            "ResNet45",
+            "ResNet_ASTER",
+            "MicroNet",
+            "EfficientNetb3_PREN",
+            "SVTRNet",
+            "ViTSTR",
+            "ResNet32",
+            "ResNetRFL",
+            "DenseNet",
+            "ShallowCNN",
+            "PPLCNetV3",
+            "PPHGNet_small",
+            "ViTParseQ",
+            "ViT",
+            "RepSVTR",
+            "SVTRv2",
+            "ResNetV2",
+            "HybridTransformer",
+            "DonutSwinModel",
+            "Vary_VIT_B",
+            "PPHGNetV2_B4",
+            "PPHGNetV2_B4_Formula",
+            "PPHGNetV2_B6_Formula",
+            "Vary_VIT_B_Formula",
+        ]
+    elif model_type == "e2e":
+        from .e2e_resnet_vd_pg import ResNet
+
+        support_dict = ["ResNet"]
+    elif model_type == "kie":
+        from .kie_unet_sdmgr import Kie_backbone
+        from .vqa_layoutlm import (
+            LayoutLMForSer,
+            LayoutLMv2ForSer,
+            LayoutLMv2ForRe,
+            LayoutXLMForSer,
+            LayoutXLMForRe,
+        )
+
+        support_dict = [
+            "Kie_backbone",
+            "LayoutLMForSer",
+            "LayoutLMv2ForSer",
+            "LayoutLMv2ForRe",
+            "LayoutXLMForSer",
+            "LayoutXLMForRe",
+        ]
+    elif model_type == "table":
+        from .table_resnet_vd import ResNet
+        from .table_mobilenet_v3 import MobileNetV3
+        from .rec_vary_vit import Vary_VIT_B
+
+        support_dict = ["ResNet", "MobileNetV3", "Vary_VIT_B"]
+    else:
+        raise NotImplementedError
+
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "when model typs is {}, backbone only support {}".format(
+            model_type, support_dict
+        )
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/modeling/backbones/det_mobilenet_v3.py

@@ -0,0 +1,289 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+from ppocr.modeling.backbones.rec_hgnet import MeanPool2D
+
+__all__ = ["MobileNetV3"]
+
+
+def make_divisible(v, divisor=8, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class MobileNetV3(nn.Layer):
+    def __init__(
+        self, in_channels=3, model_name="large", scale=0.5, disable_se=False, **kwargs
+    ):
+        """
+        the MobilenetV3 backbone network for detection module.
+        Args:
+            params(dict): the super parameters for build network
+        """
+        super(MobileNetV3, self).__init__()
+
+        self.disable_se = disable_se
+
+        if model_name == "large":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, False, "relu", 1],
+                [3, 64, 24, False, "relu", 2],
+                [3, 72, 24, False, "relu", 1],
+                [5, 72, 40, True, "relu", 2],
+                [5, 120, 40, True, "relu", 1],
+                [5, 120, 40, True, "relu", 1],
+                [3, 240, 80, False, "hardswish", 2],
+                [3, 200, 80, False, "hardswish", 1],
+                [3, 184, 80, False, "hardswish", 1],
+                [3, 184, 80, False, "hardswish", 1],
+                [3, 480, 112, True, "hardswish", 1],
+                [3, 672, 112, True, "hardswish", 1],
+                [5, 672, 160, True, "hardswish", 2],
+                [5, 960, 160, True, "hardswish", 1],
+                [5, 960, 160, True, "hardswish", 1],
+            ]
+            cls_ch_squeeze = 960
+        elif model_name == "small":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, True, "relu", 2],
+                [3, 72, 24, False, "relu", 2],
+                [3, 88, 24, False, "relu", 1],
+                [5, 96, 40, True, "hardswish", 2],
+                [5, 240, 40, True, "hardswish", 1],
+                [5, 240, 40, True, "hardswish", 1],
+                [5, 120, 48, True, "hardswish", 1],
+                [5, 144, 48, True, "hardswish", 1],
+                [5, 288, 96, True, "hardswish", 2],
+                [5, 576, 96, True, "hardswish", 1],
+                [5, 576, 96, True, "hardswish", 1],
+            ]
+            cls_ch_squeeze = 576
+        else:
+            raise NotImplementedError(
+                "mode[" + model_name + "_model] is not implemented!"
+            )
+
+        supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25]
+        assert (
+            scale in supported_scale
+        ), "supported scale are {} but input scale is {}".format(supported_scale, scale)
+        inplanes = 16
+        # conv1
+        self.conv = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=make_divisible(inplanes * scale),
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=1,
+            if_act=True,
+            act="hardswish",
+        )
+
+        self.stages = []
+        self.out_channels = []
+        block_list = []
+        i = 0
+        inplanes = make_divisible(inplanes * scale)
+        for k, exp, c, se, nl, s in cfg:
+            se = se and not self.disable_se
+            start_idx = 2 if model_name == "large" else 0
+            if s == 2 and i > start_idx:
+                self.out_channels.append(inplanes)
+                self.stages.append(nn.Sequential(*block_list))
+                block_list = []
+            block_list.append(
+                ResidualUnit(
+                    in_channels=inplanes,
+                    mid_channels=make_divisible(scale * exp),
+                    out_channels=make_divisible(scale * c),
+                    kernel_size=k,
+                    stride=s,
+                    use_se=se,
+                    act=nl,
+                )
+            )
+            inplanes = make_divisible(scale * c)
+            i += 1
+        block_list.append(
+            ConvBNLayer(
+                in_channels=inplanes,
+                out_channels=make_divisible(scale * cls_ch_squeeze),
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                groups=1,
+                if_act=True,
+                act="hardswish",
+            )
+        )
+        self.stages.append(nn.Sequential(*block_list))
+        self.out_channels.append(make_divisible(scale * cls_ch_squeeze))
+        for i, stage in enumerate(self.stages):
+            self.add_sublayer(sublayer=stage, name="stage{}".format(i))
+
+    def forward(self, x):
+        x = self.conv(x)
+        out_list = []
+        for stage in self.stages:
+            x = stage(x)
+            out_list.append(x)
+        return out_list
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups=1,
+        if_act=True,
+        act=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm(num_channels=out_channels, act=None)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.if_act:
+            if self.act == "relu":
+                x = F.relu(x)
+            elif self.act == "hardswish":
+                x = F.hardswish(x)
+            else:
+                print(
+                    "The activation function({}) is selected incorrectly.".format(
+                        self.act
+                    )
+                )
+                exit()
+        return x
+
+
+class ResidualUnit(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        mid_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        use_se,
+        act=None,
+    ):
+        super(ResidualUnit, self).__init__()
+        self.if_shortcut = stride == 1 and in_channels == out_channels
+        self.if_se = use_se
+
+        self.expand_conv = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=mid_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=True,
+            act=act,
+        )
+        self.bottleneck_conv = ConvBNLayer(
+            in_channels=mid_channels,
+            out_channels=mid_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=int((kernel_size - 1) // 2),
+            groups=mid_channels,
+            if_act=True,
+            act=act,
+        )
+        if self.if_se:
+            self.mid_se = SEModule(mid_channels)
+        self.linear_conv = ConvBNLayer(
+            in_channels=mid_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=False,
+            act=None,
+        )
+
+    def forward(self, inputs):
+        x = self.expand_conv(inputs)
+        x = self.bottleneck_conv(x)
+        if self.if_se:
+            x = self.mid_se(x)
+        x = self.linear_conv(x)
+        if self.if_shortcut:
+            x = paddle.add(inputs, x)
+        return x
+
+
+class SEModule(nn.Layer):
+    def __init__(self, in_channels, reduction=4):
+        super(SEModule, self).__init__()
+        if "npu" in paddle.device.get_device():
+            self.avg_pool = MeanPool2D(1, 1)
+        else:
+            self.avg_pool = nn.AdaptiveAvgPool2D(1)
+        self.conv1 = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=in_channels // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.conv2 = nn.Conv2D(
+            in_channels=in_channels // reduction,
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+    def forward(self, inputs):
+        outputs = self.avg_pool(inputs)
+        outputs = self.conv1(outputs)
+        outputs = F.relu(outputs)
+        outputs = self.conv2(outputs)
+        outputs = F.hardsigmoid(outputs, slope=0.2, offset=0.5)
+        return inputs * outputs

ppocr/modeling/backbones/det_pp_lcnet.py

@@ -0,0 +1,274 @@
+# copyright (c) 2021 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, division, print_function
+
+import os
+import paddle
+import paddle.nn as nn
+from paddle import ParamAttr
+from paddle.nn import AdaptiveAvgPool2D, BatchNorm, Conv2D, Dropout, Linear
+from paddle.regularizer import L2Decay
+from paddle.nn.initializer import KaimingNormal
+from paddle.utils.download import get_path_from_url
+
+MODEL_URLS = {
+    "PPLCNet_x0.25": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x0_25_pretrained.pdparams",
+    "PPLCNet_x0.35": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x0_35_pretrained.pdparams",
+    "PPLCNet_x0.5": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x0_5_pretrained.pdparams",
+    "PPLCNet_x0.75": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x0_75_pretrained.pdparams",
+    "PPLCNet_x1.0": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x1_0_pretrained.pdparams",
+    "PPLCNet_x1.5": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x1_5_pretrained.pdparams",
+    "PPLCNet_x2.0": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x2_0_pretrained.pdparams",
+    "PPLCNet_x2.5": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNet_x2_5_pretrained.pdparams",
+}
+
+MODEL_STAGES_PATTERN = {
+    "PPLCNet": ["blocks2", "blocks3", "blocks4", "blocks5", "blocks6"]
+}
+
+__all__ = list(MODEL_URLS.keys())
+
+# Each element(list) represents a depthwise block, which is composed of k, in_c, out_c, s, use_se.
+# k: kernel_size
+# in_c: input channel number in depthwise block
+# out_c: output channel number in depthwise block
+# s: stride in depthwise block
+# use_se: whether to use SE block
+
+NET_CONFIG = {
+    "blocks2":
+    # k, in_c, out_c, s, use_se
+    [[3, 16, 32, 1, False]],
+    "blocks3": [[3, 32, 64, 2, False], [3, 64, 64, 1, False]],
+    "blocks4": [[3, 64, 128, 2, False], [3, 128, 128, 1, False]],
+    "blocks5": [
+        [3, 128, 256, 2, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+    ],
+    "blocks6": [[5, 256, 512, 2, True], [5, 512, 512, 1, True]],
+}
+
+
+def make_divisible(v, divisor=8, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(self, num_channels, filter_size, num_filters, stride, num_groups=1):
+        super().__init__()
+
+        self.conv = Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=(filter_size - 1) // 2,
+            groups=num_groups,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+            bias_attr=False,
+        )
+
+        self.bn = BatchNorm(
+            num_filters,
+            param_attr=ParamAttr(regularizer=L2Decay(0.0)),
+            bias_attr=ParamAttr(regularizer=L2Decay(0.0)),
+        )
+        self.hardswish = nn.Hardswish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.hardswish(x)
+        return x
+
+
+class DepthwiseSeparable(nn.Layer):
+    def __init__(self, num_channels, num_filters, stride, dw_size=3, use_se=False):
+        super().__init__()
+        self.use_se = use_se
+        self.dw_conv = ConvBNLayer(
+            num_channels=num_channels,
+            num_filters=num_channels,
+            filter_size=dw_size,
+            stride=stride,
+            num_groups=num_channels,
+        )
+        if use_se:
+            self.se = SEModule(num_channels)
+        self.pw_conv = ConvBNLayer(
+            num_channels=num_channels, filter_size=1, num_filters=num_filters, stride=1
+        )
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        if self.use_se:
+            x = self.se(x)
+        x = self.pw_conv(x)
+        return x
+
+
+class SEModule(nn.Layer):
+    def __init__(self, channel, reduction=4):
+        super().__init__()
+        self.avg_pool = AdaptiveAvgPool2D(1)
+        self.conv1 = Conv2D(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.relu = nn.ReLU()
+        self.conv2 = Conv2D(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.hardsigmoid = nn.Hardsigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.hardsigmoid(x)
+        x = paddle.multiply(x=identity, y=x)
+        return x
+
+
+class PPLCNet(nn.Layer):
+    def __init__(self, in_channels=3, scale=1.0, pretrained=False, use_ssld=False):
+        super().__init__()
+        self.out_channels = [
+            int(NET_CONFIG["blocks3"][-1][2] * scale),
+            int(NET_CONFIG["blocks4"][-1][2] * scale),
+            int(NET_CONFIG["blocks5"][-1][2] * scale),
+            int(NET_CONFIG["blocks6"][-1][2] * scale),
+        ]
+        self.scale = scale
+
+        self.conv1 = ConvBNLayer(
+            num_channels=in_channels,
+            filter_size=3,
+            num_filters=make_divisible(16 * scale),
+            stride=2,
+        )
+
+        self.blocks2 = nn.Sequential(
+            *[
+                DepthwiseSeparable(
+                    num_channels=make_divisible(in_c * scale),
+                    num_filters=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(NET_CONFIG["blocks2"])
+            ]
+        )
+
+        self.blocks3 = nn.Sequential(
+            *[
+                DepthwiseSeparable(
+                    num_channels=make_divisible(in_c * scale),
+                    num_filters=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(NET_CONFIG["blocks3"])
+            ]
+        )
+
+        self.blocks4 = nn.Sequential(
+            *[
+                DepthwiseSeparable(
+                    num_channels=make_divisible(in_c * scale),
+                    num_filters=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(NET_CONFIG["blocks4"])
+            ]
+        )
+
+        self.blocks5 = nn.Sequential(
+            *[
+                DepthwiseSeparable(
+                    num_channels=make_divisible(in_c * scale),
+                    num_filters=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(NET_CONFIG["blocks5"])
+            ]
+        )
+
+        self.blocks6 = nn.Sequential(
+            *[
+                DepthwiseSeparable(
+                    num_channels=make_divisible(in_c * scale),
+                    num_filters=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(NET_CONFIG["blocks6"])
+            ]
+        )
+
+        if pretrained:
+            self._load_pretrained(
+                MODEL_URLS["PPLCNet_x{}".format(scale)], use_ssld=use_ssld
+            )
+
+    def forward(self, x):
+        outs = []
+        x = self.conv1(x)
+        x = self.blocks2(x)
+        x = self.blocks3(x)
+        outs.append(x)
+        x = self.blocks4(x)
+        outs.append(x)
+        x = self.blocks5(x)
+        outs.append(x)
+        x = self.blocks6(x)
+        outs.append(x)
+        return outs
+
+    def _load_pretrained(self, pretrained_url, use_ssld=False):
+        if use_ssld:
+            pretrained_url = pretrained_url.replace("_pretrained", "_ssld_pretrained")
+        print(pretrained_url)
+        local_weight_path = get_path_from_url(
+            pretrained_url, os.path.expanduser("~/.paddleclas/weights")
+        )
+        param_state_dict = paddle.load(local_weight_path)
+        self.set_dict(param_state_dict)
+        return

ppocr/modeling/backbones/det_pp_lcnet_v2.py

@@ -0,0 +1,358 @@
+# copyright (c) 2024 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, division, print_function
+import os
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+from paddle.nn import AdaptiveAvgPool2D, BatchNorm2D, Conv2D, Dropout, Linear
+from paddle.regularizer import L2Decay
+from paddle.nn.initializer import KaimingNormal
+from paddle.utils.download import get_path_from_url
+
+MODEL_URLS = {
+    "PPLCNetV2_small": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNetV2_small_ssld_pretrained.pdparams",
+    "PPLCNetV2_base": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNetV2_base_ssld_pretrained.pdparams",
+    "PPLCNetV2_large": "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/PPLCNetV2_large_ssld_pretrained.pdparams",
+}
+
+__all__ = list(MODEL_URLS.keys())
+
+NET_CONFIG = {
+    # in_channels, kernel_size, split_pw, use_rep, use_se, use_shortcut
+    "stage1": [64, 3, False, False, False, False],
+    "stage2": [128, 3, False, False, False, False],
+    "stage3": [256, 5, True, True, True, False],
+    "stage4": [512, 5, False, True, False, True],
+}
+
+
+def make_divisible(v, divisor=8, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, groups=1, use_act=True
+    ):
+        super().__init__()
+        self.use_act = use_act
+        self.conv = 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(initializer=KaimingNormal()),
+            bias_attr=False,
+        )
+
+        self.bn = BatchNorm2D(
+            out_channels,
+            weight_attr=ParamAttr(regularizer=L2Decay(0.0)),
+            bias_attr=ParamAttr(regularizer=L2Decay(0.0)),
+        )
+        if self.use_act:
+            self.act = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.use_act:
+            x = self.act(x)
+        return x
+
+
+class SEModule(nn.Layer):
+    def __init__(self, channel, reduction=4):
+        super().__init__()
+        self.avg_pool = AdaptiveAvgPool2D(1)
+        self.conv1 = Conv2D(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.relu = nn.ReLU()
+        self.conv2 = Conv2D(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.hardsigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.hardsigmoid(x)
+        x = paddle.multiply(x=identity, y=x)
+        return x
+
+
+class RepDepthwiseSeparable(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        dw_size=3,
+        split_pw=False,
+        use_rep=False,
+        use_se=False,
+        use_shortcut=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.is_repped = False
+
+        self.dw_size = dw_size
+        self.split_pw = split_pw
+        self.use_rep = use_rep
+        self.use_se = use_se
+        self.use_shortcut = (
+            True
+            if use_shortcut and stride == 1 and in_channels == out_channels
+            else False
+        )
+
+        if self.use_rep:
+            self.dw_conv_list = nn.LayerList()
+            for kernel_size in range(self.dw_size, 0, -2):
+                if kernel_size == 1 and stride != 1:
+                    continue
+                dw_conv = ConvBNLayer(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    groups=in_channels,
+                    use_act=False,
+                )
+                self.dw_conv_list.append(dw_conv)
+            self.dw_conv = nn.Conv2D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=dw_size,
+                stride=stride,
+                padding=(dw_size - 1) // 2,
+                groups=in_channels,
+            )
+        else:
+            self.dw_conv = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=dw_size,
+                stride=stride,
+                groups=in_channels,
+            )
+
+        self.act = nn.ReLU()
+
+        if use_se:
+            self.se = SEModule(in_channels)
+
+        if self.split_pw:
+            pw_ratio = 0.5
+            self.pw_conv_1 = ConvBNLayer(
+                in_channels=in_channels,
+                kernel_size=1,
+                out_channels=int(out_channels * pw_ratio),
+                stride=1,
+            )
+            self.pw_conv_2 = ConvBNLayer(
+                in_channels=int(out_channels * pw_ratio),
+                kernel_size=1,
+                out_channels=out_channels,
+                stride=1,
+            )
+        else:
+            self.pw_conv = ConvBNLayer(
+                in_channels=in_channels,
+                kernel_size=1,
+                out_channels=out_channels,
+                stride=1,
+            )
+
+    def forward(self, x):
+        if self.use_rep:
+            input_x = x
+            if self.is_repped:
+                x = self.act(self.dw_conv(x))
+            else:
+                y = self.dw_conv_list[0](x)
+                for dw_conv in self.dw_conv_list[1:]:
+                    y += dw_conv(x)
+                x = self.act(y)
+        else:
+            x = self.dw_conv(x)
+
+        if self.use_se:
+            x = self.se(x)
+        if self.split_pw:
+            x = self.pw_conv_1(x)
+            x = self.pw_conv_2(x)
+        else:
+            x = self.pw_conv(x)
+        if self.use_shortcut:
+            x = x + input_x
+        return x
+
+    def re_parameterize(self):
+        if self.use_rep:
+            self.is_repped = True
+            kernel, bias = self._get_equivalent_kernel_bias()
+            self.dw_conv.weight.set_value(kernel)
+            self.dw_conv.bias.set_value(bias)
+
+    def _get_equivalent_kernel_bias(self):
+        kernel_sum = 0
+        bias_sum = 0
+        for dw_conv in self.dw_conv_list:
+            kernel, bias = self._fuse_bn_tensor(dw_conv)
+            kernel = self._pad_tensor(kernel, to_size=self.dw_size)
+            kernel_sum += kernel
+            bias_sum += bias
+        return kernel_sum, bias_sum
+
+    def _fuse_bn_tensor(self, branch):
+        kernel = branch.conv.weight
+        running_mean = branch.bn._mean
+        running_var = branch.bn._variance
+        gamma = branch.bn.weight
+        beta = branch.bn.bias
+        eps = branch.bn._epsilon
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape((-1, 1, 1, 1))
+        return kernel * t, beta - running_mean * gamma / std
+
+    def _pad_tensor(self, tensor, to_size):
+        from_size = tensor.shape[-1]
+        if from_size == to_size:
+            return tensor
+        pad = (to_size - from_size) // 2
+        return F.pad(tensor, [pad, pad, pad, pad])
+
+
+class PPLCNetV2(nn.Layer):
+    def __init__(self, scale, depths, out_indx=[1, 2, 3, 4], **kwargs):
+        super().__init__(**kwargs)
+        self.scale = scale
+        self.out_channels = [
+            # int(NET_CONFIG["blocks3"][-1][2] * scale),
+            int(NET_CONFIG["stage1"][0] * scale * 2),
+            int(NET_CONFIG["stage2"][0] * scale * 2),
+            int(NET_CONFIG["stage3"][0] * scale * 2),
+            int(NET_CONFIG["stage4"][0] * scale * 2),
+        ]
+        self.stem = nn.Sequential(
+            *[
+                ConvBNLayer(
+                    in_channels=3,
+                    kernel_size=3,
+                    out_channels=make_divisible(32 * scale),
+                    stride=2,
+                ),
+                RepDepthwiseSeparable(
+                    in_channels=make_divisible(32 * scale),
+                    out_channels=make_divisible(64 * scale),
+                    stride=1,
+                    dw_size=3,
+                ),
+            ]
+        )
+        self.out_indx = out_indx
+        # stages
+        self.stages = nn.LayerList()
+        for depth_idx, k in enumerate(NET_CONFIG):
+            (
+                in_channels,
+                kernel_size,
+                split_pw,
+                use_rep,
+                use_se,
+                use_shortcut,
+            ) = NET_CONFIG[k]
+            self.stages.append(
+                nn.Sequential(
+                    *[
+                        RepDepthwiseSeparable(
+                            in_channels=make_divisible(
+                                (in_channels if i == 0 else in_channels * 2) * scale
+                            ),
+                            out_channels=make_divisible(in_channels * 2 * scale),
+                            stride=2 if i == 0 else 1,
+                            dw_size=kernel_size,
+                            split_pw=split_pw,
+                            use_rep=use_rep,
+                            use_se=use_se,
+                            use_shortcut=use_shortcut,
+                        )
+                        for i in range(depths[depth_idx])
+                    ]
+                )
+            )
+
+        # if pretrained:
+        self._load_pretrained(MODEL_URLS["PPLCNetV2_base"], use_ssld=True)
+
+    def forward(self, x):
+        x = self.stem(x)
+        i = 1
+        outs = []
+        for stage in self.stages:
+            x = stage(x)
+            if i in self.out_indx:
+                outs.append(x)
+            i += 1
+        return outs
+
+    def _load_pretrained(self, pretrained_url, use_ssld=False):
+        print(pretrained_url)
+        local_weight_path = get_path_from_url(
+            pretrained_url, os.path.expanduser("~/.paddleclas/weights")
+        )
+        param_state_dict = paddle.load(local_weight_path)
+        self.set_dict(param_state_dict)
+        print("load pretrain ssd success!")
+        return
+
+
+def PPLCNetV2_base(in_channels=3, **kwargs):
+    """
+    PPLCNetV2_base
+    Args:
+        pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise.
+                    If str, means the path of the pretrained model.
+        use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True.
+    Returns:
+        model: nn.Layer. Specific `PPLCNetV2_base` model depends on args.
+    """
+    model = PPLCNetV2(scale=1.0, depths=[2, 2, 6, 2], **kwargs)
+    return model

ppocr/modeling/backbones/det_resnet.py

@@ -0,0 +1,235 @@
+# 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 numpy as np
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+from paddle.nn.initializer import Uniform
+
+import math
+
+from paddle.vision.ops import DeformConv2D
+from paddle.regularizer import L2Decay
+from paddle.nn.initializer import Normal, Constant, XavierUniform
+from .det_resnet_vd import DeformableConvV2, ConvBNLayer
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(self, num_channels, num_filters, stride, shortcut=True, is_dcn=False):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=1,
+            act="relu",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=num_filters,
+            out_channels=num_filters,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            is_dcn=is_dcn,
+            dcn_groups=1,
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels=num_filters,
+            out_channels=num_filters * 4,
+            kernel_size=1,
+            act=None,
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=num_channels,
+                out_channels=num_filters * 4,
+                kernel_size=1,
+                stride=stride,
+            )
+
+        self.shortcut = shortcut
+
+        self._num_channels_out = num_filters * 4
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self, num_channels, num_filters, stride, shortcut=True, name=None):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=num_filters, out_channels=num_filters, kernel_size=3, act=None
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=num_channels,
+                out_channels=num_filters,
+                kernel_size=1,
+                stride=stride,
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet(nn.Layer):
+    def __init__(self, in_channels=3, layers=50, out_indices=None, dcn_stage=None):
+        super(ResNet, self).__init__()
+
+        self.layers = layers
+        self.input_image_channel = in_channels
+
+        supported_layers = [18, 34, 50, 101, 152]
+        assert (
+            layers in supported_layers
+        ), "supported layers are {} but input layer is {}".format(
+            supported_layers, layers
+        )
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            depth = [3, 4, 6, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        num_channels = [64, 256, 512, 1024] if layers >= 50 else [64, 64, 128, 256]
+        num_filters = [64, 128, 256, 512]
+
+        self.dcn_stage = (
+            dcn_stage if dcn_stage is not None else [False, False, False, False]
+        )
+        self.out_indices = out_indices if out_indices is not None else [0, 1, 2, 3]
+
+        self.conv = ConvBNLayer(
+            in_channels=self.input_image_channel,
+            out_channels=64,
+            kernel_size=7,
+            stride=2,
+            act="relu",
+        )
+        self.pool2d_max = MaxPool2D(
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+
+        self.stages = []
+        self.out_channels = []
+        if layers >= 50:
+            for block in range(len(depth)):
+                shortcut = False
+                block_list = []
+                is_dcn = self.dcn_stage[block]
+                for i in range(depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+                    bottleneck_block = self.add_sublayer(
+                        conv_name,
+                        BottleneckBlock(
+                            num_channels=(
+                                num_channels[block]
+                                if i == 0
+                                else num_filters[block] * 4
+                            ),
+                            num_filters=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            is_dcn=is_dcn,
+                        ),
+                    )
+                    block_list.append(bottleneck_block)
+                    shortcut = True
+                if block in self.out_indices:
+                    self.out_channels.append(num_filters[block] * 4)
+                self.stages.append(nn.Sequential(*block_list))
+        else:
+            for block in range(len(depth)):
+                shortcut = False
+                block_list = []
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    basic_block = self.add_sublayer(
+                        conv_name,
+                        BasicBlock(
+                            num_channels=(
+                                num_channels[block] if i == 0 else num_filters[block]
+                            ),
+                            num_filters=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                        ),
+                    )
+                    block_list.append(basic_block)
+                    shortcut = True
+                if block in self.out_indices:
+                    self.out_channels.append(num_filters[block])
+                self.stages.append(nn.Sequential(*block_list))
+
+    def forward(self, inputs):
+        y = self.conv(inputs)
+        y = self.pool2d_max(y)
+        out = []
+        for i, block in enumerate(self.stages):
+            y = block(y)
+            if i in self.out_indices:
+                out.append(y)
+        return out

ppocr/modeling/backbones/det_resnet_vd.py

@@ -0,0 +1,369 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from paddle.vision.ops import DeformConv2D
+from paddle.regularizer import L2Decay
+from paddle.nn.initializer import Normal, Constant, XavierUniform
+
+__all__ = ["ResNet_vd", "ConvBNLayer", "DeformableConvV2"]
+
+
+class DeformableConvV2(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        weight_attr=None,
+        bias_attr=None,
+        lr_scale=1,
+        regularizer=None,
+        skip_quant=False,
+        dcn_bias_regularizer=L2Decay(0.0),
+        dcn_bias_lr_scale=2.0,
+    ):
+        super(DeformableConvV2, self).__init__()
+        self.offset_channel = 2 * kernel_size**2 * groups
+        self.mask_channel = kernel_size**2 * groups
+
+        if bias_attr:
+            # in FCOS-DCN head, specifically need learning_rate and regularizer
+            dcn_bias_attr = ParamAttr(
+                initializer=Constant(value=0),
+                regularizer=dcn_bias_regularizer,
+                learning_rate=dcn_bias_lr_scale,
+            )
+        else:
+            # in ResNet backbone, do not need bias
+            dcn_bias_attr = False
+        self.conv_dcn = DeformConv2D(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2 * dilation,
+            dilation=dilation,
+            deformable_groups=groups,
+            weight_attr=weight_attr,
+            bias_attr=dcn_bias_attr,
+        )
+
+        if lr_scale == 1 and regularizer is None:
+            offset_bias_attr = ParamAttr(initializer=Constant(0.0))
+        else:
+            offset_bias_attr = ParamAttr(
+                initializer=Constant(0.0),
+                learning_rate=lr_scale,
+                regularizer=regularizer,
+            )
+        self.conv_offset = nn.Conv2D(
+            in_channels,
+            groups * 3 * kernel_size**2,
+            kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            weight_attr=ParamAttr(initializer=Constant(0.0)),
+            bias_attr=offset_bias_attr,
+        )
+        if skip_quant:
+            self.conv_offset.skip_quant = True
+
+    def forward(self, x):
+        offset_mask = self.conv_offset(x)
+        offset, mask = paddle.split(
+            offset_mask,
+            num_or_sections=[self.offset_channel, self.mask_channel],
+            axis=1,
+        )
+        mask = F.sigmoid(mask)
+        y = self.conv_dcn(x, offset, mask=mask)
+        return y
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        dcn_groups=1,
+        is_vd_mode=False,
+        act=None,
+        is_dcn=False,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True
+        )
+        if not is_dcn:
+            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,
+                bias_attr=False,
+            )
+        else:
+            self._conv = DeformableConvV2(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=(kernel_size - 1) // 2,
+                groups=dcn_groups,  # groups,
+                bias_attr=False,
+            )
+        self._batch_norm = nn.BatchNorm(out_channels, act=act)
+
+    def forward(self, inputs):
+        if self.is_vd_mode:
+            inputs = self._pool2d_avg(inputs)
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        is_dcn=False,
+    ):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act="relu",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            is_dcn=is_dcn,
+            dcn_groups=2,
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+    ):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels, out_channels=out_channels, kernel_size=3, act=None
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet_vd(nn.Layer):
+    def __init__(
+        self, in_channels=3, layers=50, dcn_stage=None, out_indices=None, **kwargs
+    ):
+        super(ResNet_vd, self).__init__()
+
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert (
+            layers in supported_layers
+        ), "supported layers are {} but input layer is {}".format(
+            supported_layers, layers
+        )
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            depth = [3, 4, 6, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        num_channels = [64, 256, 512, 1024] if layers >= 50 else [64, 64, 128, 256]
+        num_filters = [64, 128, 256, 512]
+
+        self.dcn_stage = (
+            dcn_stage if dcn_stage is not None else [False, False, False, False]
+        )
+        self.out_indices = out_indices if out_indices is not None else [0, 1, 2, 3]
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=32,
+            kernel_size=3,
+            stride=2,
+            act="relu",
+        )
+        self.conv1_2 = ConvBNLayer(
+            in_channels=32, out_channels=32, kernel_size=3, stride=1, act="relu"
+        )
+        self.conv1_3 = ConvBNLayer(
+            in_channels=32, out_channels=64, kernel_size=3, stride=1, act="relu"
+        )
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        self.stages = []
+        self.out_channels = []
+        if layers >= 50:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                is_dcn = self.dcn_stage[block]
+                for i in range(depth[block]):
+                    bottleneck_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BottleneckBlock(
+                            in_channels=(
+                                num_channels[block]
+                                if i == 0
+                                else num_filters[block] * 4
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            is_dcn=is_dcn,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(bottleneck_block)
+                if block in self.out_indices:
+                    self.out_channels.append(num_filters[block] * 4)
+                self.stages.append(nn.Sequential(*block_list))
+        else:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    basic_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BasicBlock(
+                            in_channels=(
+                                num_channels[block] if i == 0 else num_filters[block]
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(basic_block)
+                if block in self.out_indices:
+                    self.out_channels.append(num_filters[block])
+                self.stages.append(nn.Sequential(*block_list))
+
+    def forward(self, inputs):
+        y = self.conv1_1(inputs)
+        y = self.conv1_2(y)
+        y = self.conv1_3(y)
+        y = self.pool2d_max(y)
+        out = []
+        for i, block in enumerate(self.stages):
+            y = block(y)
+            if i in self.out_indices:
+                out.append(y)
+        return out

ppocr/modeling/backbones/det_resnet_vd_sast.py

@@ -0,0 +1,314 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+__all__ = ["ResNet_SAST"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        is_vd_mode=False,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True
+        )
+        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,
+        )
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = 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, inputs):
+        if self.is_vd_mode:
+            inputs = self._pool2d_avg(inputs)
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2b",
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+            name=name + "_branch2c",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None,
+            name=name + "_branch2b",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet_SAST(nn.Layer):
+    def __init__(self, in_channels=3, layers=50, **kwargs):
+        super(ResNet_SAST, self).__init__()
+
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert (
+            layers in supported_layers
+        ), "supported layers are {} but input layer is {}".format(
+            supported_layers, layers
+        )
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            # depth = [3, 4, 6, 3]
+            depth = [3, 4, 6, 3, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        # num_channels = [64, 256, 512,
+        #                 1024] if layers >= 50 else [64, 64, 128, 256]
+        # num_filters = [64, 128, 256, 512]
+        num_channels = (
+            [64, 256, 512, 1024, 2048] if layers >= 50 else [64, 64, 128, 256]
+        )
+        num_filters = [64, 128, 256, 512, 512]
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=32,
+            kernel_size=3,
+            stride=2,
+            act="relu",
+            name="conv1_1",
+        )
+        self.conv1_2 = ConvBNLayer(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv1_2",
+        )
+        self.conv1_3 = ConvBNLayer(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv1_3",
+        )
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        self.stages = []
+        self.out_channels = [3, 64]
+        if layers >= 50:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+                    bottleneck_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BottleneckBlock(
+                            in_channels=(
+                                num_channels[block]
+                                if i == 0
+                                else num_filters[block] * 4
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(bottleneck_block)
+                self.out_channels.append(num_filters[block] * 4)
+                self.stages.append(nn.Sequential(*block_list))
+        else:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    basic_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BasicBlock(
+                            in_channels=(
+                                num_channels[block] if i == 0 else num_filters[block]
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(basic_block)
+                self.out_channels.append(num_filters[block])
+                self.stages.append(nn.Sequential(*block_list))
+
+    def forward(self, inputs):
+        out = [inputs]
+        y = self.conv1_1(inputs)
+        y = self.conv1_2(y)
+        y = self.conv1_3(y)
+        out.append(y)
+        y = self.pool2d_max(y)
+        for block in self.stages:
+            y = block(y)
+            out.append(y)
+        return out

ppocr/modeling/backbones/e2e_resnet_vd_pg.py

@@ -0,0 +1,292 @@
+# copyright (c) 2021 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 paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+__all__ = ["ResNet"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        is_vd_mode=False,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True
+        )
+        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,
+        )
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = 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, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2b",
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+            name=name + "_branch2c",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=stride,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None,
+            name=name + "_branch2b",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet(nn.Layer):
+    def __init__(self, in_channels=3, layers=50, **kwargs):
+        super(ResNet, self).__init__()
+
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert (
+            layers in supported_layers
+        ), "supported layers are {} but input layer is {}".format(
+            supported_layers, layers
+        )
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            # depth = [3, 4, 6, 3]
+            depth = [3, 4, 6, 3, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        num_channels = (
+            [64, 256, 512, 1024, 2048] if layers >= 50 else [64, 64, 128, 256]
+        )
+        num_filters = [64, 128, 256, 512, 512]
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=7,
+            stride=2,
+            act="relu",
+            name="conv1_1",
+        )
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        self.stages = []
+        self.out_channels = [3, 64]
+        # num_filters = [64, 128, 256, 512, 512]
+        if layers >= 50:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+                    bottleneck_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BottleneckBlock(
+                            in_channels=(
+                                num_channels[block]
+                                if i == 0
+                                else num_filters[block] * 4
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(bottleneck_block)
+                self.out_channels.append(num_filters[block] * 4)
+                self.stages.append(nn.Sequential(*block_list))
+        else:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    basic_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BasicBlock(
+                            in_channels=(
+                                num_channels[block] if i == 0 else num_filters[block]
+                            ),
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    block_list.append(basic_block)
+                self.out_channels.append(num_filters[block])
+                self.stages.append(nn.Sequential(*block_list))
+
+    def forward(self, inputs):
+        out = [inputs]
+        y = self.conv1_1(inputs)
+        out.append(y)
+        y = self.pool2d_max(y)
+        for block in self.stages:
+            y = block(y)
+            out.append(y)
+        return out

ppocr/modeling/backbones/kie_unet_sdmgr.py

@@ -0,0 +1,199 @@
+# copyright (c) 2021 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 paddle
+from paddle import nn
+import numpy as np
+import cv2
+
+__all__ = ["Kie_backbone"]
+
+
+class Encoder(nn.Layer):
+    def __init__(self, num_channels, num_filters):
+        super(Encoder, self).__init__()
+        self.conv1 = nn.Conv2D(
+            num_channels,
+            num_filters,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm(num_filters, act="relu")
+
+        self.conv2 = nn.Conv2D(
+            num_filters,
+            num_filters,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn2 = nn.BatchNorm(num_filters, act="relu")
+
+        self.pool = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, inputs):
+        x = self.conv1(inputs)
+        x = self.bn1(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x_pooled = self.pool(x)
+        return x, x_pooled
+
+
+class Decoder(nn.Layer):
+    def __init__(self, num_channels, num_filters):
+        super(Decoder, self).__init__()
+
+        self.conv1 = nn.Conv2D(
+            num_channels,
+            num_filters,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm(num_filters, act="relu")
+
+        self.conv2 = nn.Conv2D(
+            num_filters,
+            num_filters,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn2 = nn.BatchNorm(num_filters, act="relu")
+
+        self.conv0 = nn.Conv2D(
+            num_channels,
+            num_filters,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias_attr=False,
+        )
+        self.bn0 = nn.BatchNorm(num_filters, act="relu")
+
+    def forward(self, inputs_prev, inputs):
+        x = self.conv0(inputs)
+        x = self.bn0(x)
+        x = paddle.nn.functional.interpolate(
+            x, scale_factor=2, mode="bilinear", align_corners=False
+        )
+        x = paddle.concat([inputs_prev, x], axis=1)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        return x
+
+
+class UNet(nn.Layer):
+    def __init__(self):
+        super(UNet, self).__init__()
+        self.down1 = Encoder(num_channels=3, num_filters=16)
+        self.down2 = Encoder(num_channels=16, num_filters=32)
+        self.down3 = Encoder(num_channels=32, num_filters=64)
+        self.down4 = Encoder(num_channels=64, num_filters=128)
+        self.down5 = Encoder(num_channels=128, num_filters=256)
+
+        self.up1 = Decoder(32, 16)
+        self.up2 = Decoder(64, 32)
+        self.up3 = Decoder(128, 64)
+        self.up4 = Decoder(256, 128)
+        self.out_channels = 16
+
+    def forward(self, inputs):
+        x1, _ = self.down1(inputs)
+        _, x2 = self.down2(x1)
+        _, x3 = self.down3(x2)
+        _, x4 = self.down4(x3)
+        _, x5 = self.down5(x4)
+
+        x = self.up4(x4, x5)
+        x = self.up3(x3, x)
+        x = self.up2(x2, x)
+        x = self.up1(x1, x)
+        return x
+
+
+class Kie_backbone(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super(Kie_backbone, self).__init__()
+        self.out_channels = 16
+        self.img_feat = UNet()
+        self.maxpool = nn.MaxPool2D(kernel_size=7)
+
+    def bbox2roi(self, bbox_list):
+        rois_list = []
+        rois_num = []
+        for img_id, bboxes in enumerate(bbox_list):
+            rois_num.append(bboxes.shape[0])
+            rois_list.append(bboxes)
+        rois = paddle.concat(rois_list, 0)
+        rois_num = paddle.to_tensor(rois_num, dtype="int32")
+        return rois, rois_num
+
+    def pre_process(self, img, relations, texts, gt_bboxes, tag, img_size):
+        img, relations, texts, gt_bboxes, tag, img_size = (
+            img.numpy(),
+            relations.numpy(),
+            texts.numpy(),
+            gt_bboxes.numpy(),
+            tag.numpy().tolist(),
+            img_size.numpy(),
+        )
+        temp_relations, temp_texts, temp_gt_bboxes = [], [], []
+        h, w = int(np.max(img_size[:, 0])), int(np.max(img_size[:, 1]))
+        img = paddle.to_tensor(img[:, :, :h, :w])
+        batch = len(tag)
+        for i in range(batch):
+            num, recoder_len = tag[i][0], tag[i][1]
+            temp_relations.append(
+                paddle.to_tensor(relations[i, :num, :num, :], dtype="float32")
+            )
+            temp_texts.append(
+                paddle.to_tensor(texts[i, :num, :recoder_len], dtype="float32")
+            )
+            temp_gt_bboxes.append(
+                paddle.to_tensor(gt_bboxes[i, :num, ...], dtype="float32")
+            )
+        return img, temp_relations, temp_texts, temp_gt_bboxes
+
+    def forward(self, inputs):
+        img = inputs[0]
+        relations, texts, gt_bboxes, tag, img_size = (
+            inputs[1],
+            inputs[2],
+            inputs[3],
+            inputs[5],
+            inputs[-1],
+        )
+        img, relations, texts, gt_bboxes = self.pre_process(
+            img, relations, texts, gt_bboxes, tag, img_size
+        )
+        x = self.img_feat(img)
+        boxes, rois_num = self.bbox2roi(gt_bboxes)
+        feats = paddle.vision.ops.roi_align(
+            x, boxes, spatial_scale=1.0, output_size=7, boxes_num=rois_num
+        )
+        feats = self.maxpool(feats).squeeze(-1).squeeze(-1)
+        return [relations, texts, feats]

ppocr/modeling/backbones/rec_densenet.py

@@ -0,0 +1,150 @@
+# 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/LBH1024/CAN/models/densenet.py
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class Bottleneck(nn.Layer):
+    def __init__(self, nChannels, growthRate, use_dropout):
+        super(Bottleneck, self).__init__()
+        interChannels = 4 * growthRate
+        self.bn1 = nn.BatchNorm2D(interChannels)
+        self.conv1 = nn.Conv2D(
+            nChannels, interChannels, kernel_size=1, bias_attr=None
+        )  # Xavier initialization
+        self.bn2 = nn.BatchNorm2D(growthRate)
+        self.conv2 = nn.Conv2D(
+            interChannels, growthRate, kernel_size=3, padding=1, bias_attr=None
+        )  # Xavier initialization
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(p=0.2)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        if self.use_dropout:
+            out = self.dropout(out)
+        out = F.relu(self.bn2(self.conv2(out)))
+        if self.use_dropout:
+            out = self.dropout(out)
+        out = paddle.concat([x, out], 1)
+        return out
+
+
+class SingleLayer(nn.Layer):
+    def __init__(self, nChannels, growthRate, use_dropout):
+        super(SingleLayer, self).__init__()
+        self.bn1 = nn.BatchNorm2D(nChannels)
+        self.conv1 = nn.Conv2D(
+            nChannels, growthRate, kernel_size=3, padding=1, bias_attr=False
+        )
+
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(p=0.2)
+
+    def forward(self, x):
+        out = self.conv1(F.relu(x))
+        if self.use_dropout:
+            out = self.dropout(out)
+
+        out = paddle.concat([x, out], 1)
+        return out
+
+
+class Transition(nn.Layer):
+    def __init__(self, nChannels, out_channels, use_dropout):
+        super(Transition, self).__init__()
+        self.bn1 = nn.BatchNorm2D(out_channels)
+        self.conv1 = nn.Conv2D(nChannels, out_channels, kernel_size=1, bias_attr=False)
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(p=0.2)
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        if self.use_dropout:
+            out = self.dropout(out)
+        out = F.avg_pool2d(out, 2, ceil_mode=True, exclusive=False)
+        return out
+
+
+class DenseNet(nn.Layer):
+    def __init__(
+        self, growthRate, reduction, bottleneck, use_dropout, input_channel, **kwargs
+    ):
+        super(DenseNet, self).__init__()
+
+        nDenseBlocks = 16
+        nChannels = 2 * growthRate
+
+        self.conv1 = nn.Conv2D(
+            input_channel,
+            nChannels,
+            kernel_size=7,
+            padding=3,
+            stride=2,
+            bias_attr=False,
+        )
+        self.dense1 = self._make_dense(
+            nChannels, growthRate, nDenseBlocks, bottleneck, use_dropout
+        )
+        nChannels += nDenseBlocks * growthRate
+        out_channels = int(math.floor(nChannels * reduction))
+        self.trans1 = Transition(nChannels, out_channels, use_dropout)
+
+        nChannels = out_channels
+        self.dense2 = self._make_dense(
+            nChannels, growthRate, nDenseBlocks, bottleneck, use_dropout
+        )
+        nChannels += nDenseBlocks * growthRate
+        out_channels = int(math.floor(nChannels * reduction))
+        self.trans2 = Transition(nChannels, out_channels, use_dropout)
+
+        nChannels = out_channels
+        self.dense3 = self._make_dense(
+            nChannels, growthRate, nDenseBlocks, bottleneck, use_dropout
+        )
+        self.out_channels = out_channels
+
+    def _make_dense(self, nChannels, growthRate, nDenseBlocks, bottleneck, use_dropout):
+        layers = []
+        for i in range(int(nDenseBlocks)):
+            if bottleneck:
+                layers.append(Bottleneck(nChannels, growthRate, use_dropout))
+            else:
+                layers.append(SingleLayer(nChannels, growthRate, use_dropout))
+            nChannels += growthRate
+        return nn.Sequential(*layers)
+
+    def forward(self, inputs):
+        x, x_m, y = inputs
+        out = self.conv1(x)
+        out = F.relu(out)
+        out = F.max_pool2d(out, 2, ceil_mode=True)
+        out = self.dense1(out)
+        out = self.trans1(out)
+        out = self.dense2(out)
+        out = self.trans2(out)
+        out = self.dense3(out)
+        return out, x_m, y

ppocr/modeling/backbones/rec_efficientb3_pren.py

@@ -0,0 +1,305 @@
+# 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.
+"""
+Code is refer from:
+https://github.com/RuijieJ/pren/blob/main/Nets/EfficientNet.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import re
+import collections
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+__all__ = ["EfficientNetb3_PREN"]
+
+GlobalParams = collections.namedtuple(
+    "GlobalParams",
+    [
+        "batch_norm_momentum",
+        "batch_norm_epsilon",
+        "dropout_rate",
+        "num_classes",
+        "width_coefficient",
+        "depth_coefficient",
+        "depth_divisor",
+        "min_depth",
+        "drop_connect_rate",
+        "image_size",
+    ],
+)
+
+BlockArgs = collections.namedtuple(
+    "BlockArgs",
+    [
+        "kernel_size",
+        "num_repeat",
+        "input_filters",
+        "output_filters",
+        "expand_ratio",
+        "id_skip",
+        "stride",
+        "se_ratio",
+    ],
+)
+
+
+class BlockDecoder:
+    @staticmethod
+    def _decode_block_string(block_string):
+        assert isinstance(block_string, str)
+
+        ops = block_string.split("_")
+        options = {}
+        for op in ops:
+            splits = re.split(r"(\d.*)", op)
+            if len(splits) >= 2:
+                key, value = splits[:2]
+                options[key] = value
+
+        assert ("s" in options and len(options["s"]) == 1) or (
+            len(options["s"]) == 2 and options["s"][0] == options["s"][1]
+        )
+
+        return BlockArgs(
+            kernel_size=int(options["k"]),
+            num_repeat=int(options["r"]),
+            input_filters=int(options["i"]),
+            output_filters=int(options["o"]),
+            expand_ratio=int(options["e"]),
+            id_skip=("noskip" not in block_string),
+            se_ratio=float(options["se"]) if "se" in options else None,
+            stride=[int(options["s"][0])],
+        )
+
+    @staticmethod
+    def decode(string_list):
+        assert isinstance(string_list, list)
+        blocks_args = []
+        for block_string in string_list:
+            blocks_args.append(BlockDecoder._decode_block_string(block_string))
+        return blocks_args
+
+
+def efficientnet(
+    width_coefficient=None,
+    depth_coefficient=None,
+    dropout_rate=0.2,
+    drop_connect_rate=0.2,
+    image_size=None,
+    num_classes=1000,
+):
+    blocks_args = [
+        "r1_k3_s11_e1_i32_o16_se0.25",
+        "r2_k3_s22_e6_i16_o24_se0.25",
+        "r2_k5_s22_e6_i24_o40_se0.25",
+        "r3_k3_s22_e6_i40_o80_se0.25",
+        "r3_k5_s11_e6_i80_o112_se0.25",
+        "r4_k5_s22_e6_i112_o192_se0.25",
+        "r1_k3_s11_e6_i192_o320_se0.25",
+    ]
+    blocks_args = BlockDecoder.decode(blocks_args)
+
+    global_params = GlobalParams(
+        batch_norm_momentum=0.99,
+        batch_norm_epsilon=1e-3,
+        dropout_rate=dropout_rate,
+        drop_connect_rate=drop_connect_rate,
+        num_classes=num_classes,
+        width_coefficient=width_coefficient,
+        depth_coefficient=depth_coefficient,
+        depth_divisor=8,
+        min_depth=None,
+        image_size=image_size,
+    )
+    return blocks_args, global_params
+
+
+class EffUtils:
+    @staticmethod
+    def round_filters(filters, global_params):
+        """Calculate and round number of filters based on depth multiplier."""
+        multiplier = global_params.width_coefficient
+        if not multiplier:
+            return filters
+        divisor = global_params.depth_divisor
+        min_depth = global_params.min_depth
+        filters *= multiplier
+        min_depth = min_depth or divisor
+        new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)
+        if new_filters < 0.9 * filters:
+            new_filters += divisor
+        return int(new_filters)
+
+    @staticmethod
+    def round_repeats(repeats, global_params):
+        """Round number of filters based on depth multiplier."""
+        multiplier = global_params.depth_coefficient
+        if not multiplier:
+            return repeats
+        return int(math.ceil(multiplier * repeats))
+
+
+class MbConvBlock(nn.Layer):
+    def __init__(self, block_args):
+        super(MbConvBlock, self).__init__()
+        self._block_args = block_args
+        self.has_se = (self._block_args.se_ratio is not None) and (
+            0 < self._block_args.se_ratio <= 1
+        )
+        self.id_skip = block_args.id_skip
+
+        # expansion phase
+        self.inp = self._block_args.input_filters
+        oup = self._block_args.input_filters * self._block_args.expand_ratio
+        if self._block_args.expand_ratio != 1:
+            self._expand_conv = nn.Conv2D(self.inp, oup, 1, bias_attr=False)
+            self._bn0 = nn.BatchNorm(oup)
+
+        # depthwise conv phase
+        k = self._block_args.kernel_size
+        s = self._block_args.stride
+        if isinstance(s, list):
+            s = s[0]
+        self._depthwise_conv = nn.Conv2D(
+            oup,
+            oup,
+            groups=oup,
+            kernel_size=k,
+            stride=s,
+            padding="same",
+            bias_attr=False,
+        )
+        self._bn1 = nn.BatchNorm(oup)
+
+        # squeeze and excitation layer, if desired
+        if self.has_se:
+            num_squeezed_channels = max(
+                1, int(self._block_args.input_filters * self._block_args.se_ratio)
+            )
+            self._se_reduce = nn.Conv2D(oup, num_squeezed_channels, 1)
+            self._se_expand = nn.Conv2D(num_squeezed_channels, oup, 1)
+
+        # output phase and some util class
+        self.final_oup = self._block_args.output_filters
+        self._project_conv = nn.Conv2D(oup, self.final_oup, 1, bias_attr=False)
+        self._bn2 = nn.BatchNorm(self.final_oup)
+        self._swish = nn.Swish()
+
+    def _drop_connect(self, inputs, p, training):
+        if not training:
+            return inputs
+        batch_size = inputs.shape[0]
+        keep_prob = 1 - p
+        random_tensor = keep_prob
+        random_tensor += paddle.rand([batch_size, 1, 1, 1], dtype=inputs.dtype)
+        random_tensor = paddle.to_tensor(random_tensor, place=inputs.place)
+        binary_tensor = paddle.floor(random_tensor)
+        output = inputs / keep_prob * binary_tensor
+        return output
+
+    def forward(self, inputs, drop_connect_rate=None):
+        # expansion and depthwise conv
+        x = inputs
+        if self._block_args.expand_ratio != 1:
+            x = self._swish(self._bn0(self._expand_conv(inputs)))
+        x = self._swish(self._bn1(self._depthwise_conv(x)))
+
+        # squeeze and excitation
+        if self.has_se:
+            x_squeezed = F.adaptive_avg_pool2d(x, 1)
+            x_squeezed = self._se_expand(self._swish(self._se_reduce(x_squeezed)))
+            x = F.sigmoid(x_squeezed) * x
+        x = self._bn2(self._project_conv(x))
+
+        # skip connection and drop connect
+        if self.id_skip and self._block_args.stride == 1 and self.inp == self.final_oup:
+            if drop_connect_rate:
+                x = self._drop_connect(x, p=drop_connect_rate, training=self.training)
+            x = x + inputs
+        return x
+
+
+class EfficientNetb3_PREN(nn.Layer):
+    def __init__(self, in_channels):
+        super(EfficientNetb3_PREN, self).__init__()
+        """
+        the fllowing are efficientnetb3's superparams,
+        they means efficientnetb3 network's width, depth, resolution and
+        dropout respectively, to fit for text recognition task, the resolution
+        here is changed from 300 to 64.
+        """
+        w, d, s, p = 1.2, 1.4, 64, 0.3
+        self._blocks_args, self._global_params = efficientnet(
+            width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s
+        )
+        self.out_channels = []
+        # stem
+        out_channels = EffUtils.round_filters(32, self._global_params)
+        self._conv_stem = nn.Conv2D(
+            in_channels, out_channels, 3, 2, padding="same", bias_attr=False
+        )
+        self._bn0 = nn.BatchNorm(out_channels)
+
+        # build blocks
+        self._blocks = []
+        # to extract three feature maps for fpn based on efficientnetb3 backbone
+        self._concerned_block_idxes = [7, 17, 25]
+        _concerned_idx = 0
+        for i, block_args in enumerate(self._blocks_args):
+            block_args = block_args._replace(
+                input_filters=EffUtils.round_filters(
+                    block_args.input_filters, self._global_params
+                ),
+                output_filters=EffUtils.round_filters(
+                    block_args.output_filters, self._global_params
+                ),
+                num_repeat=EffUtils.round_repeats(
+                    block_args.num_repeat, self._global_params
+                ),
+            )
+            self._blocks.append(self.add_sublayer(f"{i}-0", MbConvBlock(block_args)))
+            _concerned_idx += 1
+            if _concerned_idx in self._concerned_block_idxes:
+                self.out_channels.append(block_args.output_filters)
+            if block_args.num_repeat > 1:
+                block_args = block_args._replace(
+                    input_filters=block_args.output_filters, stride=1
+                )
+            for j in range(block_args.num_repeat - 1):
+                self._blocks.append(
+                    self.add_sublayer(f"{i}-{j+1}", MbConvBlock(block_args))
+                )
+                _concerned_idx += 1
+                if _concerned_idx in self._concerned_block_idxes:
+                    self.out_channels.append(block_args.output_filters)
+
+        self._swish = nn.Swish()
+
+    def forward(self, inputs):
+        outs = []
+        x = self._swish(self._bn0(self._conv_stem(inputs)))
+        for idx, block in enumerate(self._blocks):
+            drop_connect_rate = self._global_params.drop_connect_rate
+            if drop_connect_rate:
+                drop_connect_rate *= float(idx) / len(self._blocks)
+            x = block(x, drop_connect_rate=drop_connect_rate)
+            if idx in self._concerned_block_idxes:
+                outs.append(x)
+        return outs

ppocr/modeling/backbones/rec_hgnet.py

@@ -0,0 +1,385 @@
+# 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.
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn.initializer import KaimingNormal, Constant
+from paddle.nn import Conv2D, BatchNorm2D, ReLU, AdaptiveAvgPool2D, MaxPool2D
+from paddle.regularizer import L2Decay
+from paddle import ParamAttr
+
+kaiming_normal_ = KaimingNormal()
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+class MeanPool2D(nn.Layer):
+    def __init__(self, w, h):
+        super().__init__()
+        self.w = w
+        self.h = h
+
+    def forward(self, feat):
+        batch_size, channels, _, _ = feat.shape
+        feat_flat = paddle.reshape(feat, [batch_size, channels, -1])
+        feat_mean = paddle.mean(feat_flat, axis=2)
+        feat_mean = paddle.reshape(feat_mean, [batch_size, channels, self.w, self.h])
+        return feat_mean
+
+
+class ConvBNAct(nn.Layer):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, groups=1, use_act=True
+    ):
+        super().__init__()
+        self.use_act = use_act
+        self.conv = Conv2D(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            bias_attr=False,
+        )
+        self.bn = BatchNorm2D(
+            out_channels,
+            weight_attr=ParamAttr(regularizer=L2Decay(0.0)),
+            bias_attr=ParamAttr(regularizer=L2Decay(0.0)),
+        )
+        if self.use_act:
+            self.act = ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.use_act:
+            x = self.act(x)
+        return x
+
+
+class ESEModule(nn.Layer):
+    def __init__(self, channels):
+        super().__init__()
+        if "npu" in paddle.device.get_device():
+            self.avg_pool = MeanPool2D(1, 1)
+        else:
+            self.avg_pool = AdaptiveAvgPool2D(1)
+        self.conv = Conv2D(
+            in_channels=channels,
+            out_channels=channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv(x)
+        x = self.sigmoid(x)
+        return paddle.multiply(x=identity, y=x)
+
+
+class HG_Block(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        mid_channels,
+        out_channels,
+        layer_num,
+        identity=False,
+    ):
+        super().__init__()
+        self.identity = identity
+
+        self.layers = nn.LayerList()
+        self.layers.append(
+            ConvBNAct(
+                in_channels=in_channels,
+                out_channels=mid_channels,
+                kernel_size=3,
+                stride=1,
+            )
+        )
+        for _ in range(layer_num - 1):
+            self.layers.append(
+                ConvBNAct(
+                    in_channels=mid_channels,
+                    out_channels=mid_channels,
+                    kernel_size=3,
+                    stride=1,
+                )
+            )
+
+        # feature aggregation
+        total_channels = in_channels + layer_num * mid_channels
+        self.aggregation_conv = ConvBNAct(
+            in_channels=total_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+        )
+        self.att = ESEModule(out_channels)
+
+    def forward(self, x):
+        identity = x
+        output = []
+        output.append(x)
+        for layer in self.layers:
+            x = layer(x)
+            output.append(x)
+        x = paddle.concat(output, axis=1)
+        x = self.aggregation_conv(x)
+        x = self.att(x)
+        if self.identity:
+            x += identity
+        return x
+
+
+class HG_Stage(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        mid_channels,
+        out_channels,
+        block_num,
+        layer_num,
+        downsample=True,
+        stride=[2, 1],
+    ):
+        super().__init__()
+        self.downsample = downsample
+        if downsample:
+            self.downsample = ConvBNAct(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=3,
+                stride=stride,
+                groups=in_channels,
+                use_act=False,
+            )
+
+        blocks_list = []
+        blocks_list.append(
+            HG_Block(in_channels, mid_channels, out_channels, layer_num, identity=False)
+        )
+        for _ in range(block_num - 1):
+            blocks_list.append(
+                HG_Block(
+                    out_channels, mid_channels, out_channels, layer_num, identity=True
+                )
+            )
+        self.blocks = nn.Sequential(*blocks_list)
+
+    def forward(self, x):
+        if self.downsample:
+            x = self.downsample(x)
+        x = self.blocks(x)
+        return x
+
+
+class PPHGNet(nn.Layer):
+    """
+    PPHGNet
+    Args:
+        stem_channels: list. Stem channel list of PPHGNet.
+        stage_config: dict. The configuration of each stage of PPHGNet. such as the number of channels, stride, etc.
+        layer_num: int. Number of layers of HG_Block.
+        use_last_conv: boolean. Whether to use a 1x1 convolutional layer before the classification layer.
+        class_expand: int=2048. Number of channels for the last 1x1 convolutional layer.
+        dropout_prob: float. Parameters of dropout, 0.0 means dropout is not used.
+        class_num: int=1000. The number of classes.
+    Returns:
+        model: nn.Layer. Specific PPHGNet model depends on args.
+    """
+
+    def __init__(
+        self,
+        stem_channels,
+        stage_config,
+        layer_num,
+        in_channels=3,
+        det=False,
+        out_indices=None,
+    ):
+        super().__init__()
+        self.det = det
+        self.out_indices = out_indices if out_indices is not None else [0, 1, 2, 3]
+
+        # stem
+        stem_channels.insert(0, in_channels)
+        self.stem = nn.Sequential(
+            *[
+                ConvBNAct(
+                    in_channels=stem_channels[i],
+                    out_channels=stem_channels[i + 1],
+                    kernel_size=3,
+                    stride=2 if i == 0 else 1,
+                )
+                for i in range(len(stem_channels) - 1)
+            ]
+        )
+
+        if self.det:
+            self.pool = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+        # stages
+        self.stages = nn.LayerList()
+        self.out_channels = []
+        for block_id, k in enumerate(stage_config):
+            (
+                in_channels,
+                mid_channels,
+                out_channels,
+                block_num,
+                downsample,
+                stride,
+            ) = stage_config[k]
+            self.stages.append(
+                HG_Stage(
+                    in_channels,
+                    mid_channels,
+                    out_channels,
+                    block_num,
+                    layer_num,
+                    downsample,
+                    stride,
+                )
+            )
+            if block_id in self.out_indices:
+                self.out_channels.append(out_channels)
+
+        if not self.det:
+            self.out_channels = stage_config["stage4"][2]
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for m in self.sublayers():
+            if isinstance(m, nn.Conv2D):
+                kaiming_normal_(m.weight)
+            elif isinstance(m, (nn.BatchNorm2D)):
+                ones_(m.weight)
+                zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                zeros_(m.bias)
+
+    def forward(self, x):
+        x = self.stem(x)
+        if self.det:
+            x = self.pool(x)
+
+        out = []
+        for i, stage in enumerate(self.stages):
+            x = stage(x)
+            if self.det and i in self.out_indices:
+                out.append(x)
+        if self.det:
+            return out
+
+        if self.training:
+            x = F.adaptive_avg_pool2d(x, [1, 40])
+        else:
+            x = F.avg_pool2d(x, [3, 2])
+        return x
+
+
+def PPHGNet_tiny(pretrained=False, use_ssld=False, **kwargs):
+    """
+    PPHGNet_tiny
+    Args:
+        pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise.
+                    If str, means the path of the pretrained model.
+        use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True.
+    Returns:
+        model: nn.Layer. Specific `PPHGNet_tiny` model depends on args.
+    """
+    stage_config = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [96, 96, 224, 1, False, [2, 1]],
+        "stage2": [224, 128, 448, 1, True, [1, 2]],
+        "stage3": [448, 160, 512, 2, True, [2, 1]],
+        "stage4": [512, 192, 768, 1, True, [2, 1]],
+    }
+
+    model = PPHGNet(
+        stem_channels=[48, 48, 96], stage_config=stage_config, layer_num=5, **kwargs
+    )
+    return model
+
+
+def PPHGNet_small(pretrained=False, use_ssld=False, det=False, **kwargs):
+    """
+    PPHGNet_small
+    Args:
+        pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise.
+                    If str, means the path of the pretrained model.
+        use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True.
+    Returns:
+        model: nn.Layer. Specific `PPHGNet_small` model depends on args.
+    """
+    stage_config_det = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [128, 128, 256, 1, False, 2],
+        "stage2": [256, 160, 512, 1, True, 2],
+        "stage3": [512, 192, 768, 2, True, 2],
+        "stage4": [768, 224, 1024, 1, True, 2],
+    }
+
+    stage_config_rec = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [128, 128, 256, 1, True, [2, 1]],
+        "stage2": [256, 160, 512, 1, True, [1, 2]],
+        "stage3": [512, 192, 768, 2, True, [2, 1]],
+        "stage4": [768, 224, 1024, 1, True, [2, 1]],
+    }
+
+    model = PPHGNet(
+        stem_channels=[64, 64, 128],
+        stage_config=stage_config_det if det else stage_config_rec,
+        layer_num=6,
+        det=det,
+        **kwargs,
+    )
+    return model
+
+
+def PPHGNet_base(pretrained=False, use_ssld=True, **kwargs):
+    """
+    PPHGNet_base
+    Args:
+        pretrained: bool=False or str. If `True` load pretrained parameters, `False` otherwise.
+                    If str, means the path of the pretrained model.
+        use_ssld: bool=False. Whether using distillation pretrained model when pretrained=True.
+    Returns:
+        model: nn.Layer. Specific `PPHGNet_base` model depends on args.
+    """
+    stage_config = {
+        # in_channels, mid_channels, out_channels, blocks, downsample
+        "stage1": [160, 192, 320, 1, False, [2, 1]],
+        "stage2": [320, 224, 640, 2, True, [1, 2]],
+        "stage3": [640, 256, 960, 3, True, [2, 1]],
+        "stage4": [960, 288, 1280, 2, True, [2, 1]],
+    }
+
+    model = PPHGNet(
+        stem_channels=[96, 96, 160],
+        stage_config=stage_config,
+        layer_num=7,
+        dropout_prob=0.2,
+        **kwargs,
+    )
+    return model

ppocr/modeling/backbones/rec_hybridvit.py

@@ -0,0 +1,529 @@
+# copyright (c) 2024 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/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer_hybrid.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from itertools import repeat
+import collections
+import math
+from functools import partial
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from ppocr.modeling.backbones.rec_resnetv2 import (
+    ResNetV2,
+    StdConv2dSame,
+    DropPath,
+    get_padding,
+)
+from paddle.nn.initializer import (
+    TruncatedNormal,
+    Constant,
+    Normal,
+    KaimingUniform,
+    XavierUniform,
+)
+
+normal_ = Normal(mean=0.0, std=1e-6)
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+kaiming_normal_ = KaimingUniform(nonlinearity="relu")
+trunc_normal_ = TruncatedNormal(std=0.02)
+xavier_uniform_ = XavierUniform()
+
+
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = _ntuple
+
+
+class Conv2dAlign(nn.Conv2D):
+    """Conv2d with Weight Standardization. Used for BiT ResNet-V2 models.
+
+    Paper: `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization` -
+        https://arxiv.org/abs/1903.10520v2
+    """
+
+    def __init__(
+        self,
+        in_channel,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bias=True,
+        eps=1e-6,
+    ):
+
+        super().__init__(
+            in_channel,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias_attr=bias,
+            weight_attr=True,
+        )
+        self.eps = eps
+
+    def forward(self, x):
+        x = F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self._stride,
+            self._padding,
+            self._dilation,
+            self._groups,
+        )
+        return x
+
+
+class HybridEmbed(nn.Layer):
+    """CNN Feature Map Embedding
+    Extract feature map from CNN, flatten, project to embedding dim.
+    """
+
+    def __init__(
+        self,
+        backbone,
+        img_size=224,
+        patch_size=1,
+        feature_size=None,
+        in_chans=3,
+        embed_dim=768,
+    ):
+        super().__init__()
+        assert isinstance(backbone, nn.Layer)
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.backbone = backbone
+        feature_dim = 1024
+        feature_size = (42, 12)
+        patch_size = (1, 1)
+        assert (
+            feature_size[0] % patch_size[0] == 0
+            and feature_size[1] % patch_size[1] == 0
+        )
+
+        self.grid_size = (
+            feature_size[0] // patch_size[0],
+            feature_size[1] // patch_size[1],
+        )
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.proj = nn.Conv2D(
+            feature_dim,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            weight_attr=True,
+            bias_attr=True,
+        )
+
+    def forward(self, x):
+
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+        x = self.proj(x).flatten(2).transpose([0, 2, 1])
+
+        return x
+
+
+class myLinear(nn.Linear):
+    def __init__(self, in_channel, out_channels, weight_attr=True, bias_attr=True):
+        super().__init__(
+            in_channel, out_channels, weight_attr=weight_attr, bias_attr=bias_attr
+        )
+
+    def forward(self, x):
+        return paddle.matmul(x, self.weight, transpose_y=True) + self.bias
+
+
+class Attention(nn.Layer):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = myLinear(dim, dim, weight_attr=True, bias_attr=True)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape([B, N, 3, self.num_heads, C // self.num_heads])
+            .transpose([2, 0, 3, 1, 4])
+        )
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose([0, 1, 3, 2])) * self.scale
+
+        attn = F.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose([0, 2, 1, 3]).reshape([B, N, C])
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Mlp(nn.Layer):
+    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        drop_probs = to_2tuple(drop)
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.drop1 = nn.Dropout(drop_probs[0])
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop2 = nn.Dropout(drop_probs[1])
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop1(x)
+        x = self.fc2(x)
+        x = self.drop2(x)
+        return x
+
+
+class Block(nn.Layer):
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class HybridTransformer(nn.Layer):
+    """Implementation of HybridTransformer.
+
+    Args:
+      x: input images with shape [N, 1, H, W]
+      label: LaTeX-OCR labels with shape [N, L] , L is the max sequence length
+      attention_mask: LaTeX-OCR attention mask with shape [N, L]  , L is the max sequence length
+
+    Returns:
+      The encoded features with shape [N, 1, H//16, W//16]
+    """
+
+    def __init__(
+        self,
+        backbone_layers=[2, 3, 7],
+        input_channel=1,
+        is_predict=False,
+        is_export=False,
+        img_size=(224, 224),
+        patch_size=16,
+        num_classes=1000,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        representation_size=None,
+        distilled=False,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        embed_layer=None,
+        norm_layer=None,
+        act_layer=None,
+        weight_init="",
+        **kwargs,
+    ):
+        super(HybridTransformer, self).__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
+        self.num_tokens = 2 if distilled else 1
+        norm_layer = norm_layer or partial(nn.LayerNorm, epsilon=1e-6)
+        act_layer = act_layer or nn.GELU
+        self.height, self.width = img_size
+        self.patch_size = patch_size
+        backbone = ResNetV2(
+            layers=backbone_layers,
+            num_classes=0,
+            global_pool="",
+            in_chans=input_channel,
+            preact=False,
+            stem_type="same",
+            conv_layer=StdConv2dSame,
+            is_export=is_export,
+        )
+        min_patch_size = 2 ** (len(backbone_layers) + 1)
+        self.patch_embed = HybridEmbed(
+            img_size=img_size,
+            patch_size=patch_size // min_patch_size,
+            in_chans=input_channel,
+            embed_dim=embed_dim,
+            backbone=backbone,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = paddle.create_parameter([1, 1, embed_dim], dtype="float32")
+        self.dist_token = (
+            paddle.create_parameter(
+                [1, 1, embed_dim],
+                dtype="float32",
+            )
+            if distilled
+            else None
+        )
+        self.pos_embed = paddle.create_parameter(
+            [1, num_patches + self.num_tokens, embed_dim], dtype="float32"
+        )
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        zeros_(self.cls_token)
+        if self.dist_token is not None:
+            zeros_(self.dist_token)
+        zeros_(self.pos_embed)
+
+        dpr = [
+            x.item() for x in paddle.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.Sequential(
+            *[
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=act_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        # Representation layer
+        if representation_size and not distilled:
+            self.num_features = representation_size
+            self.pre_logits = nn.Sequential(
+                ("fc", nn.Linear(embed_dim, representation_size)), ("act", nn.Tanh())
+            )
+        else:
+            self.pre_logits = nn.Identity()
+
+        # Classifier head(s)
+        self.head = (
+            nn.Linear(self.num_features, num_classes)
+            if num_classes > 0
+            else nn.Identity()
+        )
+        self.head_dist = None
+        if distilled:
+            self.head_dist = (
+                nn.Linear(self.embed_dim, self.num_classes)
+                if num_classes > 0
+                else nn.Identity()
+            )
+        self.init_weights(weight_init)
+        self.out_channels = embed_dim
+        self.is_predict = is_predict
+        self.is_export = is_export
+
+    def init_weights(self, mode=""):
+        assert mode in ("jax", "jax_nlhb", "nlhb", "")
+        head_bias = -math.log(self.num_classes) if "nlhb" in mode else 0.0
+        trunc_normal_(self.pos_embed)
+        trunc_normal_(self.cls_token)
+        self.apply(_init_vit_weights)
+
+    def _init_weights(self, m):
+        # this fn left here for compat with downstream users
+        _init_vit_weights(m)
+
+    def load_pretrained(self, checkpoint_path, prefix=""):
+        raise NotImplementedError
+
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token", "dist_token"}
+
+    def get_classifier(self):
+        if self.dist_token is None:
+            return self.head
+        else:
+            return self.head, self.head_dist
+
+    def reset_classifier(self, num_classes, global_pool=""):
+        self.num_classes = num_classes
+        self.head = (
+            nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
+        if self.num_tokens == 2:
+            self.head_dist = (
+                nn.Linear(self.embed_dim, self.num_classes)
+                if num_classes > 0
+                else nn.Identity()
+            )
+
+    def forward_features(self, x):
+        B, c, h, w = x.shape
+        x = self.patch_embed(x)
+        cls_tokens = self.cls_token.expand(
+            [B, -1, -1]
+        )  # stole cls_tokens impl from Phil Wang, thanks
+        x = paddle.concat((cls_tokens, x), axis=1)
+        h, w = h // self.patch_size, w // self.patch_size
+        repeat_tensor = (
+            paddle.arange(h) * (self.width // self.patch_size - w)
+        ).reshape([-1, 1])
+        repeat_tensor = paddle.repeat_interleave(
+            repeat_tensor, paddle.to_tensor(w), axis=1
+        ).reshape([-1])
+        pos_emb_ind = repeat_tensor + paddle.arange(h * w)
+        pos_emb_ind = paddle.concat(
+            (paddle.zeros([1], dtype="int64"), pos_emb_ind + 1), axis=0
+        ).cast(paddle.int64)
+        x += self.pos_embed[:, pos_emb_ind]
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.norm(x)
+        return x
+
+    def forward(self, input_data):
+
+        if self.training:
+            x, label, attention_mask = input_data
+        else:
+            if isinstance(input_data, list):
+                x = input_data[0]
+            else:
+                x = input_data
+        x = self.forward_features(x)
+        x = self.head(x)
+        if self.training:
+            return x, label, attention_mask
+        else:
+            return x
+
+
+def _init_vit_weights(
+    module: nn.Layer, name: str = "", head_bias: float = 0.0, jax_impl: bool = False
+):
+    """ViT weight initialization
+    * When called without n, head_bias, jax_impl args it will behave exactly the same
+      as my original init for compatibility with prev hparam / downstream use cases (ie DeiT).
+    * When called w/ valid n (module name) and jax_impl=True, will (hopefully) match JAX impl
+    """
+    if isinstance(module, nn.Linear):
+        if name.startswith("head"):
+            zeros_(module.weight)
+            constant_ = Constant(value=head_bias)
+            constant_(module.bias, head_bias)
+        elif name.startswith("pre_logits"):
+            zeros_(module.bias)
+        else:
+            if jax_impl:
+                xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    if "mlp" in name:
+                        normal_(module.bias)
+                    else:
+                        zeros_(module.bias)
+            else:
+                trunc_normal_(module.weight)
+                if module.bias is not None:
+                    zeros_(module.bias)
+    elif jax_impl and isinstance(module, nn.Conv2D):
+        # NOTE conv was left to pytorch default in my original init
+        if module.bias is not None:
+            zeros_(module.bias)
+    elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2D)):
+        zeros_(module.bias)
+        ones_(module.weight)

ppocr/modeling/backbones/rec_lcnetv3.py

@@ -0,0 +1,558 @@
+# copyright (c) 2021 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 paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+from paddle.nn.initializer import Constant, KaimingNormal
+from paddle.nn import (
+    AdaptiveAvgPool2D,
+    BatchNorm2D,
+    Conv2D,
+    Dropout,
+    Hardsigmoid,
+    Hardswish,
+    Identity,
+    Linear,
+    ReLU,
+)
+from paddle.regularizer import L2Decay
+from ppocr.modeling.backbones.rec_hgnet import MeanPool2D
+
+NET_CONFIG_det = {
+    "blocks2":
+    # k, in_c, out_c, s, use_se
+    [[3, 16, 32, 1, False]],
+    "blocks3": [[3, 32, 64, 2, False], [3, 64, 64, 1, False]],
+    "blocks4": [[3, 64, 128, 2, False], [3, 128, 128, 1, False]],
+    "blocks5": [
+        [3, 128, 256, 2, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+    ],
+    "blocks6": [
+        [5, 256, 512, 2, True],
+        [5, 512, 512, 1, True],
+        [5, 512, 512, 1, False],
+        [5, 512, 512, 1, False],
+    ],
+}
+
+NET_CONFIG_rec = {
+    "blocks2":
+    # k, in_c, out_c, s, use_se
+    [[3, 16, 32, 1, False]],
+    "blocks3": [[3, 32, 64, 1, False], [3, 64, 64, 1, False]],
+    "blocks4": [[3, 64, 128, (2, 1), False], [3, 128, 128, 1, False]],
+    "blocks5": [
+        [3, 128, 256, (1, 2), False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+    ],
+    "blocks6": [
+        [5, 256, 512, (2, 1), True],
+        [5, 512, 512, 1, True],
+        [5, 512, 512, (2, 1), False],
+        [5, 512, 512, 1, False],
+    ],
+}
+
+
+def make_divisible(v, divisor=16, min_value=None):
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class LearnableAffineBlock(nn.Layer):
+    def __init__(self, scale_value=1.0, bias_value=0.0, lr_mult=1.0, lab_lr=0.1):
+        super().__init__()
+        self.scale = self.create_parameter(
+            shape=[
+                1,
+            ],
+            default_initializer=Constant(value=scale_value),
+            attr=ParamAttr(learning_rate=lr_mult * lab_lr),
+        )
+        self.add_parameter("scale", self.scale)
+        self.bias = self.create_parameter(
+            shape=[
+                1,
+            ],
+            default_initializer=Constant(value=bias_value),
+            attr=ParamAttr(learning_rate=lr_mult * lab_lr),
+        )
+        self.add_parameter("bias", self.bias)
+
+    def forward(self, x):
+        return self.scale * x + self.bias
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, groups=1, lr_mult=1.0
+    ):
+        super().__init__()
+        self.conv = 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(initializer=KaimingNormal(), learning_rate=lr_mult),
+            bias_attr=False,
+        )
+
+        self.bn = BatchNorm2D(
+            out_channels,
+            weight_attr=ParamAttr(regularizer=L2Decay(0.0), learning_rate=lr_mult),
+            bias_attr=ParamAttr(regularizer=L2Decay(0.0), learning_rate=lr_mult),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class Act(nn.Layer):
+    def __init__(self, act="hswish", lr_mult=1.0, lab_lr=0.1):
+        super().__init__()
+        if act == "hswish":
+            self.act = Hardswish()
+        else:
+            assert act == "relu"
+            self.act = ReLU()
+        self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr)
+
+    def forward(self, x):
+        return self.lab(self.act(x))
+
+
+class LearnableRepLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        num_conv_branches=1,
+        lr_mult=1.0,
+        lab_lr=0.1,
+    ):
+        super().__init__()
+        self.is_repped = False
+        self.groups = groups
+        self.stride = stride
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_conv_branches = num_conv_branches
+        self.padding = (kernel_size - 1) // 2
+
+        self.identity = (
+            BatchNorm2D(
+                num_features=in_channels,
+                weight_attr=ParamAttr(learning_rate=lr_mult),
+                bias_attr=ParamAttr(learning_rate=lr_mult),
+            )
+            if out_channels == in_channels and stride == 1
+            else None
+        )
+
+        self.conv_kxk = nn.LayerList(
+            [
+                ConvBNLayer(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride,
+                    groups=groups,
+                    lr_mult=lr_mult,
+                )
+                for _ in range(self.num_conv_branches)
+            ]
+        )
+
+        self.conv_1x1 = (
+            ConvBNLayer(
+                in_channels, out_channels, 1, stride, groups=groups, lr_mult=lr_mult
+            )
+            if kernel_size > 1
+            else None
+        )
+
+        self.lab = LearnableAffineBlock(lr_mult=lr_mult, lab_lr=lab_lr)
+        self.act = Act(lr_mult=lr_mult, lab_lr=lab_lr)
+
+    def forward(self, x):
+        # for export
+        if self.is_repped:
+            out = self.lab(self.reparam_conv(x))
+            if self.stride != 2:
+                out = self.act(out)
+            return out
+
+        out = 0
+        if self.identity is not None:
+            out += self.identity(x)
+
+        if self.conv_1x1 is not None:
+            out += self.conv_1x1(x)
+
+        for conv in self.conv_kxk:
+            out += conv(x)
+
+        out = self.lab(out)
+        if self.stride != 2:
+            out = self.act(out)
+        return out
+
+    def rep(self):
+        if self.is_repped:
+            return
+        kernel, bias = self._get_kernel_bias()
+        self.reparam_conv = Conv2D(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=self.kernel_size,
+            stride=self.stride,
+            padding=self.padding,
+            groups=self.groups,
+        )
+        self.reparam_conv.weight.set_value(kernel)
+        self.reparam_conv.bias.set_value(bias)
+        self.is_repped = True
+
+    def _pad_kernel_1x1_to_kxk(self, kernel1x1, pad):
+        if not isinstance(kernel1x1, paddle.Tensor):
+            return 0
+        else:
+            return nn.functional.pad(kernel1x1, [pad, pad, pad, pad])
+
+    def _get_kernel_bias(self):
+        kernel_conv_1x1, bias_conv_1x1 = self._fuse_bn_tensor(self.conv_1x1)
+        kernel_conv_1x1 = self._pad_kernel_1x1_to_kxk(
+            kernel_conv_1x1, self.kernel_size // 2
+        )
+
+        kernel_identity, bias_identity = self._fuse_bn_tensor(self.identity)
+
+        kernel_conv_kxk = 0
+        bias_conv_kxk = 0
+        for conv in self.conv_kxk:
+            kernel, bias = self._fuse_bn_tensor(conv)
+            kernel_conv_kxk += kernel
+            bias_conv_kxk += bias
+
+        kernel_reparam = kernel_conv_kxk + kernel_conv_1x1 + kernel_identity
+        bias_reparam = bias_conv_kxk + bias_conv_1x1 + bias_identity
+        return kernel_reparam, bias_reparam
+
+    def _fuse_bn_tensor(self, branch):
+        if not branch:
+            return 0, 0
+        elif isinstance(branch, ConvBNLayer):
+            kernel = branch.conv.weight
+            running_mean = branch.bn._mean
+            running_var = branch.bn._variance
+            gamma = branch.bn.weight
+            beta = branch.bn.bias
+            eps = branch.bn._epsilon
+        else:
+            assert isinstance(branch, BatchNorm2D)
+            if not hasattr(self, "id_tensor"):
+                input_dim = self.in_channels // self.groups
+                kernel_value = paddle.zeros(
+                    (self.in_channels, input_dim, self.kernel_size, self.kernel_size),
+                    dtype=branch.weight.dtype,
+                )
+                for i in range(self.in_channels):
+                    kernel_value[
+                        i, i % input_dim, self.kernel_size // 2, self.kernel_size // 2
+                    ] = 1
+                self.id_tensor = kernel_value
+            kernel = self.id_tensor
+            running_mean = branch._mean
+            running_var = branch._variance
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch._epsilon
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape((-1, 1, 1, 1))
+        return kernel * t, beta - running_mean * gamma / std
+
+
+class SELayer(nn.Layer):
+    def __init__(self, channel, reduction=4, lr_mult=1.0):
+        super().__init__()
+        if "npu" in paddle.device.get_device():
+            self.avg_pool = MeanPool2D(1, 1)
+        else:
+            self.avg_pool = AdaptiveAvgPool2D(1)
+        self.conv1 = Conv2D(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            weight_attr=ParamAttr(learning_rate=lr_mult),
+            bias_attr=ParamAttr(learning_rate=lr_mult),
+        )
+        self.relu = ReLU()
+        self.conv2 = Conv2D(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            weight_attr=ParamAttr(learning_rate=lr_mult),
+            bias_attr=ParamAttr(learning_rate=lr_mult),
+        )
+        self.hardsigmoid = Hardsigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.hardsigmoid(x)
+        x = paddle.multiply(x=identity, y=x)
+        return x
+
+
+class LCNetV3Block(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        dw_size,
+        use_se=False,
+        conv_kxk_num=4,
+        lr_mult=1.0,
+        lab_lr=0.1,
+    ):
+        super().__init__()
+        self.use_se = use_se
+        self.dw_conv = LearnableRepLayer(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=dw_size,
+            stride=stride,
+            groups=in_channels,
+            num_conv_branches=conv_kxk_num,
+            lr_mult=lr_mult,
+            lab_lr=lab_lr,
+        )
+        if use_se:
+            self.se = SELayer(in_channels, lr_mult=lr_mult)
+        self.pw_conv = LearnableRepLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            num_conv_branches=conv_kxk_num,
+            lr_mult=lr_mult,
+            lab_lr=lab_lr,
+        )
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        if self.use_se:
+            x = self.se(x)
+        x = self.pw_conv(x)
+        return x
+
+
+class PPLCNetV3(nn.Layer):
+    def __init__(
+        self,
+        scale=1.0,
+        conv_kxk_num=4,
+        lr_mult_list=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        lab_lr=0.1,
+        det=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.scale = scale
+        self.lr_mult_list = lr_mult_list
+        self.det = det
+
+        self.net_config = NET_CONFIG_det if self.det else NET_CONFIG_rec
+
+        assert isinstance(
+            self.lr_mult_list, (list, tuple)
+        ), "lr_mult_list should be in (list, tuple) but got {}".format(
+            type(self.lr_mult_list)
+        )
+        assert (
+            len(self.lr_mult_list) == 6
+        ), "lr_mult_list length should be 6 but got {}".format(len(self.lr_mult_list))
+
+        self.conv1 = ConvBNLayer(
+            in_channels=3,
+            out_channels=make_divisible(16 * scale),
+            kernel_size=3,
+            stride=2,
+            lr_mult=self.lr_mult_list[0],
+        )
+
+        self.blocks2 = nn.Sequential(
+            *[
+                LCNetV3Block(
+                    in_channels=make_divisible(in_c * scale),
+                    out_channels=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                    conv_kxk_num=conv_kxk_num,
+                    lr_mult=self.lr_mult_list[1],
+                    lab_lr=lab_lr,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks2"])
+            ]
+        )
+
+        self.blocks3 = nn.Sequential(
+            *[
+                LCNetV3Block(
+                    in_channels=make_divisible(in_c * scale),
+                    out_channels=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                    conv_kxk_num=conv_kxk_num,
+                    lr_mult=self.lr_mult_list[2],
+                    lab_lr=lab_lr,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks3"])
+            ]
+        )
+
+        self.blocks4 = nn.Sequential(
+            *[
+                LCNetV3Block(
+                    in_channels=make_divisible(in_c * scale),
+                    out_channels=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                    conv_kxk_num=conv_kxk_num,
+                    lr_mult=self.lr_mult_list[3],
+                    lab_lr=lab_lr,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks4"])
+            ]
+        )
+
+        self.blocks5 = nn.Sequential(
+            *[
+                LCNetV3Block(
+                    in_channels=make_divisible(in_c * scale),
+                    out_channels=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                    conv_kxk_num=conv_kxk_num,
+                    lr_mult=self.lr_mult_list[4],
+                    lab_lr=lab_lr,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks5"])
+            ]
+        )
+
+        self.blocks6 = nn.Sequential(
+            *[
+                LCNetV3Block(
+                    in_channels=make_divisible(in_c * scale),
+                    out_channels=make_divisible(out_c * scale),
+                    dw_size=k,
+                    stride=s,
+                    use_se=se,
+                    conv_kxk_num=conv_kxk_num,
+                    lr_mult=self.lr_mult_list[5],
+                    lab_lr=lab_lr,
+                )
+                for i, (k, in_c, out_c, s, se) in enumerate(self.net_config["blocks6"])
+            ]
+        )
+        self.out_channels = make_divisible(512 * scale)
+
+        if self.det:
+            mv_c = [16, 24, 56, 480]
+            self.out_channels = [
+                make_divisible(self.net_config["blocks3"][-1][2] * scale),
+                make_divisible(self.net_config["blocks4"][-1][2] * scale),
+                make_divisible(self.net_config["blocks5"][-1][2] * scale),
+                make_divisible(self.net_config["blocks6"][-1][2] * scale),
+            ]
+
+            self.layer_list = nn.LayerList(
+                [
+                    nn.Conv2D(self.out_channels[0], int(mv_c[0] * scale), 1, 1, 0),
+                    nn.Conv2D(self.out_channels[1], int(mv_c[1] * scale), 1, 1, 0),
+                    nn.Conv2D(self.out_channels[2], int(mv_c[2] * scale), 1, 1, 0),
+                    nn.Conv2D(self.out_channels[3], int(mv_c[3] * scale), 1, 1, 0),
+                ]
+            )
+            self.out_channels = [
+                int(mv_c[0] * scale),
+                int(mv_c[1] * scale),
+                int(mv_c[2] * scale),
+                int(mv_c[3] * scale),
+            ]
+
+    def forward(self, x):
+        out_list = []
+        x = self.conv1(x)
+
+        x = self.blocks2(x)
+        x = self.blocks3(x)
+        out_list.append(x)
+        x = self.blocks4(x)
+        out_list.append(x)
+        x = self.blocks5(x)
+        out_list.append(x)
+        x = self.blocks6(x)
+        out_list.append(x)
+
+        if self.det:
+            out_list[0] = self.layer_list[0](out_list[0])
+            out_list[1] = self.layer_list[1](out_list[1])
+            out_list[2] = self.layer_list[2](out_list[2])
+            out_list[3] = self.layer_list[3](out_list[3])
+            return out_list
+
+        if self.training:
+            x = F.adaptive_avg_pool2d(x, [1, 40])
+        else:
+            x = F.avg_pool2d(x, [3, 2])
+        return x

ppocr/modeling/backbones/rec_micronet.py

@@ -0,0 +1,605 @@
+# copyright (c) 2021 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/liyunsheng13/micronet/blob/main/backbone/micronet.py
+https://github.com/liyunsheng13/micronet/blob/main/backbone/activation.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+import paddle.nn as nn
+
+from ppocr.modeling.backbones.det_mobilenet_v3 import make_divisible
+
+M0_cfgs = [
+    # s, n, c, ks, c1, c2, g1, g2, c3, g3, g4, y1, y2, y3, r
+    [2, 1, 8, 3, 2, 2, 0, 4, 8, 2, 2, 2, 0, 1, 1],
+    [2, 1, 12, 3, 2, 2, 0, 8, 12, 4, 4, 2, 2, 1, 1],
+    [2, 1, 16, 5, 2, 2, 0, 12, 16, 4, 4, 2, 2, 1, 1],
+    [1, 1, 32, 5, 1, 4, 4, 4, 32, 4, 4, 2, 2, 1, 1],
+    [2, 1, 64, 5, 1, 4, 8, 8, 64, 8, 8, 2, 2, 1, 1],
+    [1, 1, 96, 3, 1, 4, 8, 8, 96, 8, 8, 2, 2, 1, 2],
+    [1, 1, 384, 3, 1, 4, 12, 12, 0, 0, 0, 2, 2, 1, 2],
+]
+M1_cfgs = [
+    # s, n, c, ks, c1, c2, g1, g2, c3, g3, g4
+    [2, 1, 8, 3, 2, 2, 0, 6, 8, 2, 2, 2, 0, 1, 1],
+    [2, 1, 16, 3, 2, 2, 0, 8, 16, 4, 4, 2, 2, 1, 1],
+    [2, 1, 16, 5, 2, 2, 0, 16, 16, 4, 4, 2, 2, 1, 1],
+    [1, 1, 32, 5, 1, 6, 4, 4, 32, 4, 4, 2, 2, 1, 1],
+    [2, 1, 64, 5, 1, 6, 8, 8, 64, 8, 8, 2, 2, 1, 1],
+    [1, 1, 96, 3, 1, 6, 8, 8, 96, 8, 8, 2, 2, 1, 2],
+    [1, 1, 576, 3, 1, 6, 12, 12, 0, 0, 0, 2, 2, 1, 2],
+]
+M2_cfgs = [
+    # s, n, c, ks, c1, c2, g1, g2, c3, g3, g4
+    [2, 1, 12, 3, 2, 2, 0, 8, 12, 4, 4, 2, 0, 1, 1],
+    [2, 1, 16, 3, 2, 2, 0, 12, 16, 4, 4, 2, 2, 1, 1],
+    [1, 1, 24, 3, 2, 2, 0, 16, 24, 4, 4, 2, 2, 1, 1],
+    [2, 1, 32, 5, 1, 6, 6, 6, 32, 4, 4, 2, 2, 1, 1],
+    [1, 1, 32, 5, 1, 6, 8, 8, 32, 4, 4, 2, 2, 1, 2],
+    [1, 1, 64, 5, 1, 6, 8, 8, 64, 8, 8, 2, 2, 1, 2],
+    [2, 1, 96, 5, 1, 6, 8, 8, 96, 8, 8, 2, 2, 1, 2],
+    [1, 1, 128, 3, 1, 6, 12, 12, 128, 8, 8, 2, 2, 1, 2],
+    [1, 1, 768, 3, 1, 6, 16, 16, 0, 0, 0, 2, 2, 1, 2],
+]
+M3_cfgs = [
+    # s, n, c, ks, c1, c2, g1, g2, c3, g3, g4
+    [2, 1, 16, 3, 2, 2, 0, 12, 16, 4, 4, 0, 2, 0, 1],
+    [2, 1, 24, 3, 2, 2, 0, 16, 24, 4, 4, 0, 2, 0, 1],
+    [1, 1, 24, 3, 2, 2, 0, 24, 24, 4, 4, 0, 2, 0, 1],
+    [2, 1, 32, 5, 1, 6, 6, 6, 32, 4, 4, 0, 2, 0, 1],
+    [1, 1, 32, 5, 1, 6, 8, 8, 32, 4, 4, 0, 2, 0, 2],
+    [1, 1, 64, 5, 1, 6, 8, 8, 48, 8, 8, 0, 2, 0, 2],
+    [1, 1, 80, 5, 1, 6, 8, 8, 80, 8, 8, 0, 2, 0, 2],
+    [1, 1, 80, 5, 1, 6, 10, 10, 80, 8, 8, 0, 2, 0, 2],
+    [1, 1, 120, 5, 1, 6, 10, 10, 120, 10, 10, 0, 2, 0, 2],
+    [1, 1, 120, 5, 1, 6, 12, 12, 120, 10, 10, 0, 2, 0, 2],
+    [1, 1, 144, 3, 1, 6, 12, 12, 144, 12, 12, 0, 2, 0, 2],
+    [1, 1, 432, 3, 1, 3, 12, 12, 0, 0, 0, 0, 2, 0, 2],
+]
+
+
+def get_micronet_config(mode):
+    return eval(mode + "_cfgs")
+
+
+class MaxGroupPooling(nn.Layer):
+    def __init__(self, channel_per_group=2):
+        super(MaxGroupPooling, self).__init__()
+        self.channel_per_group = channel_per_group
+
+    def forward(self, x):
+        if self.channel_per_group == 1:
+            return x
+        # max op
+        b, c, h, w = x.shape
+
+        # reshape
+        y = paddle.reshape(x, [b, c // self.channel_per_group, -1, h, w])
+        out = paddle.max(y, axis=2)
+        return out
+
+
+class SpatialSepConvSF(nn.Layer):
+    def __init__(self, inp, oups, kernel_size, stride):
+        super(SpatialSepConvSF, self).__init__()
+
+        oup1, oup2 = oups
+        self.conv = nn.Sequential(
+            nn.Conv2D(
+                inp,
+                oup1,
+                (kernel_size, 1),
+                (stride, 1),
+                (kernel_size // 2, 0),
+                bias_attr=False,
+                groups=1,
+            ),
+            nn.BatchNorm2D(oup1),
+            nn.Conv2D(
+                oup1,
+                oup1 * oup2,
+                (1, kernel_size),
+                (1, stride),
+                (0, kernel_size // 2),
+                bias_attr=False,
+                groups=oup1,
+            ),
+            nn.BatchNorm2D(oup1 * oup2),
+            ChannelShuffle(oup1),
+        )
+
+    def forward(self, x):
+        out = self.conv(x)
+        return out
+
+
+class ChannelShuffle(nn.Layer):
+    def __init__(self, groups):
+        super(ChannelShuffle, self).__init__()
+        self.groups = groups
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        channels_per_group = c // self.groups
+
+        # reshape
+        x = paddle.reshape(x, [b, self.groups, channels_per_group, h, w])
+
+        x = paddle.transpose(x, (0, 2, 1, 3, 4))
+        out = paddle.reshape(x, [b, -1, h, w])
+
+        return out
+
+
+class StemLayer(nn.Layer):
+    def __init__(self, inp, oup, stride, groups=(4, 4)):
+        super(StemLayer, self).__init__()
+
+        g1, g2 = groups
+        self.stem = nn.Sequential(
+            SpatialSepConvSF(inp, groups, 3, stride),
+            MaxGroupPooling(2) if g1 * g2 == 2 * oup else nn.ReLU6(),
+        )
+
+    def forward(self, x):
+        out = self.stem(x)
+        return out
+
+
+class DepthSpatialSepConv(nn.Layer):
+    def __init__(self, inp, expand, kernel_size, stride):
+        super(DepthSpatialSepConv, self).__init__()
+
+        exp1, exp2 = expand
+
+        hidden_dim = inp * exp1
+        oup = inp * exp1 * exp2
+
+        self.conv = nn.Sequential(
+            nn.Conv2D(
+                inp,
+                inp * exp1,
+                (kernel_size, 1),
+                (stride, 1),
+                (kernel_size // 2, 0),
+                bias_attr=False,
+                groups=inp,
+            ),
+            nn.BatchNorm2D(inp * exp1),
+            nn.Conv2D(
+                hidden_dim,
+                oup,
+                (1, kernel_size),
+                1,
+                (0, kernel_size // 2),
+                bias_attr=False,
+                groups=hidden_dim,
+            ),
+            nn.BatchNorm2D(oup),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class GroupConv(nn.Layer):
+    def __init__(self, inp, oup, groups=2):
+        super(GroupConv, self).__init__()
+        self.inp = inp
+        self.oup = oup
+        self.groups = groups
+        self.conv = nn.Sequential(
+            nn.Conv2D(inp, oup, 1, 1, 0, bias_attr=False, groups=self.groups[0]),
+            nn.BatchNorm2D(oup),
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class DepthConv(nn.Layer):
+    def __init__(self, inp, oup, kernel_size, stride):
+        super(DepthConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2D(
+                inp,
+                oup,
+                kernel_size,
+                stride,
+                kernel_size // 2,
+                bias_attr=False,
+                groups=inp,
+            ),
+            nn.BatchNorm2D(oup),
+        )
+
+    def forward(self, x):
+        out = self.conv(x)
+        return out
+
+
+class DYShiftMax(nn.Layer):
+    def __init__(
+        self,
+        inp,
+        oup,
+        reduction=4,
+        act_max=1.0,
+        act_relu=True,
+        init_a=[0.0, 0.0],
+        init_b=[0.0, 0.0],
+        relu_before_pool=False,
+        g=None,
+        expansion=False,
+    ):
+        super(DYShiftMax, self).__init__()
+        self.oup = oup
+        self.act_max = act_max * 2
+        self.act_relu = act_relu
+        self.avg_pool = nn.Sequential(
+            nn.ReLU() if relu_before_pool == True else nn.Sequential(),
+            nn.AdaptiveAvgPool2D(1),
+        )
+
+        self.exp = 4 if act_relu else 2
+        self.init_a = init_a
+        self.init_b = init_b
+
+        # determine squeeze
+        squeeze = make_divisible(inp // reduction, 4)
+        if squeeze < 4:
+            squeeze = 4
+
+        self.fc = nn.Sequential(
+            nn.Linear(inp, squeeze),
+            nn.ReLU(),
+            nn.Linear(squeeze, oup * self.exp),
+            nn.Hardsigmoid(),
+        )
+
+        if g is None:
+            g = 1
+        self.g = g[1]
+        if self.g != 1 and expansion:
+            self.g = inp // self.g
+
+        self.gc = inp // self.g
+        index = paddle.to_tensor([range(inp)])
+        index = paddle.reshape(index, [1, inp, 1, 1])
+        index = paddle.reshape(index, [1, self.g, self.gc, 1, 1])
+        indexgs = paddle.split(index, [1, self.g - 1], axis=1)
+        indexgs = paddle.concat((indexgs[1], indexgs[0]), axis=1)
+        indexes = paddle.split(indexgs, [1, self.gc - 1], axis=2)
+        indexes = paddle.concat((indexes[1], indexes[0]), axis=2)
+        self.index = paddle.reshape(indexes, [inp])
+        self.expansion = expansion
+
+    def forward(self, x):
+        x_in = x
+        x_out = x
+
+        b, c, _, _ = x_in.shape
+        y = self.avg_pool(x_in)
+        y = paddle.reshape(y, [b, c])
+        y = self.fc(y)
+        y = paddle.reshape(y, [b, self.oup * self.exp, 1, 1])
+        y = (y - 0.5) * self.act_max
+
+        n2, c2, h2, w2 = x_out.shape
+        x2 = paddle.to_tensor(x_out.numpy()[:, self.index.numpy(), :, :])
+
+        if self.exp == 4:
+            temp = y.shape
+            a1, b1, a2, b2 = paddle.split(y, temp[1] // self.oup, axis=1)
+
+            a1 = a1 + self.init_a[0]
+            a2 = a2 + self.init_a[1]
+
+            b1 = b1 + self.init_b[0]
+            b2 = b2 + self.init_b[1]
+
+            z1 = x_out * a1 + x2 * b1
+            z2 = x_out * a2 + x2 * b2
+
+            out = paddle.maximum(z1, z2)
+
+        elif self.exp == 2:
+            temp = y.shape
+            a1, b1 = paddle.split(y, temp[1] // self.oup, axis=1)
+            a1 = a1 + self.init_a[0]
+            b1 = b1 + self.init_b[0]
+            out = x_out * a1 + x2 * b1
+
+        return out
+
+
+class DYMicroBlock(nn.Layer):
+    def __init__(
+        self,
+        inp,
+        oup,
+        kernel_size=3,
+        stride=1,
+        ch_exp=(2, 2),
+        ch_per_group=4,
+        groups_1x1=(1, 1),
+        depthsep=True,
+        shuffle=False,
+        activation_cfg=None,
+    ):
+        super(DYMicroBlock, self).__init__()
+
+        self.identity = stride == 1 and inp == oup
+
+        y1, y2, y3 = activation_cfg["dy"]
+        act_reduction = 8 * activation_cfg["ratio"]
+        init_a = activation_cfg["init_a"]
+        init_b = activation_cfg["init_b"]
+
+        t1 = ch_exp
+        gs1 = ch_per_group
+        hidden_fft, g1, g2 = groups_1x1
+        hidden_dim2 = inp * t1[0] * t1[1]
+
+        if gs1[0] == 0:
+            self.layers = nn.Sequential(
+                DepthSpatialSepConv(inp, t1, kernel_size, stride),
+                (
+                    DYShiftMax(
+                        hidden_dim2,
+                        hidden_dim2,
+                        act_max=2.0,
+                        act_relu=True if y2 == 2 else False,
+                        init_a=init_a,
+                        reduction=act_reduction,
+                        init_b=init_b,
+                        g=gs1,
+                        expansion=False,
+                    )
+                    if y2 > 0
+                    else nn.ReLU6()
+                ),
+                ChannelShuffle(gs1[1]) if shuffle else nn.Sequential(),
+                (
+                    ChannelShuffle(hidden_dim2 // 2)
+                    if shuffle and y2 != 0
+                    else nn.Sequential()
+                ),
+                GroupConv(hidden_dim2, oup, (g1, g2)),
+                (
+                    DYShiftMax(
+                        oup,
+                        oup,
+                        act_max=2.0,
+                        act_relu=False,
+                        init_a=[1.0, 0.0],
+                        reduction=act_reduction // 2,
+                        init_b=[0.0, 0.0],
+                        g=(g1, g2),
+                        expansion=False,
+                    )
+                    if y3 > 0
+                    else nn.Sequential()
+                ),
+                ChannelShuffle(g2) if shuffle else nn.Sequential(),
+                (
+                    ChannelShuffle(oup // 2)
+                    if shuffle and oup % 2 == 0 and y3 != 0
+                    else nn.Sequential()
+                ),
+            )
+        elif g2 == 0:
+            self.layers = nn.Sequential(
+                GroupConv(inp, hidden_dim2, gs1),
+                (
+                    DYShiftMax(
+                        hidden_dim2,
+                        hidden_dim2,
+                        act_max=2.0,
+                        act_relu=False,
+                        init_a=[1.0, 0.0],
+                        reduction=act_reduction,
+                        init_b=[0.0, 0.0],
+                        g=gs1,
+                        expansion=False,
+                    )
+                    if y3 > 0
+                    else nn.Sequential()
+                ),
+            )
+        else:
+            self.layers = nn.Sequential(
+                GroupConv(inp, hidden_dim2, gs1),
+                (
+                    DYShiftMax(
+                        hidden_dim2,
+                        hidden_dim2,
+                        act_max=2.0,
+                        act_relu=True if y1 == 2 else False,
+                        init_a=init_a,
+                        reduction=act_reduction,
+                        init_b=init_b,
+                        g=gs1,
+                        expansion=False,
+                    )
+                    if y1 > 0
+                    else nn.ReLU6()
+                ),
+                ChannelShuffle(gs1[1]) if shuffle else nn.Sequential(),
+                (
+                    DepthSpatialSepConv(hidden_dim2, (1, 1), kernel_size, stride)
+                    if depthsep
+                    else DepthConv(hidden_dim2, hidden_dim2, kernel_size, stride)
+                ),
+                nn.Sequential(),
+                (
+                    DYShiftMax(
+                        hidden_dim2,
+                        hidden_dim2,
+                        act_max=2.0,
+                        act_relu=True if y2 == 2 else False,
+                        init_a=init_a,
+                        reduction=act_reduction,
+                        init_b=init_b,
+                        g=gs1,
+                        expansion=True,
+                    )
+                    if y2 > 0
+                    else nn.ReLU6()
+                ),
+                (
+                    ChannelShuffle(hidden_dim2 // 4)
+                    if shuffle and y1 != 0 and y2 != 0
+                    else (
+                        nn.Sequential()
+                        if y1 == 0 and y2 == 0
+                        else ChannelShuffle(hidden_dim2 // 2)
+                    )
+                ),
+                GroupConv(hidden_dim2, oup, (g1, g2)),
+                (
+                    DYShiftMax(
+                        oup,
+                        oup,
+                        act_max=2.0,
+                        act_relu=False,
+                        init_a=[1.0, 0.0],
+                        reduction=(
+                            act_reduction // 2 if oup < hidden_dim2 else act_reduction
+                        ),
+                        init_b=[0.0, 0.0],
+                        g=(g1, g2),
+                        expansion=False,
+                    )
+                    if y3 > 0
+                    else nn.Sequential()
+                ),
+                ChannelShuffle(g2) if shuffle else nn.Sequential(),
+                ChannelShuffle(oup // 2) if shuffle and y3 != 0 else nn.Sequential(),
+            )
+
+    def forward(self, x):
+        identity = x
+        out = self.layers(x)
+
+        if self.identity:
+            out = out + identity
+
+        return out
+
+
+class MicroNet(nn.Layer):
+    """
+    the MicroNet backbone network for recognition module.
+    Args:
+        mode(str): {'M0', 'M1', 'M2', 'M3'}
+            Four models are proposed based on four different computational costs (4M, 6M, 12M, 21M MAdds)
+            Default: 'M3'.
+    """
+
+    def __init__(self, mode="M3", **kwargs):
+        super(MicroNet, self).__init__()
+
+        self.cfgs = get_micronet_config(mode)
+
+        activation_cfg = {}
+        if mode == "M0":
+            input_channel = 4
+            stem_groups = 2, 2
+            out_ch = 384
+            activation_cfg["init_a"] = 1.0, 1.0
+            activation_cfg["init_b"] = 0.0, 0.0
+        elif mode == "M1":
+            input_channel = 6
+            stem_groups = 3, 2
+            out_ch = 576
+            activation_cfg["init_a"] = 1.0, 1.0
+            activation_cfg["init_b"] = 0.0, 0.0
+        elif mode == "M2":
+            input_channel = 8
+            stem_groups = 4, 2
+            out_ch = 768
+            activation_cfg["init_a"] = 1.0, 1.0
+            activation_cfg["init_b"] = 0.0, 0.0
+        elif mode == "M3":
+            input_channel = 12
+            stem_groups = 4, 3
+            out_ch = 432
+            activation_cfg["init_a"] = 1.0, 0.5
+            activation_cfg["init_b"] = 0.0, 0.5
+        else:
+            raise NotImplementedError("mode[" + mode + "_model] is not implemented!")
+
+        layers = [StemLayer(3, input_channel, stride=2, groups=stem_groups)]
+
+        for idx, val in enumerate(self.cfgs):
+            s, n, c, ks, c1, c2, g1, g2, c3, g3, g4, y1, y2, y3, r = val
+
+            t1 = (c1, c2)
+            gs1 = (g1, g2)
+            gs2 = (c3, g3, g4)
+            activation_cfg["dy"] = [y1, y2, y3]
+            activation_cfg["ratio"] = r
+
+            output_channel = c
+            layers.append(
+                DYMicroBlock(
+                    input_channel,
+                    output_channel,
+                    kernel_size=ks,
+                    stride=s,
+                    ch_exp=t1,
+                    ch_per_group=gs1,
+                    groups_1x1=gs2,
+                    depthsep=True,
+                    shuffle=True,
+                    activation_cfg=activation_cfg,
+                )
+            )
+            input_channel = output_channel
+            for i in range(1, n):
+                layers.append(
+                    DYMicroBlock(
+                        input_channel,
+                        output_channel,
+                        kernel_size=ks,
+                        stride=1,
+                        ch_exp=t1,
+                        ch_per_group=gs1,
+                        groups_1x1=gs2,
+                        depthsep=True,
+                        shuffle=True,
+                        activation_cfg=activation_cfg,
+                    )
+                )
+                input_channel = output_channel
+        self.features = nn.Sequential(*layers)
+
+        self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+
+        self.out_channels = make_divisible(out_ch)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.pool(x)
+        return x

ppocr/modeling/backbones/rec_mobilenet_v3.py

@@ -0,0 +1,156 @@
+# 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.
+
+from paddle import nn
+
+from ppocr.modeling.backbones.det_mobilenet_v3 import (
+    ResidualUnit,
+    ConvBNLayer,
+    make_divisible,
+)
+
+__all__ = ["MobileNetV3"]
+
+
+class MobileNetV3(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        model_name="small",
+        scale=0.5,
+        large_stride=None,
+        small_stride=None,
+        disable_se=False,
+        **kwargs,
+    ):
+        super(MobileNetV3, self).__init__()
+        self.disable_se = disable_se
+        if small_stride is None:
+            small_stride = [2, 2, 2, 2]
+        if large_stride is None:
+            large_stride = [1, 2, 2, 2]
+
+        assert isinstance(
+            large_stride, list
+        ), "large_stride type must " "be list but got {}".format(type(large_stride))
+        assert isinstance(
+            small_stride, list
+        ), "small_stride type must " "be list but got {}".format(type(small_stride))
+        assert (
+            len(large_stride) == 4
+        ), "large_stride length must be " "4 but got {}".format(len(large_stride))
+        assert (
+            len(small_stride) == 4
+        ), "small_stride length must be " "4 but got {}".format(len(small_stride))
+
+        if model_name == "large":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, False, "relu", large_stride[0]],
+                [3, 64, 24, False, "relu", (large_stride[1], 1)],
+                [3, 72, 24, False, "relu", 1],
+                [5, 72, 40, True, "relu", (large_stride[2], 1)],
+                [5, 120, 40, True, "relu", 1],
+                [5, 120, 40, True, "relu", 1],
+                [3, 240, 80, False, "hardswish", 1],
+                [3, 200, 80, False, "hardswish", 1],
+                [3, 184, 80, False, "hardswish", 1],
+                [3, 184, 80, False, "hardswish", 1],
+                [3, 480, 112, True, "hardswish", 1],
+                [3, 672, 112, True, "hardswish", 1],
+                [5, 672, 160, True, "hardswish", (large_stride[3], 1)],
+                [5, 960, 160, True, "hardswish", 1],
+                [5, 960, 160, True, "hardswish", 1],
+            ]
+            cls_ch_squeeze = 960
+        elif model_name == "small":
+            cfg = [
+                # k, exp, c,  se,     nl,  s,
+                [3, 16, 16, True, "relu", (small_stride[0], 1)],
+                [3, 72, 24, False, "relu", (small_stride[1], 1)],
+                [3, 88, 24, False, "relu", 1],
+                [5, 96, 40, True, "hardswish", (small_stride[2], 1)],
+                [5, 240, 40, True, "hardswish", 1],
+                [5, 240, 40, True, "hardswish", 1],
+                [5, 120, 48, True, "hardswish", 1],
+                [5, 144, 48, True, "hardswish", 1],
+                [5, 288, 96, True, "hardswish", (small_stride[3], 1)],
+                [5, 576, 96, True, "hardswish", 1],
+                [5, 576, 96, True, "hardswish", 1],
+            ]
+            cls_ch_squeeze = 576
+        else:
+            raise NotImplementedError(
+                "mode[" + model_name + "_model] is not implemented!"
+            )
+
+        supported_scale = [0.35, 0.5, 0.75, 1.0, 1.25]
+        assert (
+            scale in supported_scale
+        ), "supported scales are {} but input scale is {}".format(
+            supported_scale, scale
+        )
+
+        inplanes = 16
+        # conv1
+        self.conv1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=make_divisible(inplanes * scale),
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=1,
+            if_act=True,
+            act="hardswish",
+        )
+        i = 0
+        block_list = []
+        inplanes = make_divisible(inplanes * scale)
+        for k, exp, c, se, nl, s in cfg:
+            se = se and not self.disable_se
+            block_list.append(
+                ResidualUnit(
+                    in_channels=inplanes,
+                    mid_channels=make_divisible(scale * exp),
+                    out_channels=make_divisible(scale * c),
+                    kernel_size=k,
+                    stride=s,
+                    use_se=se,
+                    act=nl,
+                )
+            )
+            inplanes = make_divisible(scale * c)
+            i += 1
+        self.blocks = nn.Sequential(*block_list)
+
+        self.conv2 = ConvBNLayer(
+            in_channels=inplanes,
+            out_channels=make_divisible(scale * cls_ch_squeeze),
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1,
+            if_act=True,
+            act="hardswish",
+        )
+
+        self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.out_channels = make_divisible(scale * cls_ch_squeeze)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.blocks(x)
+        x = self.conv2(x)
+        x = self.pool(x)
+        return x

ppocr/modeling/backbones/rec_mv1_enhance.py

@@ -0,0 +1,283 @@
+# copyright (c) 2021 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/PaddlePaddle/PaddleClas/blob/develop/ppcls/arch/backbone/legendary_models/pp_lcnet.py
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import numpy as np
+import paddle
+from paddle import ParamAttr, reshape, transpose
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+from paddle.nn.initializer import KaimingNormal
+from paddle.regularizer import L2Decay
+from paddle.nn.functional import hardswish, hardsigmoid
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        num_channels,
+        filter_size,
+        num_filters,
+        stride,
+        padding,
+        channels=None,
+        num_groups=1,
+        act="hard_swish",
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self._conv = Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=num_groups,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+            bias_attr=False,
+        )
+
+        self._batch_norm = BatchNorm(
+            num_filters,
+            act=act,
+            param_attr=ParamAttr(regularizer=L2Decay(0.0)),
+            bias_attr=ParamAttr(regularizer=L2Decay(0.0)),
+        )
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class DepthwiseSeparable(nn.Layer):
+    def __init__(
+        self,
+        num_channels,
+        num_filters1,
+        num_filters2,
+        num_groups,
+        stride,
+        scale,
+        dw_size=3,
+        padding=1,
+        use_se=False,
+    ):
+        super(DepthwiseSeparable, self).__init__()
+        self.use_se = use_se
+        self._depthwise_conv = ConvBNLayer(
+            num_channels=num_channels,
+            num_filters=int(num_filters1 * scale),
+            filter_size=dw_size,
+            stride=stride,
+            padding=padding,
+            num_groups=int(num_groups * scale),
+        )
+        if use_se:
+            self._se = SEModule(int(num_filters1 * scale))
+        self._pointwise_conv = ConvBNLayer(
+            num_channels=int(num_filters1 * scale),
+            filter_size=1,
+            num_filters=int(num_filters2 * scale),
+            stride=1,
+            padding=0,
+        )
+
+    def forward(self, inputs):
+        y = self._depthwise_conv(inputs)
+        if self.use_se:
+            y = self._se(y)
+        y = self._pointwise_conv(y)
+        return y
+
+
+class MobileNetV1Enhance(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        scale=0.5,
+        last_conv_stride=1,
+        last_pool_type="max",
+        last_pool_kernel_size=[3, 2],
+        **kwargs,
+    ):
+        super().__init__()
+        self.scale = scale
+        self.block_list = []
+
+        self.conv1 = ConvBNLayer(
+            num_channels=3,
+            filter_size=3,
+            channels=3,
+            num_filters=int(32 * scale),
+            stride=2,
+            padding=1,
+        )
+
+        conv2_1 = DepthwiseSeparable(
+            num_channels=int(32 * scale),
+            num_filters1=32,
+            num_filters2=64,
+            num_groups=32,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv2_1)
+
+        conv2_2 = DepthwiseSeparable(
+            num_channels=int(64 * scale),
+            num_filters1=64,
+            num_filters2=128,
+            num_groups=64,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv2_2)
+
+        conv3_1 = DepthwiseSeparable(
+            num_channels=int(128 * scale),
+            num_filters1=128,
+            num_filters2=128,
+            num_groups=128,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv3_1)
+
+        conv3_2 = DepthwiseSeparable(
+            num_channels=int(128 * scale),
+            num_filters1=128,
+            num_filters2=256,
+            num_groups=128,
+            stride=(2, 1),
+            scale=scale,
+        )
+        self.block_list.append(conv3_2)
+
+        conv4_1 = DepthwiseSeparable(
+            num_channels=int(256 * scale),
+            num_filters1=256,
+            num_filters2=256,
+            num_groups=256,
+            stride=1,
+            scale=scale,
+        )
+        self.block_list.append(conv4_1)
+
+        conv4_2 = DepthwiseSeparable(
+            num_channels=int(256 * scale),
+            num_filters1=256,
+            num_filters2=512,
+            num_groups=256,
+            stride=(2, 1),
+            scale=scale,
+        )
+        self.block_list.append(conv4_2)
+
+        for _ in range(5):
+            conv5 = DepthwiseSeparable(
+                num_channels=int(512 * scale),
+                num_filters1=512,
+                num_filters2=512,
+                num_groups=512,
+                stride=1,
+                dw_size=5,
+                padding=2,
+                scale=scale,
+                use_se=False,
+            )
+            self.block_list.append(conv5)
+
+        conv5_6 = DepthwiseSeparable(
+            num_channels=int(512 * scale),
+            num_filters1=512,
+            num_filters2=1024,
+            num_groups=512,
+            stride=(2, 1),
+            dw_size=5,
+            padding=2,
+            scale=scale,
+            use_se=True,
+        )
+        self.block_list.append(conv5_6)
+
+        conv6 = DepthwiseSeparable(
+            num_channels=int(1024 * scale),
+            num_filters1=1024,
+            num_filters2=1024,
+            num_groups=1024,
+            stride=last_conv_stride,
+            dw_size=5,
+            padding=2,
+            use_se=True,
+            scale=scale,
+        )
+        self.block_list.append(conv6)
+
+        self.block_list = nn.Sequential(*self.block_list)
+        if last_pool_type == "avg":
+            self.pool = nn.AvgPool2D(
+                kernel_size=last_pool_kernel_size,
+                stride=last_pool_kernel_size,
+                padding=0,
+            )
+        else:
+            self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.out_channels = int(1024 * scale)
+
+    def forward(self, inputs):
+        y = self.conv1(inputs)
+        y = self.block_list(y)
+        y = self.pool(y)
+        return y
+
+
+class SEModule(nn.Layer):
+    def __init__(self, channel, reduction=4):
+        super(SEModule, self).__init__()
+        self.avg_pool = AdaptiveAvgPool2D(1)
+        self.conv1 = Conv2D(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            weight_attr=ParamAttr(),
+            bias_attr=ParamAttr(),
+        )
+        self.conv2 = Conv2D(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            weight_attr=ParamAttr(),
+            bias_attr=ParamAttr(),
+        )
+
+    def forward(self, inputs):
+        outputs = self.avg_pool(inputs)
+        outputs = self.conv1(outputs)
+        outputs = F.relu(outputs)
+        outputs = self.conv2(outputs)
+        outputs = hardsigmoid(outputs)
+        return paddle.multiply(x=inputs, y=outputs)

ppocr/modeling/backbones/rec_nrtr_mtb.py

@@ -0,0 +1,47 @@
+# 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.
+
+from paddle import nn
+import paddle
+
+
+class MTB(nn.Layer):
+    def __init__(self, cnn_num, in_channels):
+        super(MTB, self).__init__()
+        self.block = nn.Sequential()
+        self.out_channels = in_channels
+        self.cnn_num = cnn_num
+        if self.cnn_num == 2:
+            for i in range(self.cnn_num):
+                self.block.add_sublayer(
+                    "conv_{}".format(i),
+                    nn.Conv2D(
+                        in_channels=in_channels if i == 0 else 32 * (2 ** (i - 1)),
+                        out_channels=32 * (2**i),
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                    ),
+                )
+                self.block.add_sublayer("relu_{}".format(i), nn.ReLU())
+                self.block.add_sublayer("bn_{}".format(i), nn.BatchNorm2D(32 * (2**i)))
+
+    def forward(self, images):
+        x = self.block(images)
+        if self.cnn_num == 2:
+            # (b, w, h, c)
+            x = paddle.transpose(x, [0, 3, 2, 1])
+            x_shape = x.shape
+            x = paddle.reshape(x, [x_shape[0], x_shape[1], x_shape[2] * x_shape[3]])
+        return x

ppocr/modeling/backbones/rec_repvit.py

@@ -0,0 +1,363 @@
+# copyright (c) 2024 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/THU-MIG/RepViT
+"""
+
+import paddle.nn as nn
+import paddle
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+trunc_normal_ = TruncatedNormal(std=0.02)
+normal_ = Normal
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+def _make_divisible(v, divisor, min_value=None):
+    """
+    This function is taken from the original tf repo.
+    It ensures that all layers have a channel number that is divisible by 8
+    It can be seen here:
+    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
+    :param v:
+    :param divisor:
+    :param min_value:
+    :return:
+    """
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+# from timm.models.layers import SqueezeExcite
+
+
+def make_divisible(v, divisor=8, min_value=None, round_limit=0.9):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v
+
+
+class SEModule(nn.Layer):
+    """SE Module as defined in original SE-Nets with a few additions
+    Additions include:
+        * divisor can be specified to keep channels % div == 0 (default: 8)
+        * reduction channels can be specified directly by arg (if rd_channels is set)
+        * reduction channels can be specified by float rd_ratio (default: 1/16)
+        * global max pooling can be added to the squeeze aggregation
+        * customizable activation, normalization, and gate layer
+    """
+
+    def __init__(
+        self,
+        channels,
+        rd_ratio=1.0 / 16,
+        rd_channels=None,
+        rd_divisor=8,
+        act_layer=nn.ReLU,
+    ):
+        super(SEModule, self).__init__()
+        if not rd_channels:
+            rd_channels = make_divisible(
+                channels * rd_ratio, rd_divisor, round_limit=0.0
+            )
+        self.fc1 = nn.Conv2D(channels, rd_channels, kernel_size=1, bias_attr=True)
+        self.act = act_layer()
+        self.fc2 = nn.Conv2D(rd_channels, channels, kernel_size=1, bias_attr=True)
+
+    def forward(self, x):
+        x_se = x.mean((2, 3), keepdim=True)
+        x_se = self.fc1(x_se)
+        x_se = self.act(x_se)
+        x_se = self.fc2(x_se)
+        return x * nn.functional.sigmoid(x_se)
+
+
+class Conv2D_BN(nn.Sequential):
+    def __init__(
+        self,
+        a,
+        b,
+        ks=1,
+        stride=1,
+        pad=0,
+        dilation=1,
+        groups=1,
+        bn_weight_init=1,
+        resolution=-10000,
+    ):
+        super().__init__()
+        self.add_sublayer(
+            "c", nn.Conv2D(a, b, ks, stride, pad, dilation, groups, bias_attr=False)
+        )
+        self.add_sublayer("bn", nn.BatchNorm2D(b))
+        if bn_weight_init == 1:
+            ones_(self.bn.weight)
+        else:
+            zeros_(self.bn.weight)
+        zeros_(self.bn.bias)
+
+    @paddle.no_grad()
+    def fuse(self):
+        c, bn = self.c, self.bn
+        w = bn.weight / (bn._variance + bn._epsilon) ** 0.5
+        w = c.weight * w[:, None, None, None]
+        b = bn.bias - bn._mean * bn.weight / (bn._variance + bn._epsilon) ** 0.5
+        m = nn.Conv2D(
+            w.shape[1] * self.c._groups,
+            w.shape[0],
+            w.shape[2:],
+            stride=self.c._stride,
+            padding=self.c._padding,
+            dilation=self.c._dilation,
+            groups=self.c._groups,
+        )
+        m.weight.set_value(w)
+        m.bias.set_value(b)
+        return m
+
+
+class Residual(nn.Layer):
+    def __init__(self, m, drop=0.0):
+        super().__init__()
+        self.m = m
+        self.drop = drop
+
+    def forward(self, x):
+        if self.training and self.drop > 0:
+            return (
+                x
+                + self.m(x)
+                * paddle.rand(x.size(0), 1, 1, 1)
+                .ge_(self.drop)
+                .div(1 - self.drop)
+                .detach()
+            )
+        else:
+            return x + self.m(x)
+
+    @paddle.no_grad()
+    def fuse(self):
+        if isinstance(self.m, Conv2D_BN):
+            m = self.m.fuse()
+            assert m._groups == m.in_channels
+            identity = paddle.ones([m.weight.shape[0], m.weight.shape[1], 1, 1])
+            identity = nn.functional.pad(identity, [1, 1, 1, 1])
+            m.weight += identity
+            return m
+        elif isinstance(self.m, nn.Conv2D):
+            m = self.m
+            assert m._groups != m.in_channels
+            identity = paddle.ones([m.weight.shape[0], m.weight.shape[1], 1, 1])
+            identity = nn.functional.pad(identity, [1, 1, 1, 1])
+            m.weight += identity
+            return m
+        else:
+            return self
+
+
+class RepVGGDW(nn.Layer):
+    def __init__(self, ed) -> None:
+        super().__init__()
+        self.conv = Conv2D_BN(ed, ed, 3, 1, 1, groups=ed)
+        self.conv1 = nn.Conv2D(ed, ed, 1, 1, 0, groups=ed)
+        self.dim = ed
+        self.bn = nn.BatchNorm2D(ed)
+
+    def forward(self, x):
+        return self.bn((self.conv(x) + self.conv1(x)) + x)
+
+    @paddle.no_grad()
+    def fuse(self):
+        conv = self.conv.fuse()
+        conv1 = self.conv1
+
+        conv_w = conv.weight
+        conv_b = conv.bias
+        conv1_w = conv1.weight
+        conv1_b = conv1.bias
+
+        conv1_w = nn.functional.pad(conv1_w, [1, 1, 1, 1])
+
+        identity = nn.functional.pad(
+            paddle.ones([conv1_w.shape[0], conv1_w.shape[1], 1, 1]), [1, 1, 1, 1]
+        )
+
+        final_conv_w = conv_w + conv1_w + identity
+        final_conv_b = conv_b + conv1_b
+
+        conv.weight.set_value(final_conv_w)
+        conv.bias.set_value(final_conv_b)
+
+        bn = self.bn
+        w = bn.weight / (bn._variance + bn._epsilon) ** 0.5
+        w = conv.weight * w[:, None, None, None]
+        b = (
+            bn.bias
+            + (conv.bias - bn._mean) * bn.weight / (bn._variance + bn._epsilon) ** 0.5
+        )
+        conv.weight.set_value(w)
+        conv.bias.set_value(b)
+        return conv
+
+
+class RepViTBlock(nn.Layer):
+    def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
+        super(RepViTBlock, self).__init__()
+
+        self.identity = stride == 1 and inp == oup
+        assert hidden_dim == 2 * inp
+
+        if stride != 1:
+            self.token_mixer = nn.Sequential(
+                Conv2D_BN(
+                    inp, inp, kernel_size, stride, (kernel_size - 1) // 2, groups=inp
+                ),
+                SEModule(inp, 0.25) if use_se else nn.Identity(),
+                Conv2D_BN(inp, oup, ks=1, stride=1, pad=0),
+            )
+            self.channel_mixer = Residual(
+                nn.Sequential(
+                    # pw
+                    Conv2D_BN(oup, 2 * oup, 1, 1, 0),
+                    nn.GELU() if use_hs else nn.GELU(),
+                    # pw-linear
+                    Conv2D_BN(2 * oup, oup, 1, 1, 0, bn_weight_init=0),
+                )
+            )
+        else:
+            assert self.identity
+            self.token_mixer = nn.Sequential(
+                RepVGGDW(inp),
+                SEModule(inp, 0.25) if use_se else nn.Identity(),
+            )
+            self.channel_mixer = Residual(
+                nn.Sequential(
+                    # pw
+                    Conv2D_BN(inp, hidden_dim, 1, 1, 0),
+                    nn.GELU() if use_hs else nn.GELU(),
+                    # pw-linear
+                    Conv2D_BN(hidden_dim, oup, 1, 1, 0, bn_weight_init=0),
+                )
+            )
+
+    def forward(self, x):
+        return self.channel_mixer(self.token_mixer(x))
+
+
+class RepViT(nn.Layer):
+    def __init__(self, cfgs, in_channels=3, out_indices=None):
+        super(RepViT, self).__init__()
+        # setting of inverted residual blocks
+        self.cfgs = cfgs
+
+        # building first layer
+        input_channel = self.cfgs[0][2]
+        patch_embed = nn.Sequential(
+            Conv2D_BN(in_channels, input_channel // 2, 3, 2, 1),
+            nn.GELU(),
+            Conv2D_BN(input_channel // 2, input_channel, 3, 2, 1),
+        )
+        layers = [patch_embed]
+        # building inverted residual blocks
+        block = RepViTBlock
+        for k, t, c, use_se, use_hs, s in self.cfgs:
+            output_channel = _make_divisible(c, 8)
+            exp_size = _make_divisible(input_channel * t, 8)
+            layers.append(
+                block(input_channel, exp_size, output_channel, k, s, use_se, use_hs)
+            )
+            input_channel = output_channel
+        self.features = nn.LayerList(layers)
+        self.out_indices = out_indices
+        if out_indices is not None:
+            self.out_channels = [self.cfgs[ids - 1][2] for ids in out_indices]
+        else:
+            self.out_channels = self.cfgs[-1][2]
+
+    def forward(self, x):
+        if self.out_indices is not None:
+            return self.forward_det(x)
+        return self.forward_rec(x)
+
+    def forward_det(self, x):
+        outs = []
+        for i, f in enumerate(self.features):
+            x = f(x)
+            if i in self.out_indices:
+                outs.append(x)
+        return outs
+
+    def forward_rec(self, x):
+        for f in self.features:
+            x = f(x)
+        h = x.shape[2]
+        x = nn.functional.avg_pool2d(x, [h, 2])
+        return x
+
+
+def RepSVTR(in_channels=3):
+    """
+    Constructs a MobileNetV3-Large model
+    """
+    # k, t, c, SE, HS, s
+    cfgs = [
+        [3, 2, 96, 1, 0, 1],
+        [3, 2, 96, 0, 0, 1],
+        [3, 2, 96, 0, 0, 1],
+        [3, 2, 192, 0, 1, (2, 1)],
+        [3, 2, 192, 1, 1, 1],
+        [3, 2, 192, 0, 1, 1],
+        [3, 2, 192, 1, 1, 1],
+        [3, 2, 192, 0, 1, 1],
+        [3, 2, 192, 1, 1, 1],
+        [3, 2, 192, 0, 1, 1],
+        [3, 2, 384, 0, 1, (2, 1)],
+        [3, 2, 384, 1, 1, 1],
+        [3, 2, 384, 0, 1, 1],
+    ]
+    return RepViT(cfgs, in_channels=in_channels)
+
+
+def RepSVTR_det(in_channels=3, out_indices=[2, 5, 10, 13]):
+    """
+    Constructs a MobileNetV3-Large model
+    """
+    # k, t, c, SE, HS, s
+    cfgs = [
+        [3, 2, 48, 1, 0, 1],
+        [3, 2, 48, 0, 0, 1],
+        [3, 2, 96, 0, 0, 2],
+        [3, 2, 96, 1, 0, 1],
+        [3, 2, 96, 0, 0, 1],
+        [3, 2, 192, 0, 1, 2],
+        [3, 2, 192, 1, 1, 1],
+        [3, 2, 192, 0, 1, 1],
+        [3, 2, 192, 1, 1, 1],
+        [3, 2, 192, 0, 1, 1],
+        [3, 2, 384, 0, 1, 2],
+        [3, 2, 384, 1, 1, 1],
+        [3, 2, 384, 0, 1, 1],
+    ]
+    return RepViT(cfgs, in_channels=in_channels, out_indices=out_indices)

ppocr/modeling/backbones/rec_resnet_31.py

@@ -0,0 +1,318 @@
+# copyright (c) 2021 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/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/layers/conv_layer.py
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/backbones/resnet31_ocr.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+
+__all__ = ["ResNet31"]
+
+
+def conv3x3(in_channel, out_channel, stride=1, conv_weight_attr=None):
+    return nn.Conv2D(
+        in_channel,
+        out_channel,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        weight_attr=conv_weight_attr,
+        bias_attr=False,
+    )
+
+
+class BasicBlock(nn.Layer):
+    expansion = 1
+
+    def __init__(
+        self,
+        in_channels,
+        channels,
+        stride=1,
+        downsample=False,
+        conv_weight_attr=None,
+        bn_weight_attr=None,
+    ):
+        super().__init__()
+        self.conv1 = conv3x3(
+            in_channels, channels, stride, conv_weight_attr=conv_weight_attr
+        )
+        self.bn1 = nn.BatchNorm2D(channels, weight_attr=bn_weight_attr)
+        self.relu = nn.ReLU()
+        self.conv2 = conv3x3(channels, channels, conv_weight_attr=conv_weight_attr)
+        self.bn2 = nn.BatchNorm2D(channels, weight_attr=bn_weight_attr)
+        self.downsample = downsample
+        if downsample:
+            self.downsample = nn.Sequential(
+                nn.Conv2D(
+                    in_channels,
+                    channels * self.expansion,
+                    1,
+                    stride,
+                    weight_attr=conv_weight_attr,
+                    bias_attr=False,
+                ),
+                nn.BatchNorm2D(channels * self.expansion, weight_attr=bn_weight_attr),
+            )
+        else:
+            self.downsample = nn.Sequential()
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet31(nn.Layer):
+    """
+    Args:
+        in_channels (int): Number of channels of input image tensor.
+        layers (list[int]): List of BasicBlock number for each stage.
+        channels (list[int]): List of out_channels of Conv2d layer.
+        out_indices (None | Sequence[int]): Indices of output stages.
+        last_stage_pool (bool): If True, add `MaxPool2d` layer to last stage.
+        init_type (None | str): the config to control the initialization.
+    """
+
+    def __init__(
+        self,
+        in_channels=3,
+        layers=[1, 2, 5, 3],
+        channels=[64, 128, 256, 256, 512, 512, 512],
+        out_indices=None,
+        last_stage_pool=False,
+        init_type=None,
+    ):
+        super(ResNet31, self).__init__()
+        assert isinstance(in_channels, int)
+        assert isinstance(last_stage_pool, bool)
+
+        self.out_indices = out_indices
+        self.last_stage_pool = last_stage_pool
+
+        conv_weight_attr = None
+        bn_weight_attr = None
+
+        if init_type is not None:
+            support_dict = ["KaimingNormal"]
+            assert init_type in support_dict, Exception(
+                "resnet31 only support {}".format(support_dict)
+            )
+            conv_weight_attr = nn.initializer.KaimingNormal()
+            bn_weight_attr = ParamAttr(
+                initializer=nn.initializer.Uniform(), learning_rate=1
+            )
+
+        # conv 1 (Conv Conv)
+        self.conv1_1 = nn.Conv2D(
+            in_channels,
+            channels[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn1_1 = nn.BatchNorm2D(channels[0], weight_attr=bn_weight_attr)
+        self.relu1_1 = nn.ReLU()
+
+        self.conv1_2 = nn.Conv2D(
+            channels[0],
+            channels[1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn1_2 = nn.BatchNorm2D(channels[1], weight_attr=bn_weight_attr)
+        self.relu1_2 = nn.ReLU()
+
+        # conv 2 (Max-pooling, Residual block, Conv)
+        self.pool2 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0, ceil_mode=True)
+        self.block2 = self._make_layer(
+            channels[1],
+            channels[2],
+            layers[0],
+            conv_weight_attr=conv_weight_attr,
+            bn_weight_attr=bn_weight_attr,
+        )
+        self.conv2 = nn.Conv2D(
+            channels[2],
+            channels[2],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn2 = nn.BatchNorm2D(channels[2], weight_attr=bn_weight_attr)
+        self.relu2 = nn.ReLU()
+
+        # conv 3 (Max-pooling, Residual block, Conv)
+        self.pool3 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0, ceil_mode=True)
+        self.block3 = self._make_layer(
+            channels[2],
+            channels[3],
+            layers[1],
+            conv_weight_attr=conv_weight_attr,
+            bn_weight_attr=bn_weight_attr,
+        )
+        self.conv3 = nn.Conv2D(
+            channels[3],
+            channels[3],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn3 = nn.BatchNorm2D(channels[3], weight_attr=bn_weight_attr)
+        self.relu3 = nn.ReLU()
+
+        # conv 4 (Max-pooling, Residual block, Conv)
+        self.pool4 = nn.MaxPool2D(
+            kernel_size=(2, 1), stride=(2, 1), padding=0, ceil_mode=True
+        )
+        self.block4 = self._make_layer(
+            channels[3],
+            channels[4],
+            layers[2],
+            conv_weight_attr=conv_weight_attr,
+            bn_weight_attr=bn_weight_attr,
+        )
+        self.conv4 = nn.Conv2D(
+            channels[4],
+            channels[4],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn4 = nn.BatchNorm2D(channels[4], weight_attr=bn_weight_attr)
+        self.relu4 = nn.ReLU()
+
+        # conv 5 ((Max-pooling), Residual block, Conv)
+        self.pool5 = None
+        if self.last_stage_pool:
+            self.pool5 = nn.MaxPool2D(
+                kernel_size=2, stride=2, padding=0, ceil_mode=True
+            )
+        self.block5 = self._make_layer(
+            channels[4],
+            channels[5],
+            layers[3],
+            conv_weight_attr=conv_weight_attr,
+            bn_weight_attr=bn_weight_attr,
+        )
+        self.conv5 = nn.Conv2D(
+            channels[5],
+            channels[5],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+        )
+        self.bn5 = nn.BatchNorm2D(channels[5], weight_attr=bn_weight_attr)
+        self.relu5 = nn.ReLU()
+
+        self.out_channels = channels[-1]
+
+    def _make_layer(
+        self,
+        input_channels,
+        output_channels,
+        blocks,
+        conv_weight_attr=None,
+        bn_weight_attr=None,
+    ):
+        layers = []
+        for _ in range(blocks):
+            downsample = None
+            if input_channels != output_channels:
+                downsample = nn.Sequential(
+                    nn.Conv2D(
+                        input_channels,
+                        output_channels,
+                        kernel_size=1,
+                        stride=1,
+                        weight_attr=conv_weight_attr,
+                        bias_attr=False,
+                    ),
+                    nn.BatchNorm2D(output_channels, weight_attr=bn_weight_attr),
+                )
+
+            layers.append(
+                BasicBlock(
+                    input_channels,
+                    output_channels,
+                    downsample=downsample,
+                    conv_weight_attr=conv_weight_attr,
+                    bn_weight_attr=bn_weight_attr,
+                )
+            )
+            input_channels = output_channels
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1_1(x)
+        x = self.bn1_1(x)
+        x = self.relu1_1(x)
+
+        x = self.conv1_2(x)
+        x = self.bn1_2(x)
+        x = self.relu1_2(x)
+
+        outs = []
+        for i in range(4):
+            layer_index = i + 2
+            pool_layer = getattr(self, f"pool{layer_index}")
+            block_layer = getattr(self, f"block{layer_index}")
+            conv_layer = getattr(self, f"conv{layer_index}")
+            bn_layer = getattr(self, f"bn{layer_index}")
+            relu_layer = getattr(self, f"relu{layer_index}")
+
+            if pool_layer is not None:
+                x = pool_layer(x)
+            x = block_layer(x)
+            x = conv_layer(x)
+            x = bn_layer(x)
+            x = relu_layer(x)
+
+            outs.append(x)
+
+        if self.out_indices is not None:
+            return tuple([outs[i] for i in self.out_indices])
+
+        return x

ppocr/modeling/backbones/rec_resnet_32.py

@@ -0,0 +1,305 @@
+# 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/hikopensource/DAVAR-Lab-OCR/davarocr/davar_rcg/models/backbones/ResNet32.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle.nn as nn
+
+__all__ = ["ResNet32"]
+
+conv_weight_attr = nn.initializer.KaimingNormal()
+
+
+class ResNet32(nn.Layer):
+    """
+    Feature Extractor is proposed in  FAN Ref [1]
+
+    Ref [1]: Focusing Attention: Towards Accurate Text Recognition in Neural Images ICCV-2017
+    """
+
+    def __init__(self, in_channels, out_channels=512):
+        """
+
+        Args:
+            in_channels (int): input channel
+            output_channel (int): output channel
+        """
+        super(ResNet32, self).__init__()
+        self.out_channels = out_channels
+        self.ConvNet = ResNet(in_channels, out_channels, BasicBlock, [1, 2, 5, 3])
+
+    def forward(self, inputs):
+        """
+        Args:
+            inputs: input feature
+
+        Returns:
+            output feature
+
+        """
+        return self.ConvNet(inputs)
+
+
+class BasicBlock(nn.Layer):
+    """Res-net Basic Block"""
+
+    expansion = 1
+
+    def __init__(
+        self, inplanes, planes, stride=1, downsample=None, norm_type="BN", **kwargs
+    ):
+        """
+        Args:
+            inplanes (int): input channel
+            planes (int): channels of the middle feature
+            stride (int): stride of the convolution
+            downsample (int): type of the down_sample
+            norm_type (str): type of the normalization
+            **kwargs (None): backup parameter
+        """
+        super(BasicBlock, self).__init__()
+        self.conv1 = self._conv3x3(inplanes, planes)
+        self.bn1 = nn.BatchNorm2D(planes)
+        self.conv2 = self._conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2D(planes)
+        self.relu = nn.ReLU()
+        self.downsample = downsample
+        self.stride = stride
+
+    def _conv3x3(self, in_planes, out_planes, stride=1):
+        """
+
+        Args:
+            in_planes (int): input channel
+            out_planes (int): channels of the middle feature
+            stride (int): stride of the convolution
+        Returns:
+            nn.Layer: Conv2D with kernel = 3
+
+        """
+
+        return nn.Conv2D(
+            in_planes,
+            out_planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Layer):
+    """Res-Net network structure"""
+
+    def __init__(self, input_channel, output_channel, block, layers):
+        """
+
+        Args:
+            input_channel (int): input channel
+            output_channel (int): output channel
+            block (BasicBlock): convolution block
+            layers (list): layers of the block
+        """
+        super(ResNet, self).__init__()
+
+        self.output_channel_block = [
+            int(output_channel / 4),
+            int(output_channel / 2),
+            output_channel,
+            output_channel,
+        ]
+
+        self.inplanes = int(output_channel / 8)
+        self.conv0_1 = nn.Conv2D(
+            input_channel,
+            int(output_channel / 16),
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn0_1 = nn.BatchNorm2D(int(output_channel / 16))
+        self.conv0_2 = nn.Conv2D(
+            int(output_channel / 16),
+            self.inplanes,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn0_2 = nn.BatchNorm2D(self.inplanes)
+        self.relu = nn.ReLU()
+
+        self.maxpool1 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])
+        self.conv1 = nn.Conv2D(
+            self.output_channel_block[0],
+            self.output_channel_block[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm2D(self.output_channel_block[0])
+
+        self.maxpool2 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.layer2 = self._make_layer(
+            block, self.output_channel_block[1], layers[1], stride=1
+        )
+        self.conv2 = nn.Conv2D(
+            self.output_channel_block[1],
+            self.output_channel_block[1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn2 = nn.BatchNorm2D(self.output_channel_block[1])
+
+        self.maxpool3 = nn.MaxPool2D(kernel_size=2, stride=(2, 1), padding=(0, 1))
+        self.layer3 = self._make_layer(
+            block, self.output_channel_block[2], layers[2], stride=1
+        )
+        self.conv3 = nn.Conv2D(
+            self.output_channel_block[2],
+            self.output_channel_block[2],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn3 = nn.BatchNorm2D(self.output_channel_block[2])
+
+        self.layer4 = self._make_layer(
+            block, self.output_channel_block[3], layers[3], stride=1
+        )
+        self.conv4_1 = nn.Conv2D(
+            self.output_channel_block[3],
+            self.output_channel_block[3],
+            kernel_size=2,
+            stride=(2, 1),
+            padding=(0, 1),
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn4_1 = nn.BatchNorm2D(self.output_channel_block[3])
+        self.conv4_2 = nn.Conv2D(
+            self.output_channel_block[3],
+            self.output_channel_block[3],
+            kernel_size=2,
+            stride=1,
+            padding=0,
+            weight_attr=conv_weight_attr,
+            bias_attr=False,
+        )
+        self.bn4_2 = nn.BatchNorm2D(self.output_channel_block[3])
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        """
+
+        Args:
+            block (block): convolution block
+            planes (int): input channels
+            blocks (list): layers of the block
+            stride (int): stride of the convolution
+
+        Returns:
+            nn.Sequential: the combination of the convolution block
+
+        """
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2D(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    weight_attr=conv_weight_attr,
+                    bias_attr=False,
+                ),
+                nn.BatchNorm2D(planes * block.expansion),
+            )
+
+        layers = list()
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv0_1(x)
+        x = self.bn0_1(x)
+        x = self.relu(x)
+        x = self.conv0_2(x)
+        x = self.bn0_2(x)
+        x = self.relu(x)
+
+        x = self.maxpool1(x)
+        x = self.layer1(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.maxpool2(x)
+        x = self.layer2(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.maxpool3(x)
+        x = self.layer3(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.layer4(x)
+        x = self.conv4_1(x)
+        x = self.bn4_1(x)
+        x = self.relu(x)
+        x = self.conv4_2(x)
+        x = self.bn4_2(x)
+        x = self.relu(x)
+        return x

ppocr/modeling/backbones/rec_resnet_45.py

@@ -0,0 +1,150 @@
+# copyright (c) 2021 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/FangShancheng/ABINet/tree/main/modules
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+from paddle.nn.initializer import KaimingNormal
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+import math
+
+__all__ = ["ResNet45"]
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    return nn.Conv2D(
+        in_planes,
+        out_planes,
+        kernel_size=1,
+        stride=1,
+        weight_attr=ParamAttr(initializer=KaimingNormal()),
+        bias_attr=False,
+    )
+
+
+def conv3x3(in_channel, out_channel, stride=1):
+    return nn.Conv2D(
+        in_channel,
+        out_channel,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        weight_attr=ParamAttr(initializer=KaimingNormal()),
+        bias_attr=False,
+    )
+
+
+class BasicBlock(nn.Layer):
+    expansion = 1
+
+    def __init__(self, in_channels, channels, stride=1, downsample=None):
+        super().__init__()
+        self.conv1 = conv1x1(in_channels, channels)
+        self.bn1 = nn.BatchNorm2D(channels)
+        self.relu = nn.ReLU()
+        self.conv2 = conv3x3(channels, channels, stride)
+        self.bn2 = nn.BatchNorm2D(channels)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet45(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        block=BasicBlock,
+        layers=[3, 4, 6, 6, 3],
+        strides=[2, 1, 2, 1, 1],
+    ):
+        self.inplanes = 32
+        super(ResNet45, self).__init__()
+        self.conv1 = nn.Conv2D(
+            in_channels,
+            32,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm2D(32)
+        self.relu = nn.ReLU()
+
+        self.layer1 = self._make_layer(block, 32, layers[0], stride=strides[0])
+        self.layer2 = self._make_layer(block, 64, layers[1], stride=strides[1])
+        self.layer3 = self._make_layer(block, 128, layers[2], stride=strides[2])
+        self.layer4 = self._make_layer(block, 256, layers[3], stride=strides[3])
+        self.layer5 = self._make_layer(block, 512, layers[4], stride=strides[4])
+        self.out_channels = 512
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            # downsample = True
+            downsample = nn.Sequential(
+                nn.Conv2D(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    weight_attr=ParamAttr(initializer=KaimingNormal()),
+                    bias_attr=False,
+                ),
+                nn.BatchNorm2D(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.layer5(x)
+        return x

ppocr/modeling/backbones/rec_resnet_aster.py

@@ -0,0 +1,141 @@
+# 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/resnet_aster.py
+"""
+import paddle
+import paddle.nn as nn
+
+import sys
+import math
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2D(
+        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias_attr=False
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2D(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias_attr=False
+    )
+
+
+def get_sinusoid_encoding(n_position, feat_dim, wave_length=10000):
+    # [n_position]
+    positions = paddle.arange(0, n_position)
+    # [feat_dim]
+    dim_range = paddle.arange(0, feat_dim)
+    dim_range = paddle.pow(wave_length, 2 * (dim_range // 2) / feat_dim)
+    # [n_position, feat_dim]
+    angles = paddle.unsqueeze(positions, axis=1) / paddle.unsqueeze(dim_range, axis=0)
+    angles = paddle.cast(angles, "float32")
+    angles[:, 0::2] = paddle.sin(angles[:, 0::2])
+    angles[:, 1::2] = paddle.cos(angles[:, 1::2])
+    return angles
+
+
+class AsterBlock(nn.Layer):
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(AsterBlock, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2D(planes)
+        self.relu = nn.ReLU()
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2D(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        return out
+
+
+class ResNet_ASTER(nn.Layer):
+    """For aster or crnn"""
+
+    def __init__(self, with_lstm=True, n_group=1, in_channels=3):
+        super(ResNet_ASTER, self).__init__()
+        self.with_lstm = with_lstm
+        self.n_group = n_group
+
+        self.layer0 = nn.Sequential(
+            nn.Conv2D(
+                in_channels,
+                32,
+                kernel_size=(3, 3),
+                stride=1,
+                padding=1,
+                bias_attr=False,
+            ),
+            nn.BatchNorm2D(32),
+            nn.ReLU(),
+        )
+
+        self.inplanes = 32
+        self.layer1 = self._make_layer(32, 3, [2, 2])  # [16, 50]
+        self.layer2 = self._make_layer(64, 4, [2, 2])  # [8, 25]
+        self.layer3 = self._make_layer(128, 6, [2, 1])  # [4, 25]
+        self.layer4 = self._make_layer(256, 6, [2, 1])  # [2, 25]
+        self.layer5 = self._make_layer(512, 3, [2, 1])  # [1, 25]
+
+        if with_lstm:
+            self.rnn = nn.LSTM(512, 256, direction="bidirect", num_layers=2)
+            self.out_channels = 2 * 256
+        else:
+            self.out_channels = 512
+
+    def _make_layer(self, planes, blocks, stride):
+        downsample = None
+        if stride != [1, 1] or self.inplanes != planes:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes, stride), nn.BatchNorm2D(planes)
+            )
+
+        layers = []
+        layers.append(AsterBlock(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes
+        for _ in range(1, blocks):
+            layers.append(AsterBlock(self.inplanes, planes))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x0 = self.layer0(x)
+        x1 = self.layer1(x0)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        x5 = self.layer5(x4)
+
+        cnn_feat = x5.squeeze(2)  # [N, c, w]
+        cnn_feat = paddle.transpose(cnn_feat, perm=[0, 2, 1])
+        if self.with_lstm:
+            rnn_feat, _ = self.rnn(cnn_feat)
+            return rnn_feat
+        else:
+            return cnn_feat

ppocr/modeling/backbones/rec_resnet_fpn.py

@@ -0,0 +1,317 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import paddle
+import numpy as np
+
+__all__ = ["ResNetFPN"]
+
+
+class ResNetFPN(nn.Layer):
+    def __init__(self, in_channels=1, layers=50, **kwargs):
+        super(ResNetFPN, self).__init__()
+        supported_layers = {
+            18: {"depth": [2, 2, 2, 2], "block_class": BasicBlock},
+            34: {"depth": [3, 4, 6, 3], "block_class": BasicBlock},
+            50: {"depth": [3, 4, 6, 3], "block_class": BottleneckBlock},
+            101: {"depth": [3, 4, 23, 3], "block_class": BottleneckBlock},
+            152: {"depth": [3, 8, 36, 3], "block_class": BottleneckBlock},
+        }
+        stride_list = [(2, 2), (2, 2), (1, 1), (1, 1)]
+        num_filters = [64, 128, 256, 512]
+        self.depth = supported_layers[layers]["depth"]
+        self.F = []
+        self.conv = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=7,
+            stride=2,
+            act="relu",
+            name="conv1",
+        )
+        self.block_list = []
+        in_ch = 64
+        if layers >= 50:
+            for block in range(len(self.depth)):
+                for i in range(self.depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+                    block_list = self.add_sublayer(
+                        "bottleneckBlock_{}_{}".format(block, i),
+                        BottleneckBlock(
+                            in_channels=in_ch,
+                            out_channels=num_filters[block],
+                            stride=stride_list[block] if i == 0 else 1,
+                            name=conv_name,
+                        ),
+                    )
+                    in_ch = num_filters[block] * 4
+                    self.block_list.append(block_list)
+                self.F.append(block_list)
+        else:
+            for block in range(len(self.depth)):
+                for i in range(self.depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    if i == 0 and block != 0:
+                        stride = (2, 1)
+                    else:
+                        stride = (1, 1)
+                    basic_block = self.add_sublayer(
+                        conv_name,
+                        BasicBlock(
+                            in_channels=in_ch,
+                            out_channels=num_filters[block],
+                            stride=stride_list[block] if i == 0 else 1,
+                            is_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    in_ch = basic_block.out_channels
+                    self.block_list.append(basic_block)
+        out_ch_list = [in_ch // 4, in_ch // 2, in_ch]
+        self.base_block = []
+        self.conv_trans = []
+        self.bn_block = []
+        for i in [-2, -3]:
+            in_channels = out_ch_list[i + 1] + out_ch_list[i]
+
+            self.base_block.append(
+                self.add_sublayer(
+                    "F_{}_base_block_0".format(i),
+                    nn.Conv2D(
+                        in_channels=in_channels,
+                        out_channels=out_ch_list[i],
+                        kernel_size=1,
+                        weight_attr=ParamAttr(trainable=True),
+                        bias_attr=ParamAttr(trainable=True),
+                    ),
+                )
+            )
+            self.base_block.append(
+                self.add_sublayer(
+                    "F_{}_base_block_1".format(i),
+                    nn.Conv2D(
+                        in_channels=out_ch_list[i],
+                        out_channels=out_ch_list[i],
+                        kernel_size=3,
+                        padding=1,
+                        weight_attr=ParamAttr(trainable=True),
+                        bias_attr=ParamAttr(trainable=True),
+                    ),
+                )
+            )
+            self.base_block.append(
+                self.add_sublayer(
+                    "F_{}_base_block_2".format(i),
+                    nn.BatchNorm(
+                        num_channels=out_ch_list[i],
+                        act="relu",
+                        param_attr=ParamAttr(trainable=True),
+                        bias_attr=ParamAttr(trainable=True),
+                    ),
+                )
+            )
+        self.base_block.append(
+            self.add_sublayer(
+                "F_{}_base_block_3".format(i),
+                nn.Conv2D(
+                    in_channels=out_ch_list[i],
+                    out_channels=512,
+                    kernel_size=1,
+                    bias_attr=ParamAttr(trainable=True),
+                    weight_attr=ParamAttr(trainable=True),
+                ),
+            )
+        )
+        self.out_channels = 512
+
+    def __call__(self, x):
+        x = self.conv(x)
+        fpn_list = []
+        F = []
+        for i in range(len(self.depth)):
+            fpn_list.append(np.sum(self.depth[: i + 1]))
+
+        for i, block in enumerate(self.block_list):
+            x = block(x)
+            for number in fpn_list:
+                if i + 1 == number:
+                    F.append(x)
+        base = F[-1]
+
+        j = 0
+        for i, block in enumerate(self.base_block):
+            if i % 3 == 0 and i < 6:
+                j = j + 1
+                b, c, w, h = F[-j - 1].shape
+                if [w, h] == list(base.shape[2:]):
+                    base = base
+                else:
+                    base = self.conv_trans[j - 1](base)
+                    base = self.bn_block[j - 1](base)
+                base = paddle.concat([base, F[-j - 1]], axis=1)
+            base = block(base)
+        return base
+
+
+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=2 if stride == (1, 1) else kernel_size,
+            dilation=2 if stride == (1, 1) else 1,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + ".conv2d.output.1.w_0"),
+            bias_attr=False,
+        )
+
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name=name + ".output.1.w_0"),
+            bias_attr=ParamAttr(name=name + ".output.1.b_0"),
+            moving_mean_name=bn_name + "_mean",
+            moving_variance_name=bn_name + "_variance",
+        )
+
+    def __call__(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class ShortCut(nn.Layer):
+    def __init__(self, in_channels, out_channels, stride, name, is_first=False):
+        super(ShortCut, self).__init__()
+        self.use_conv = True
+
+        if in_channels != out_channels or stride != 1 or is_first == True:
+            if stride == (1, 1):
+                self.conv = ConvBNLayer(in_channels, out_channels, 1, 1, name=name)
+            else:  # stride==(2,2)
+                self.conv = ConvBNLayer(in_channels, out_channels, 1, stride, name=name)
+        else:
+            self.use_conv = False
+
+    def forward(self, x):
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(self, in_channels, out_channels, stride, name):
+        super(BottleneckBlock, self).__init__()
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2b",
+        )
+
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+            name=name + "_branch2c",
+        )
+
+        self.short = ShortCut(
+            in_channels=in_channels,
+            out_channels=out_channels * 4,
+            stride=stride,
+            is_first=False,
+            name=name + "_branch1",
+        )
+        self.out_channels = out_channels * 4
+
+    def forward(self, x):
+        y = self.conv0(x)
+        y = self.conv1(y)
+        y = self.conv2(y)
+        y = y + self.short(x)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self, in_channels, out_channels, stride, name, is_first):
+        super(BasicBlock, self).__init__()
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act="relu",
+            stride=stride,
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None,
+            name=name + "_branch2b",
+        )
+        self.short = ShortCut(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            stride=stride,
+            is_first=is_first,
+            name=name + "_branch1",
+        )
+        self.out_channels = out_channels
+
+    def forward(self, x):
+        y = self.conv0(x)
+        y = self.conv1(y)
+        y = y + self.short(x)
+        return F.relu(y)

ppocr/modeling/backbones/rec_resnet_rfl.py

@@ -0,0 +1,359 @@
+# 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/hikopensource/DAVAR-Lab-OCR/blob/main/davarocr/davar_rcg/models/backbones/ResNetRFL.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+import paddle.nn as nn
+
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal, KaimingNormal
+
+kaiming_init_ = KaimingNormal()
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+class BasicBlock(nn.Layer):
+    """Res-net Basic Block"""
+
+    expansion = 1
+
+    def __init__(
+        self, inplanes, planes, stride=1, downsample=None, norm_type="BN", **kwargs
+    ):
+        """
+        Args:
+            inplanes (int): input channel
+            planes (int): channels of the middle feature
+            stride (int): stride of the convolution
+            downsample (int): type of the down_sample
+            norm_type (str): type of the normalization
+            **kwargs (None): backup parameter
+        """
+        super(BasicBlock, self).__init__()
+        self.conv1 = self._conv3x3(inplanes, planes)
+        self.bn1 = nn.BatchNorm(planes)
+        self.conv2 = self._conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm(planes)
+        self.relu = nn.ReLU()
+        self.downsample = downsample
+        self.stride = stride
+
+    def _conv3x3(self, in_planes, out_planes, stride=1):
+        return nn.Conv2D(
+            in_planes,
+            out_planes,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            bias_attr=False,
+        )
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetRFL(nn.Layer):
+    def __init__(self, in_channels, out_channels=512, use_cnt=True, use_seq=True):
+        """
+
+        Args:
+            in_channels (int): input channel
+            out_channels (int): output channel
+        """
+        super(ResNetRFL, self).__init__()
+        assert use_cnt or use_seq
+        self.use_cnt, self.use_seq = use_cnt, use_seq
+        self.backbone = RFLBase(in_channels)
+
+        self.out_channels = out_channels
+        self.out_channels_block = [
+            int(self.out_channels / 4),
+            int(self.out_channels / 2),
+            self.out_channels,
+            self.out_channels,
+        ]
+        block = BasicBlock
+        layers = [1, 2, 5, 3]
+        self.inplanes = int(self.out_channels // 2)
+
+        self.relu = nn.ReLU()
+        if self.use_seq:
+            self.maxpool3 = nn.MaxPool2D(kernel_size=2, stride=(2, 1), padding=(0, 1))
+            self.layer3 = self._make_layer(
+                block, self.out_channels_block[2], layers[2], stride=1
+            )
+            self.conv3 = nn.Conv2D(
+                self.out_channels_block[2],
+                self.out_channels_block[2],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias_attr=False,
+            )
+            self.bn3 = nn.BatchNorm(self.out_channels_block[2])
+
+            self.layer4 = self._make_layer(
+                block, self.out_channels_block[3], layers[3], stride=1
+            )
+            self.conv4_1 = nn.Conv2D(
+                self.out_channels_block[3],
+                self.out_channels_block[3],
+                kernel_size=2,
+                stride=(2, 1),
+                padding=(0, 1),
+                bias_attr=False,
+            )
+            self.bn4_1 = nn.BatchNorm(self.out_channels_block[3])
+            self.conv4_2 = nn.Conv2D(
+                self.out_channels_block[3],
+                self.out_channels_block[3],
+                kernel_size=2,
+                stride=1,
+                padding=0,
+                bias_attr=False,
+            )
+            self.bn4_2 = nn.BatchNorm(self.out_channels_block[3])
+
+        if self.use_cnt:
+            self.inplanes = int(self.out_channels // 2)
+            self.v_maxpool3 = nn.MaxPool2D(kernel_size=2, stride=(2, 1), padding=(0, 1))
+            self.v_layer3 = self._make_layer(
+                block, self.out_channels_block[2], layers[2], stride=1
+            )
+            self.v_conv3 = nn.Conv2D(
+                self.out_channels_block[2],
+                self.out_channels_block[2],
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias_attr=False,
+            )
+            self.v_bn3 = nn.BatchNorm(self.out_channels_block[2])
+
+            self.v_layer4 = self._make_layer(
+                block, self.out_channels_block[3], layers[3], stride=1
+            )
+            self.v_conv4_1 = nn.Conv2D(
+                self.out_channels_block[3],
+                self.out_channels_block[3],
+                kernel_size=2,
+                stride=(2, 1),
+                padding=(0, 1),
+                bias_attr=False,
+            )
+            self.v_bn4_1 = nn.BatchNorm(self.out_channels_block[3])
+            self.v_conv4_2 = nn.Conv2D(
+                self.out_channels_block[3],
+                self.out_channels_block[3],
+                kernel_size=2,
+                stride=1,
+                padding=0,
+                bias_attr=False,
+            )
+            self.v_bn4_2 = nn.BatchNorm(self.out_channels_block[3])
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2D(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias_attr=False,
+                ),
+                nn.BatchNorm(planes * block.expansion),
+            )
+
+        layers = list()
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, inputs):
+        x_1 = self.backbone(inputs)
+
+        if self.use_cnt:
+            v_x = self.v_maxpool3(x_1)
+            v_x = self.v_layer3(v_x)
+            v_x = self.v_conv3(v_x)
+            v_x = self.v_bn3(v_x)
+            visual_feature_2 = self.relu(v_x)
+
+            v_x = self.v_layer4(visual_feature_2)
+            v_x = self.v_conv4_1(v_x)
+            v_x = self.v_bn4_1(v_x)
+            v_x = self.relu(v_x)
+            v_x = self.v_conv4_2(v_x)
+            v_x = self.v_bn4_2(v_x)
+            visual_feature_3 = self.relu(v_x)
+        else:
+            visual_feature_3 = None
+        if self.use_seq:
+            x = self.maxpool3(x_1)
+            x = self.layer3(x)
+            x = self.conv3(x)
+            x = self.bn3(x)
+            x_2 = self.relu(x)
+
+            x = self.layer4(x_2)
+            x = self.conv4_1(x)
+            x = self.bn4_1(x)
+            x = self.relu(x)
+            x = self.conv4_2(x)
+            x = self.bn4_2(x)
+            x_3 = self.relu(x)
+        else:
+            x_3 = None
+
+        return [visual_feature_3, x_3]
+
+
+class ResNetBase(nn.Layer):
+    def __init__(self, in_channels, out_channels, block, layers):
+        super(ResNetBase, self).__init__()
+
+        self.out_channels_block = [
+            int(out_channels / 4),
+            int(out_channels / 2),
+            out_channels,
+            out_channels,
+        ]
+
+        self.inplanes = int(out_channels / 8)
+        self.conv0_1 = nn.Conv2D(
+            in_channels,
+            int(out_channels / 16),
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn0_1 = nn.BatchNorm(int(out_channels / 16))
+        self.conv0_2 = nn.Conv2D(
+            int(out_channels / 16),
+            self.inplanes,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn0_2 = nn.BatchNorm(self.inplanes)
+        self.relu = nn.ReLU()
+
+        self.maxpool1 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.layer1 = self._make_layer(block, self.out_channels_block[0], layers[0])
+        self.conv1 = nn.Conv2D(
+            self.out_channels_block[0],
+            self.out_channels_block[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm(self.out_channels_block[0])
+
+        self.maxpool2 = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.layer2 = self._make_layer(
+            block, self.out_channels_block[1], layers[1], stride=1
+        )
+        self.conv2 = nn.Conv2D(
+            self.out_channels_block[1],
+            self.out_channels_block[1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            bias_attr=False,
+        )
+        self.bn2 = nn.BatchNorm(self.out_channels_block[1])
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2D(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias_attr=False,
+                ),
+                nn.BatchNorm(planes * block.expansion),
+            )
+
+        layers = list()
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv0_1(x)
+        x = self.bn0_1(x)
+        x = self.relu(x)
+        x = self.conv0_2(x)
+        x = self.bn0_2(x)
+        x = self.relu(x)
+
+        x = self.maxpool1(x)
+        x = self.layer1(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.maxpool2(x)
+        x = self.layer2(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        return x
+
+
+class RFLBase(nn.Layer):
+    """Reciprocal feature learning share backbone network"""
+
+    def __init__(self, in_channels, out_channels=512):
+        super(RFLBase, self).__init__()
+        self.ConvNet = ResNetBase(in_channels, out_channels, BasicBlock, [1, 2, 5, 3])
+
+    def forward(self, inputs):
+        return self.ConvNet(inputs)

ppocr/modeling/backbones/rec_resnet_vd.py

@@ -0,0 +1,313 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+__all__ = ["ResNet"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        is_vd_mode=False,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=stride, stride=stride, padding=0, ceil_mode=True
+        )
+        self._conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=1 if is_vd_mode else stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = 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, inputs):
+        if self.is_vd_mode:
+            inputs = self._pool2d_avg(inputs)
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2b",
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+            name=name + "_branch2c",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=stride,
+                is_vd_mode=not if_first and stride[0] != 1,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        shortcut=True,
+        if_first=False,
+        name=None,
+    ):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act="relu",
+            name=name + "_branch2a",
+        )
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None,
+            name=name + "_branch2b",
+        )
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=stride,
+                is_vd_mode=not if_first and stride[0] != 1,
+                name=name + "_branch1",
+            )
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet(nn.Layer):
+    def __init__(self, in_channels=3, layers=50, **kwargs):
+        super(ResNet, self).__init__()
+
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert (
+            layers in supported_layers
+        ), "supported layers are {} but input layer is {}".format(
+            supported_layers, layers
+        )
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            depth = [3, 4, 6, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        num_channels = [64, 256, 512, 1024] if layers >= 50 else [64, 64, 128, 256]
+        num_filters = [64, 128, 256, 512]
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv1_1",
+        )
+        self.conv1_2 = ConvBNLayer(
+            in_channels=32,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv1_2",
+        )
+        self.conv1_3 = ConvBNLayer(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv1_3",
+        )
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        self.block_list = []
+        if layers >= 50:
+            for block in range(len(depth)):
+                shortcut = False
+                for i in range(depth[block]):
+                    if layers in [101, 152, 200] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+
+                    if i == 0 and block != 0:
+                        stride = (2, 1)
+                    else:
+                        stride = (1, 1)
+                    bottleneck_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BottleneckBlock(
+                            in_channels=(
+                                num_channels[block]
+                                if i == 0
+                                else num_filters[block] * 4
+                            ),
+                            out_channels=num_filters[block],
+                            stride=stride,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    self.block_list.append(bottleneck_block)
+                self.out_channels = num_filters[block] * 4
+        else:
+            for block in range(len(depth)):
+                shortcut = False
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    if i == 0 and block != 0:
+                        stride = (2, 1)
+                    else:
+                        stride = (1, 1)
+
+                    basic_block = self.add_sublayer(
+                        "bb_%d_%d" % (block, i),
+                        BasicBlock(
+                            in_channels=(
+                                num_channels[block] if i == 0 else num_filters[block]
+                            ),
+                            out_channels=num_filters[block],
+                            stride=stride,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name,
+                        ),
+                    )
+                    shortcut = True
+                    self.block_list.append(basic_block)
+                self.out_channels = num_filters[block]
+        self.out_pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+
+    def forward(self, inputs):
+        y = self.conv1_1(inputs)
+        y = self.conv1_2(y)
+        y = self.conv1_3(y)
+        y = self.pool2d_max(y)
+        for block in self.block_list:
+            y = block(y)
+        y = self.out_pool(y)
+        return y

ppocr/modeling/backbones/rec_shallow_cnn.py

@@ -0,0 +1,82 @@
+# 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/open-mmlab/mmocr/blob/1.x/mmocr/models/textrecog/backbones/shallow_cnn.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import numpy as np
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn import MaxPool2D
+from paddle.nn.initializer import KaimingNormal, Uniform, Constant
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self, num_channels, filter_size, num_filters, stride, padding, num_groups=1
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.conv = nn.Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=num_groups,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm2D(
+            num_filters,
+            weight_attr=ParamAttr(initializer=Uniform(0, 1)),
+            bias_attr=ParamAttr(initializer=Constant(0)),
+        )
+        self.relu = nn.ReLU()
+
+    def forward(self, inputs):
+        y = self.conv(inputs)
+        y = self.bn(y)
+        y = self.relu(y)
+        return y
+
+
+class ShallowCNN(nn.Layer):
+    def __init__(self, in_channels=1, hidden_dim=512):
+        super().__init__()
+        assert isinstance(in_channels, int)
+        assert isinstance(hidden_dim, int)
+
+        self.conv1 = ConvBNLayer(in_channels, 3, hidden_dim // 2, stride=1, padding=1)
+        self.conv2 = ConvBNLayer(hidden_dim // 2, 3, hidden_dim, stride=1, padding=1)
+        self.pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+        self.out_channels = hidden_dim
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.pool(x)
+
+        x = self.conv2(x)
+        x = self.pool(x)
+
+        return x

ppocr/modeling/backbones/rec_svtrnet.py

@@ -0,0 +1,642 @@
+# 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 paddle import ParamAttr
+from paddle.nn.initializer import KaimingNormal
+import numpy as np
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+trunc_normal_ = TruncatedNormal(std=0.02)
+normal_ = Normal
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+def drop_path(x, drop_prob=0.0, training=False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ...
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = paddle.to_tensor(1 - drop_prob, dtype=x.dtype)
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + paddle.rand(shape, dtype=x.dtype)
+    random_tensor = paddle.floor(random_tensor)  # binarize
+    output = x.divide(keep_prob) * random_tensor
+    return output
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=0,
+        bias_attr=False,
+        groups=1,
+        act=nn.GELU,
+    ):
+        super().__init__()
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=paddle.ParamAttr(initializer=nn.initializer.KaimingUniform()),
+            bias_attr=bias_attr,
+        )
+        self.norm = nn.BatchNorm2D(out_channels)
+        self.act = act()
+
+    def forward(self, inputs):
+        out = self.conv(inputs)
+        out = self.norm(out)
+        out = self.act(out)
+        return out
+
+
+class DropPath(nn.Layer):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Identity(nn.Layer):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Layer):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class ConvMixer(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        HW=[8, 25],
+        local_k=[3, 3],
+    ):
+        super().__init__()
+        self.HW = HW
+        self.dim = dim
+        self.local_mixer = nn.Conv2D(
+            dim,
+            dim,
+            local_k,
+            1,
+            [local_k[0] // 2, local_k[1] // 2],
+            groups=num_heads,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+        )
+
+    def forward(self, x):
+        h = self.HW[0]
+        w = self.HW[1]
+        x = x.transpose([0, 2, 1]).reshape([0, self.dim, h, w])
+        x = self.local_mixer(x)
+        x = x.flatten(2).transpose([0, 2, 1])
+        return x
+
+
+class Attention(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        mixer="Global",
+        HW=None,
+        local_k=[7, 11],
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.HW = HW
+        if HW is not None:
+            H = HW[0]
+            W = HW[1]
+            self.N = H * W
+            self.C = dim
+        if mixer == "Local" and HW is not None:
+            hk = local_k[0]
+            wk = local_k[1]
+            mask = paddle.ones([H * W, H + hk - 1, W + wk - 1], dtype="float32")
+            for h in range(0, H):
+                for w in range(0, W):
+                    mask[h * W + w, h : h + hk, w : w + wk] = 0.0
+            mask_paddle = mask[:, hk // 2 : H + hk // 2, wk // 2 : W + wk // 2].flatten(
+                1
+            )
+            mask_inf = paddle.full([H * W, H * W], "-inf", dtype="float32")
+            mask = paddle.where(mask_paddle < 1, mask_paddle, mask_inf)
+            self.mask = mask.unsqueeze([0, 1])
+        self.mixer = mixer
+
+    def forward(self, x):
+        qkv = (
+            self.qkv(x)
+            .reshape((0, -1, 3, self.num_heads, self.head_dim))
+            .transpose((2, 0, 3, 1, 4))
+        )
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+
+        attn = q.matmul(k.transpose((0, 1, 3, 2)))
+        if self.mixer == "Local":
+            attn += self.mask
+        attn = nn.functional.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((0, -1, self.dim))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mixer="Global",
+        local_mixer=[7, 11],
+        HW=None,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-6,
+        prenorm=True,
+    ):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
+        else:
+            self.norm1 = norm_layer(dim)
+        if mixer == "Global" or mixer == "Local":
+            self.mixer = Attention(
+                dim,
+                num_heads=num_heads,
+                mixer=mixer,
+                HW=HW,
+                local_k=local_mixer,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                attn_drop=attn_drop,
+                proj_drop=drop,
+            )
+        elif mixer == "Conv":
+            self.mixer = ConvMixer(dim, num_heads=num_heads, HW=HW, local_k=local_mixer)
+        else:
+            raise TypeError("The mixer must be one of [Global, Local, Conv]")
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        if isinstance(norm_layer, str):
+            self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
+        else:
+            self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.prenorm = prenorm
+
+    def forward(self, x):
+        if self.prenorm:
+            x = self.norm1(x + self.drop_path(self.mixer(x)))
+            x = self.norm2(x + self.drop_path(self.mlp(x)))
+        else:
+            x = x + self.drop_path(self.mixer(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Layer):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=[32, 100],
+        in_channels=3,
+        embed_dim=768,
+        sub_num=2,
+        patch_size=[4, 4],
+        mode="pope",
+    ):
+        super().__init__()
+        num_patches = (img_size[1] // (2**sub_num)) * (img_size[0] // (2**sub_num))
+        self.img_size = img_size
+        self.num_patches = num_patches
+        self.embed_dim = embed_dim
+        self.norm = None
+        if mode == "pope":
+            if sub_num == 2:
+                self.proj = nn.Sequential(
+                    ConvBNLayer(
+                        in_channels=in_channels,
+                        out_channels=embed_dim // 2,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act=nn.GELU,
+                        bias_attr=None,
+                    ),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 2,
+                        out_channels=embed_dim,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act=nn.GELU,
+                        bias_attr=None,
+                    ),
+                )
+            if sub_num == 3:
+                self.proj = nn.Sequential(
+                    ConvBNLayer(
+                        in_channels=in_channels,
+                        out_channels=embed_dim // 4,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act=nn.GELU,
+                        bias_attr=None,
+                    ),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 4,
+                        out_channels=embed_dim // 2,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act=nn.GELU,
+                        bias_attr=None,
+                    ),
+                    ConvBNLayer(
+                        in_channels=embed_dim // 2,
+                        out_channels=embed_dim,
+                        kernel_size=3,
+                        stride=2,
+                        padding=1,
+                        act=nn.GELU,
+                        bias_attr=None,
+                    ),
+                )
+        elif mode == "linear":
+            self.proj = nn.Conv2D(
+                1, embed_dim, kernel_size=patch_size, stride=patch_size
+            )
+            self.num_patches = (
+                img_size[0] // patch_size[0] * img_size[1] // patch_size[1]
+            )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose((0, 2, 1))
+        return x
+
+
+class SubSample(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        types="Pool",
+        stride=[2, 1],
+        sub_norm="nn.LayerNorm",
+        act=None,
+    ):
+        super().__init__()
+        self.types = types
+        if types == "Pool":
+            self.avgpool = nn.AvgPool2D(
+                kernel_size=[3, 5], stride=stride, padding=[1, 2]
+            )
+            self.maxpool = nn.MaxPool2D(
+                kernel_size=[3, 5], stride=stride, padding=[1, 2]
+            )
+            self.proj = nn.Linear(in_channels, out_channels)
+        else:
+            self.conv = nn.Conv2D(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=stride,
+                padding=1,
+                weight_attr=ParamAttr(initializer=KaimingNormal()),
+            )
+        self.norm = eval(sub_norm)(out_channels)
+        if act is not None:
+            self.act = act()
+        else:
+            self.act = None
+
+    def forward(self, x):
+        if self.types == "Pool":
+            x1 = self.avgpool(x)
+            x2 = self.maxpool(x)
+            x = (x1 + x2) * 0.5
+            out = self.proj(x.flatten(2).transpose((0, 2, 1)))
+        else:
+            x = self.conv(x)
+            out = x.flatten(2).transpose((0, 2, 1))
+        out = self.norm(out)
+        if self.act is not None:
+            out = self.act(out)
+
+        return out
+
+
+class SVTRNet(nn.Layer):
+    def __init__(
+        self,
+        img_size=[32, 100],
+        in_channels=3,
+        embed_dim=[64, 128, 256],
+        depth=[3, 6, 3],
+        num_heads=[2, 4, 8],
+        mixer=["Local"] * 6 + ["Global"] * 6,  # Local atten, Global atten, Conv
+        local_mixer=[[7, 11], [7, 11], [7, 11]],
+        patch_merging="Conv",  # Conv, Pool, None
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        last_drop=0.1,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        norm_layer="nn.LayerNorm",
+        sub_norm="nn.LayerNorm",
+        epsilon=1e-6,
+        out_channels=192,
+        out_char_num=25,
+        block_unit="Block",
+        act="nn.GELU",
+        last_stage=True,
+        sub_num=2,
+        prenorm=True,
+        use_lenhead=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.embed_dim = embed_dim
+        self.out_channels = out_channels
+        self.prenorm = prenorm
+        patch_merging = (
+            None
+            if patch_merging != "Conv" and patch_merging != "Pool"
+            else patch_merging
+        )
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            in_channels=in_channels,
+            embed_dim=embed_dim[0],
+            sub_num=sub_num,
+        )
+        num_patches = self.patch_embed.num_patches
+        self.HW = [img_size[0] // (2**sub_num), img_size[1] // (2**sub_num)]
+        self.pos_embed = self.create_parameter(
+            shape=[1, num_patches, embed_dim[0]], default_initializer=zeros_
+        )
+        self.add_parameter("pos_embed", self.pos_embed)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        Block_unit = eval(block_unit)
+
+        dpr = np.linspace(0, drop_path_rate, sum(depth))
+        self.blocks1 = nn.LayerList(
+            [
+                Block_unit(
+                    dim=embed_dim[0],
+                    num_heads=num_heads[0],
+                    mixer=mixer[0 : depth[0]][i],
+                    HW=self.HW,
+                    local_mixer=local_mixer[0],
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    act_layer=eval(act),
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[0 : depth[0]][i],
+                    norm_layer=norm_layer,
+                    epsilon=epsilon,
+                    prenorm=prenorm,
+                )
+                for i in range(depth[0])
+            ]
+        )
+        if patch_merging is not None:
+            self.sub_sample1 = SubSample(
+                embed_dim[0],
+                embed_dim[1],
+                sub_norm=sub_norm,
+                stride=[2, 1],
+                types=patch_merging,
+            )
+            HW = [self.HW[0] // 2, self.HW[1]]
+        else:
+            HW = self.HW
+        self.patch_merging = patch_merging
+        self.blocks2 = nn.LayerList(
+            [
+                Block_unit(
+                    dim=embed_dim[1],
+                    num_heads=num_heads[1],
+                    mixer=mixer[depth[0] : depth[0] + depth[1]][i],
+                    HW=HW,
+                    local_mixer=local_mixer[1],
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    act_layer=eval(act),
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[depth[0] : depth[0] + depth[1]][i],
+                    norm_layer=norm_layer,
+                    epsilon=epsilon,
+                    prenorm=prenorm,
+                )
+                for i in range(depth[1])
+            ]
+        )
+        if patch_merging is not None:
+            self.sub_sample2 = SubSample(
+                embed_dim[1],
+                embed_dim[2],
+                sub_norm=sub_norm,
+                stride=[2, 1],
+                types=patch_merging,
+            )
+            HW = [self.HW[0] // 4, self.HW[1]]
+        else:
+            HW = self.HW
+        self.blocks3 = nn.LayerList(
+            [
+                Block_unit(
+                    dim=embed_dim[2],
+                    num_heads=num_heads[2],
+                    mixer=mixer[depth[0] + depth[1] :][i],
+                    HW=HW,
+                    local_mixer=local_mixer[2],
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    act_layer=eval(act),
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[depth[0] + depth[1] :][i],
+                    norm_layer=norm_layer,
+                    epsilon=epsilon,
+                    prenorm=prenorm,
+                )
+                for i in range(depth[2])
+            ]
+        )
+        self.last_stage = last_stage
+        if last_stage:
+            self.avg_pool = nn.AdaptiveAvgPool2D([1, out_char_num])
+            self.last_conv = nn.Conv2D(
+                in_channels=embed_dim[2],
+                out_channels=self.out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                bias_attr=False,
+            )
+            self.hardswish = nn.Hardswish()
+            self.dropout = nn.Dropout(p=last_drop, mode="downscale_in_infer")
+        if not prenorm:
+            self.norm = eval(norm_layer)(embed_dim[-1], epsilon=epsilon)
+        self.use_lenhead = use_lenhead
+        if use_lenhead:
+            self.len_conv = nn.Linear(embed_dim[2], self.out_channels)
+            self.hardswish_len = nn.Hardswish()
+            self.dropout_len = nn.Dropout(p=last_drop, mode="downscale_in_infer")
+
+        trunc_normal_(self.pos_embed)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        for blk in self.blocks1:
+            x = blk(x)
+        if self.patch_merging is not None:
+            x = self.sub_sample1(
+                x.transpose([0, 2, 1]).reshape(
+                    [0, self.embed_dim[0], self.HW[0], self.HW[1]]
+                )
+            )
+        for blk in self.blocks2:
+            x = blk(x)
+        if self.patch_merging is not None:
+            x = self.sub_sample2(
+                x.transpose([0, 2, 1]).reshape(
+                    [0, self.embed_dim[1], self.HW[0] // 2, self.HW[1]]
+                )
+            )
+        for blk in self.blocks3:
+            x = blk(x)
+        if not self.prenorm:
+            x = self.norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        if self.use_lenhead:
+            len_x = self.len_conv(x.mean(1))
+            len_x = self.dropout_len(self.hardswish_len(len_x))
+        if self.last_stage:
+            if self.patch_merging is not None:
+                h = self.HW[0] // 4
+            else:
+                h = self.HW[0]
+            x = self.avg_pool(
+                x.transpose([0, 2, 1]).reshape([0, self.embed_dim[2], h, self.HW[1]])
+            )
+            x = self.last_conv(x)
+            x = self.hardswish(x)
+            x = self.dropout(x)
+        if self.use_lenhead:
+            return x, len_x
+        return x

ppocr/modeling/backbones/rec_svtrv2.py

@@ -0,0 +1,575 @@
+# copyright (c) 2024 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 paddle import ParamAttr
+from paddle.nn.initializer import KaimingNormal
+import numpy as np
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+trunc_normal_ = TruncatedNormal(std=0.02)
+normal_ = Normal
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+def drop_path(x, drop_prob=0.0, training=False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ...
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = paddle.to_tensor(1 - drop_prob, dtype=x.dtype)
+    shape = (paddle.shape(x)[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + paddle.rand(shape, dtype=x.dtype)
+    random_tensor = paddle.floor(random_tensor)  # binarize
+    output = x.divide(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Layer):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Identity(nn.Layer):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Layer):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=0,
+        bias_attr=False,
+        groups=1,
+        act=nn.GELU,
+    ):
+        super().__init__()
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=paddle.ParamAttr(initializer=nn.initializer.KaimingUniform()),
+            bias_attr=bias_attr,
+        )
+        self.norm = nn.BatchNorm2D(out_channels)
+        self.act = act()
+
+    def forward(self, inputs):
+        out = self.conv(inputs)
+        out = self.norm(out)
+        out = self.act(out)
+        return out
+
+
+class Attention(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        qkv = (
+            self.qkv(x)
+            .reshape((0, -1, 3, self.num_heads, self.head_dim))
+            .transpose((2, 0, 3, 1, 4))
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
+        attn = nn.functional.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((0, -1, self.dim))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        epsilon=1e-6,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim, epsilon=epsilon)
+        self.mixer = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        self.norm2 = norm_layer(dim, epsilon=epsilon)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        x = self.norm1(x + self.drop_path(self.mixer(x)))
+        x = self.norm2(x + self.drop_path(self.mlp(x)))
+        return x
+
+
+class ConvBlock(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        epsilon=1e-6,
+    ):
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.norm1 = norm_layer(dim, epsilon=epsilon)
+        self.mixer = nn.Conv2D(
+            dim,
+            dim,
+            5,
+            1,
+            2,
+            groups=num_heads,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        self.norm2 = norm_layer(dim, epsilon=epsilon)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        C, H, W = x.shape[1:]
+        x = x + self.drop_path(self.mixer(x))
+        x = self.norm1(x.flatten(2).transpose([0, 2, 1]))
+        x = self.norm2(x + self.drop_path(self.mlp(x)))
+        x = x.transpose([0, 2, 1]).reshape([0, C, H, W])
+        return x
+
+
+class FlattenTranspose(nn.Layer):
+    def forward(self, x):
+        return x.flatten(2).transpose([0, 2, 1])
+
+
+class SubSample2D(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride=[2, 1],
+    ):
+        super().__init__()
+        self.conv = nn.Conv2D(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+        )
+        self.norm = nn.LayerNorm(out_channels)
+
+    def forward(self, x, sz):
+        # print(x.shape)
+        x = self.conv(x)
+        C, H, W = x.shape[1:]
+        x = self.norm(x.flatten(2).transpose([0, 2, 1]))
+        x = x.transpose([0, 2, 1]).reshape([0, C, H, W])
+        return x, [H, W]
+
+
+class SubSample1D(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride=[2, 1],
+    ):
+        super().__init__()
+        self.conv = nn.Conv2D(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=stride,
+            padding=1,
+            weight_attr=ParamAttr(initializer=KaimingNormal()),
+        )
+        self.norm = nn.LayerNorm(out_channels)
+
+    def forward(self, x, sz):
+        C = x.shape[-1]
+        x = x.transpose([0, 2, 1]).reshape([0, C, sz[0], sz[1]])
+        x = self.conv(x)
+        C, H, W = x.shape[1:]
+        x = self.norm(x.flatten(2).transpose([0, 2, 1]))
+        return x, [H, W]
+
+
+class IdentitySize(nn.Layer):
+    def forward(self, x, sz):
+        return x, sz
+
+
+class SVTRStage(nn.Layer):
+    def __init__(
+        self,
+        dim=64,
+        out_dim=256,
+        depth=3,
+        mixer=["Local"] * 3,
+        sub_k=[2, 1],
+        num_heads=2,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path=[0.1] * 3,
+        norm_layer=nn.LayerNorm,
+        act=nn.GELU,
+        eps=1e-6,
+        downsample=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.dim = dim
+
+        conv_block_num = sum([1 if mix == "Conv" else 0 for mix in mixer])
+        blocks = []
+        for i in range(depth):
+            if mixer[i] == "Conv":
+                blocks.append(
+                    ConvBlock(
+                        dim=dim,
+                        num_heads=num_heads,
+                        mlp_ratio=mlp_ratio,
+                        drop=drop_rate,
+                        act_layer=act,
+                        drop_path=drop_path[i],
+                        norm_layer=norm_layer,
+                        epsilon=eps,
+                    )
+                )
+            else:
+                blocks.append(
+                    Block(
+                        dim=dim,
+                        num_heads=num_heads,
+                        mlp_ratio=mlp_ratio,
+                        qkv_bias=qkv_bias,
+                        qk_scale=qk_scale,
+                        drop=drop_rate,
+                        act_layer=act,
+                        attn_drop=attn_drop_rate,
+                        drop_path=drop_path[i],
+                        norm_layer=norm_layer,
+                        epsilon=eps,
+                    )
+                )
+            if i == conv_block_num - 1 and mixer[-1] != "Conv":
+                blocks.append(FlattenTranspose())
+        self.blocks = nn.Sequential(*blocks)
+        if downsample:
+            if mixer[-1] == "Conv":
+                self.downsample = SubSample2D(dim, out_dim, stride=sub_k)
+            elif mixer[-1] == "Global":
+                self.downsample = SubSample1D(dim, out_dim, stride=sub_k)
+        else:
+            self.downsample = IdentitySize()
+
+    def forward(self, x, sz):
+        x = self.blocks(x)
+        x, sz = self.downsample(x, sz)
+        return x, sz
+
+
+class ADDPosEmbed(nn.Layer):
+    def __init__(self, feat_max_size=[8, 32], embed_dim=768):
+        super().__init__()
+        pos_embed = paddle.zeros(
+            [1, feat_max_size[0] * feat_max_size[1], embed_dim], dtype=paddle.float32
+        )
+        trunc_normal_(pos_embed)
+        pos_embed = pos_embed.transpose([0, 2, 1]).reshape(
+            [1, embed_dim, feat_max_size[0], feat_max_size[1]]
+        )
+        self.pos_embed = self.create_parameter(
+            [1, embed_dim, feat_max_size[0], feat_max_size[1]]
+        )
+        self.add_parameter("pos_embed", self.pos_embed)
+        self.pos_embed.set_value(pos_embed)
+
+    def forward(self, x):
+        sz = x.shape[2:]
+        x = x + self.pos_embed[:, :, : sz[0], : sz[1]]
+        return x
+
+
+class POPatchEmbed(nn.Layer):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        in_channels=3,
+        feat_max_size=[8, 32],
+        embed_dim=768,
+        use_pos_embed=False,
+        flatten=False,
+    ):
+        super().__init__()
+        patch_embed = [
+            ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=embed_dim // 2,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                act=nn.GELU,
+                bias_attr=None,
+            ),
+            ConvBNLayer(
+                in_channels=embed_dim // 2,
+                out_channels=embed_dim,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                act=nn.GELU,
+                bias_attr=None,
+            ),
+        ]
+        if use_pos_embed:
+            patch_embed.append(ADDPosEmbed(feat_max_size, embed_dim))
+        if flatten:
+            patch_embed.append(FlattenTranspose())
+        self.patch_embed = nn.Sequential(*patch_embed)
+
+    def forward(self, x):
+        sz = x.shape[2:]
+        x = self.patch_embed(x)
+        return x, [sz[0] // 4, sz[1] // 4]
+
+
+class LastStage(nn.Layer):
+    def __init__(self, in_channels, out_channels, last_drop, out_char_num):
+        super().__init__()
+        self.last_conv = nn.Linear(in_channels, out_channels, bias_attr=False)
+        self.hardswish = nn.Hardswish()
+        self.dropout = nn.Dropout(p=last_drop, mode="downscale_in_infer")
+
+    def forward(self, x, sz):
+        x = x.reshape([0, sz[0], sz[1], x.shape[-1]])
+        x = x.mean(1)
+        x = self.last_conv(x)
+        x = self.hardswish(x)
+        x = self.dropout(x)
+        return x, [1, sz[1]]
+
+
+class OutPool(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, sz):
+        C = x.shape[-1]
+        x = x.transpose([0, 2, 1]).reshape([0, C, sz[0], sz[1]])
+        x = nn.functional.avg_pool2d(x, [sz[0], 2])
+        return x, [1, sz[1] // 2]
+
+
+class Feat2D(nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, sz):
+        C = x.shape[-1]
+        x = x.transpose([0, 2, 1]).reshape([0, C, sz[0], sz[1]])
+        return x, sz
+
+
+class SVTRv2(nn.Layer):
+    def __init__(
+        self,
+        max_sz=[32, 128],
+        in_channels=3,
+        out_channels=192,
+        out_char_num=25,
+        depths=[3, 6, 3],
+        dims=[64, 128, 256],
+        mixer=[["Conv"] * 3, ["Conv"] * 3 + ["Global"] * 3, ["Global"] * 3],
+        use_pos_embed=False,
+        sub_k=[[1, 1], [2, 1], [1, 1]],
+        num_heads=[2, 4, 8],
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        last_drop=0.1,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        norm_layer=nn.LayerNorm,
+        act=nn.GELU,
+        last_stage=False,
+        eps=1e-6,
+        use_pool=False,
+        feat2d=False,
+        **kwargs,
+    ):
+        super().__init__()
+        num_stages = len(depths)
+        self.num_features = dims[-1]
+
+        feat_max_size = [max_sz[0] // 4, max_sz[1] // 4]
+        self.pope = POPatchEmbed(
+            in_channels=in_channels,
+            feat_max_size=feat_max_size,
+            embed_dim=dims[0],
+            use_pos_embed=use_pos_embed,
+            flatten=mixer[0][0] != "Conv",
+        )
+
+        dpr = np.linspace(0, drop_path_rate, sum(depths))  # stochastic depth decay rule
+
+        self.stages = nn.LayerList()
+        for i_stage in range(num_stages):
+            stage = SVTRStage(
+                dim=dims[i_stage],
+                out_dim=dims[i_stage + 1] if i_stage < num_stages - 1 else 0,
+                depth=depths[i_stage],
+                mixer=mixer[i_stage],
+                sub_k=sub_k[i_stage],
+                num_heads=num_heads[i_stage],
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_stage]) : sum(depths[: i_stage + 1])],
+                norm_layer=norm_layer,
+                act=act,
+                downsample=False if i_stage == num_stages - 1 else True,
+                eps=eps,
+            )
+            self.stages.append(stage)
+
+        self.out_channels = self.num_features
+        self.last_stage = last_stage
+        if last_stage:
+            self.out_channels = out_channels
+            self.stages.append(
+                LastStage(self.num_features, out_channels, last_drop, out_char_num)
+            )
+        if use_pool:
+            self.stages.append(OutPool())
+
+        if feat2d:
+            self.stages.append(Feat2D())
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward(self, x):
+        x, sz = self.pope(x)
+        for stage in self.stages:
+            x, sz = stage(x, sz)
+        return x

ppocr/modeling/backbones/rec_vary_vit.py

@@ -0,0 +1,616 @@
+# copyright (c) 2024 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 math
+from functools import partial
+from typing import Optional, Tuple, Type
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn.initializer import (
+    Constant,
+    KaimingUniform,
+    Normal,
+    TruncatedNormal,
+    XavierUniform,
+)
+from ppocr.modeling.backbones.rec_donut_swin import DonutSwinModelOutput
+
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+kaiming_normal_ = KaimingUniform(nonlinearity="relu")
+trunc_normal_ = TruncatedNormal(std=0.02)
+xavier_uniform_ = XavierUniform()
+
+
+class MLPBlock(nn.Layer):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mlp_dim: int,
+        act: Type[nn.Layer] = nn.GELU,
+    ) -> None:
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x):
+        return self.lin2(self.act(self.lin1(x)))
+
+
+# From https://github.com/facebookresearch/detectron2/blob/main/detectron2/layers/batch_norm.py # noqa
+# Itself from https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119  # noqa
+class LayerNorm2d(nn.Layer):
+    def __init__(self, num_channels: int, epsilon: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = paddle.create_parameter([num_channels], dtype="float32")
+        ones_(self.weight)
+        self.bias = paddle.create_parameter([num_channels], dtype="float32")
+        zeros_(self.bias)
+        self.epsilon = epsilon
+
+    def forward(self, x):
+        u = x.mean(1, keepdim=True)
+        s = (x - u).pow(2).mean(1, keepdim=True)
+        x = (x - u) / paddle.sqrt(s + self.epsilon)
+        x = self.weight[:, None, None] * x + self.bias[:, None, None]
+        return x
+
+
+# This class and its supporting functions below lightly adapted from the ViTDet backbone available at: https://github.com/facebookresearch/detectron2/blob/main/detectron2/modeling/backbone/vit.py # noqa
+class ImageEncoderViT(nn.Layer):
+    def __init__(
+        self,
+        img_size: int = 1024,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        out_chans: int = 256,
+        qkv_bias: bool = True,
+        norm_layer: Type[nn.Layer] = nn.LayerNorm,
+        act_layer: Type[nn.Layer] = nn.GELU,
+        use_abs_pos: bool = True,
+        use_rel_pos: bool = False,
+        rel_pos_zero_init: bool = True,
+        window_size: int = 0,
+        global_attn_indexes: Tuple[int, ...] = (),
+        is_formula: bool = False,
+    ) -> None:
+        """
+        Args:
+            img_size (int): Input image size.
+            patch_size (int): Patch size.
+            in_chans (int): Number of input image channels.
+            embed_dim (int): Patch embedding dimension.
+            depth (int): Depth of ViT.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            norm_layer (nn.Layer): Normalization layer.
+            act_layer (nn.Layer): Activation layer.
+            use_abs_pos (bool): If True, use absolute positional embeddings.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks.
+            global_attn_indexes (list): Indexes for blocks using global attention.
+        """
+        super().__init__()
+        self.img_size = img_size
+
+        self.patch_embed = PatchEmbed(
+            kernel_size=(patch_size, patch_size),
+            stride=(patch_size, patch_size),
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+
+        self.pos_embed = None
+        if use_abs_pos:
+            # Initialize absolute positional embedding with pretrain image size.
+            self.pos_embed = paddle.create_parameter(
+                shape=(1, img_size // patch_size, img_size // patch_size, embed_dim),
+                dtype="float32",
+            )
+            zeros_(self.pos_embed)
+
+        self.blocks = nn.LayerList()
+        for i in range(depth):
+            block = Block(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                norm_layer=norm_layer,
+                act_layer=act_layer,
+                use_rel_pos=use_rel_pos,
+                rel_pos_zero_init=rel_pos_zero_init,
+                window_size=window_size if i not in global_attn_indexes else 0,
+                input_size=(img_size // patch_size, img_size // patch_size),
+            )
+            self.blocks.append(block)
+
+        self.neck = nn.Sequential(
+            nn.Conv2D(
+                embed_dim,
+                out_chans,
+                kernel_size=1,
+                bias_attr=False,
+            ),
+            LayerNorm2d(out_chans),
+            nn.Conv2D(
+                out_chans,
+                out_chans,
+                kernel_size=3,
+                padding=1,
+                bias_attr=False,
+            ),
+            LayerNorm2d(out_chans),
+        )
+
+        self.net_2 = nn.Conv2D(
+            256, 512, kernel_size=3, stride=2, padding=1, bias_attr=False
+        )
+        self.net_3 = nn.Conv2D(
+            512, 1024, kernel_size=3, stride=2, padding=1, bias_attr=False
+        )
+        self.is_formula = is_formula
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.neck(x.transpose([0, 3, 1, 2]))
+        x = self.net_2(x)
+        if self.is_formula:
+            x = self.net_3(x)
+        return x
+
+
+class Block(nn.Layer):
+    """Transformer blocks with support of window attention and residual propagation blocks"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        norm_layer: Type[nn.Layer] = nn.LayerNorm,
+        act_layer: Type[nn.Layer] = nn.GELU,
+        use_rel_pos: bool = False,
+        rel_pos_zero_init: bool = True,
+        window_size: int = 0,
+        input_size: Optional[Tuple[int, int]] = None,
+    ) -> None:
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads in each ViT block.
+            mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+            qkv_bias (bool): If True, add a learnable bias to query, key, value.
+            norm_layer (nn.Layer): Normalization layer.
+            act_layer (nn.Layer): Activation layer.
+            use_rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            window_size (int): Window size for window attention blocks. If it equals 0, then
+                use global attention.
+            input_size (tuple(int, int) or None): Input resolution for calculating the relative
+                positional parameter size.
+        """
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            use_rel_pos=use_rel_pos,
+            rel_pos_zero_init=rel_pos_zero_init,
+            input_size=input_size if window_size == 0 else (window_size, window_size),
+        )
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = MLPBlock(
+            embedding_dim=dim, mlp_dim=int(dim * mlp_ratio), act=act_layer
+        )
+
+        self.window_size = window_size
+
+    def forward(self, x):
+        shortcut = x
+
+        x = self.norm1(x)
+        # Window partition
+        if self.window_size > 0:
+            H, W = x.shape[1], x.shape[2]
+            x, pad_hw = window_partition(x, self.window_size)
+        x = self.attn(x)
+        # Reverse window partition
+        if self.window_size > 0:
+            x = window_unpartition(x, self.window_size, pad_hw, (H, W))
+        x = shortcut + x
+        x = x + self.mlp(self.norm2(x))
+
+        return x
+
+
+class Attention(nn.Layer):
+    """Multi-head Attention block with relative position embeddings."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        use_rel_pos: bool = False,
+        rel_pos_zero_init: bool = True,
+        input_size: Optional[Tuple[int, int]] = None,
+    ) -> None:
+        """
+        Args:
+            dim (int): Number of input channels.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool):  If True, add a learnable bias to query, key, value.
+            rel_pos (bool): If True, add relative positional embeddings to the attention map.
+            rel_pos_zero_init (bool): If True, zero initialize relative positional parameters.
+            input_size (tuple(int, int) or None): Input resolution for calculating the relative
+                positional parameter size.
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        self.use_rel_pos = use_rel_pos
+        if self.use_rel_pos:
+            assert (
+                input_size is not None
+            ), "Input size must be provided if using relative positional encoding."
+            # initialize relative positional embeddings
+            self.rel_pos_h = paddle.create_parameter(
+                [2 * input_size[0] - 1, head_dim], dtype="float32"
+            )
+            zeros_(self.rel_pos_h)
+            self.rel_pos_w = paddle.create_parameter(
+                [2 * input_size[1] - 1, head_dim], dtype="float32"
+            )
+            zeros_(self.rel_pos_w)
+
+    def forward(self, x):
+
+        B, H, W, _ = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape([B, H * W, 3, self.num_heads, -1])
+            .transpose([2, 0, 3, 1, 4])
+        )
+        q, k, v = qkv.reshape([3, B * self.num_heads, H * W, -1]).unbind(0)
+        attn = (q * self.scale) @ k.transpose([0, 2, 1])
+
+        if self.use_rel_pos:
+            attn = add_decomposed_rel_pos(
+                attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W)
+            )
+        attn = F.softmax(attn, axis=-1)
+        x = (
+            (attn @ v)
+            .reshape([B, self.num_heads, H, W, -1])
+            .transpose([0, 2, 3, 1, 4])
+            .reshape([B, H, W, -1])
+        )
+        x = self.proj(x)
+
+        return x
+
+
+def window_partition(x, window_size: int):
+    """
+    Partition into non-overlapping windows with padding if needed.
+    Args:
+        x (tensor): input tokens with [B, H, W, C].
+        window_size (int): window size.
+
+    Returns:
+        windows: windows after partition with [B * num_windows, window_size, window_size, C].
+        (Hp, Wp): padded height and width before partition
+    """
+    B, H, W, C = x.shape
+
+    pad_h = (window_size - H % window_size) % window_size
+    pad_w = (window_size - W % window_size) % window_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h, 0, 0))
+    Hp, Wp = H + pad_h, W + pad_w
+
+    x = x.reshape(
+        [B, Hp // window_size, window_size, Wp // window_size, window_size, C]
+    )
+    windows = x.transpose([0, 1, 3, 2, 4, 5]).reshape([-1, window_size, window_size, C])
+    return windows, (Hp, Wp)
+
+
+def window_unpartition(
+    windows, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]
+):
+    """
+    Window unpartition into original sequences and removing padding.
+    Args:
+        windows (tensor): input tokens with [B * num_windows, window_size, window_size, C].
+        window_size (int): window size.
+        pad_hw (Tuple): padded height and width (Hp, Wp).
+        hw (Tuple): original height and width (H, W) before padding.
+
+    Returns:
+        x: unpartitioned sequences with [B, H, W, C].
+    """
+    Hp, Wp = pad_hw
+    H, W = hw
+    B = windows.shape[0] // (Hp * Wp // window_size // window_size)
+    x = windows.reshape(
+        [B, Hp // window_size, Wp // window_size, window_size, window_size, -1]
+    )
+    x = x.transpose([0, 1, 3, 2, 4, 5]).contiguous().reshape([B, Hp, Wp, -1])
+
+    if Hp > H or Wp > W:
+        x = x[:, :H, :W, :].contiguous()
+    return x
+
+
+def get_rel_pos(q_size: int, k_size: int, rel_pos):
+    """
+    Get relative positional embeddings according to the relative positions of
+        query and key sizes.
+    Args:
+        q_size (int): size of query q.
+        k_size (int): size of key k.
+        rel_pos (Tensor): relative position embeddings (L, C).
+
+    Returns:
+        Extracted positional embeddings according to relative positions.
+    """
+    max_rel_dist = int(2 * max(q_size, k_size) - 1)
+    # Interpolate rel pos if needed.
+    if rel_pos.shape[0] != max_rel_dist:
+        # Interpolate rel pos.
+        rel_pos_resized = F.interpolate(
+            rel_pos.reshape(1, rel_pos.shape[0], -1).transpose(0, 2, 1),
+            size=max_rel_dist,
+            mode="linear",
+        )
+        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).transpose(1, 0)
+    else:
+        rel_pos_resized = rel_pos
+
+    # Scale the coords with short length if shapes for q and k are different.
+    q_coords = paddle.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
+    k_coords = paddle.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
+    relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)
+
+    return rel_pos_resized[relative_coords.cast(paddle.int64)]
+
+
+def add_decomposed_rel_pos(
+    attn,
+    q,
+    rel_pos_h,
+    rel_pos_w,
+    q_size: Tuple[int, int],
+    k_size: Tuple[int, int],
+):
+    """
+    Calculate decomposed Relative Positional Embeddings from :paper:`mvitv2`.
+    https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py   # noqa B950
+    Args:
+        attn (Tensor): attention map.
+        q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).
+        rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.
+        rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.
+        q_size (Tuple): spatial sequence size of query q with (q_h, q_w).
+        k_size (Tuple): spatial sequence size of key k with (k_h, k_w).
+
+    Returns:
+        attn (Tensor): attention map with added relative positional embeddings.
+    """
+    q_h, q_w = q_size
+    k_h, k_w = k_size
+    Rh = get_rel_pos(q_h, k_h, rel_pos_h)
+    Rw = get_rel_pos(q_w, k_w, rel_pos_w)
+
+    B, _, dim = q.shape
+    r_q = q.reshape([B, q_h, q_w, dim])
+    rel_h = paddle.einsum("bhwc,hkc->bhwk", r_q, Rh)
+    rel_w = paddle.einsum("bhwc,wkc->bhwk", r_q, Rw)
+
+    attn = (
+        attn.reshape([B, q_h, q_w, k_h, k_w])
+        + rel_h[:, :, :, :, None]
+        + rel_w[:, :, :, None, :]
+    ).reshape([B, q_h * q_w, k_h * k_w])
+
+    return attn
+
+
+class PatchEmbed(nn.Layer):
+    """
+    Image to Patch Embedding.
+    """
+
+    def __init__(
+        self,
+        kernel_size: Tuple[int, int] = (16, 16),
+        stride: Tuple[int, int] = (16, 16),
+        padding: Tuple[int, int] = (0, 0),
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ) -> None:
+        """
+        Args:
+            kernel_size (Tuple): kernel size of the projection layer.
+            stride (Tuple): stride of the projection layer.
+            padding (Tuple): padding size of the projection layer.
+            in_chans (int): Number of input image channels.
+            embed_dim (int): Patch embedding dimension.
+        """
+        super().__init__()
+
+        self.proj = nn.Conv2D(
+            in_chans,
+            embed_dim,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            weight_attr=True,
+            bias_attr=True,
+        )
+
+    def forward(self, x):
+        x = self.proj(x)
+        # B C H W -> B H W C
+        x = x.transpose([0, 2, 3, 1])
+        return x
+
+
+def _build_vary(
+    encoder_embed_dim,
+    encoder_depth,
+    encoder_num_heads,
+    encoder_global_attn_indexes,
+    image_size,
+    is_formula=False,
+):
+    prompt_embed_dim = 256
+    vit_patch_size = 16
+    image_embedding_size = image_size // vit_patch_size
+    image_encoder = ImageEncoderViT(
+        depth=encoder_depth,
+        embed_dim=encoder_embed_dim,
+        img_size=image_size,
+        mlp_ratio=4,
+        norm_layer=partial(paddle.nn.LayerNorm, epsilon=1e-6),
+        num_heads=encoder_num_heads,
+        patch_size=vit_patch_size,
+        qkv_bias=True,
+        use_rel_pos=True,
+        global_attn_indexes=encoder_global_attn_indexes,
+        window_size=14,
+        out_chans=prompt_embed_dim,
+        is_formula=is_formula,
+    )
+    return image_encoder
+
+
+class Vary_VIT_B(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        image_size=768,
+        encoder_embed_dim=768,
+        encoder_depth=12,
+        encoder_num_heads=12,
+        encoder_global_attn_indexes=[2, 5, 8, 11],
+    ):
+        super().__init__()
+
+        self.vision_tower_high = _build_vary(
+            encoder_embed_dim=768,
+            encoder_depth=12,
+            encoder_num_heads=12,
+            encoder_global_attn_indexes=[2, 5, 8, 11],
+            image_size=image_size,
+        )
+
+        self.out_channels = 1024
+
+    def forward(self, input_data):
+        pixel_values = input_data
+        num_channels = pixel_values.shape[1]
+        if num_channels == 1:
+            pixel_values = paddle.repeat_interleave(pixel_values, repeats=3, axis=1)
+        cnn_feature = self.vision_tower_high(pixel_values)
+        cnn_feature = cnn_feature.flatten(2).transpose([0, 2, 1])
+        return cnn_feature
+
+
+class Vary_VIT_B_Formula(nn.Layer):
+    def __init__(
+        self,
+        in_channels=3,
+        image_size=768,
+        encoder_embed_dim=768,
+        encoder_depth=12,
+        encoder_num_heads=12,
+        encoder_global_attn_indexes=[2, 5, 8, 11],
+    ):
+        """
+        Vary_VIT_B_Formula
+        Args:
+            in_channels (int): Number of input channels. Default is 3 (for RGB images).
+            image_size (int): Size of the input image. Default is 768.
+            encoder_embed_dim (int): Dimension of the encoder's embedding. Default is 768.
+            encoder_depth (int): Number of layers (depth) in the encoder. Default is 12.
+            encoder_num_heads (int): Number of attention heads in the encoder. Default is 12.
+            encoder_global_attn_indexes (list): List of indices specifying which encoder layers use global attention. Default is [2, 5, 8, 11].
+        Returns:
+            model: nn.Layer. Specific `Vary_VIT_B_Formula` model with defined architecture.
+        """
+        super(Vary_VIT_B_Formula, self).__init__()
+
+        self.vision_tower_high = _build_vary(
+            encoder_embed_dim=encoder_embed_dim,
+            encoder_depth=encoder_depth,
+            encoder_num_heads=encoder_num_heads,
+            encoder_global_attn_indexes=[2, 5, 8, 11],
+            image_size=image_size,
+            is_formula=True,
+        )
+        self.mm_projector_vary = nn.Linear(1024, 1024)
+        self.out_channels = 1024
+
+    def forward(self, input_data):
+        if self.training:
+            pixel_values, label, attention_mask = input_data
+        else:
+            if isinstance(input_data, list):
+                pixel_values = input_data[0]
+            else:
+                pixel_values = input_data
+        num_channels = pixel_values.shape[1]
+        if num_channels == 1:
+            pixel_values = paddle.repeat_interleave(pixel_values, repeats=3, axis=1)
+
+        cnn_feature = self.vision_tower_high(pixel_values)
+        cnn_feature = cnn_feature.flatten(2).transpose([0, 2, 1])
+
+        cnn_feature = self.mm_projector_vary(cnn_feature)
+        donut_swin_output = DonutSwinModelOutput(
+            last_hidden_state=cnn_feature,
+            pooler_output=None,
+            hidden_states=None,
+            attentions=None,
+            reshaped_hidden_states=None,
+        )
+        if self.training:
+            return donut_swin_output, label, attention_mask
+        else:
+            return donut_swin_output

ppocr/modeling/backbones/rec_vit.py

@@ -0,0 +1,273 @@
+# copyright (c) 2023 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 paddle import ParamAttr
+from paddle.nn.initializer import KaimingNormal
+import numpy as np
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+trunc_normal_ = TruncatedNormal(std=0.02)
+normal_ = Normal
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+def drop_path(x, drop_prob=0.0, training=False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ...
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = paddle.to_tensor(1 - drop_prob)
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + paddle.rand(shape, dtype=x.dtype)
+    random_tensor = paddle.floor(random_tensor)  # binarize
+    output = x.divide(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Layer):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Identity(nn.Layer):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Layer):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dim = dim
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        qkv = paddle.reshape(
+            self.qkv(x), (0, -1, 3, self.num_heads, self.dim // self.num_heads)
+        ).transpose((2, 0, 3, 1, 4))
+        q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+
+        attn = q.matmul(k.transpose((0, 1, 3, 2)))
+        attn = nn.functional.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((0, -1, self.dim))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-6,
+        prenorm=True,
+    ):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
+        else:
+            self.norm1 = norm_layer(dim)
+        self.mixer = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        if isinstance(norm_layer, str):
+            self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
+        else:
+            self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.prenorm = prenorm
+
+    def forward(self, x):
+        if self.prenorm:
+            x = self.norm1(x + self.drop_path(self.mixer(x)))
+            x = self.norm2(x + self.drop_path(self.mlp(x)))
+        else:
+            x = x + self.drop_path(self.mixer(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class ViT(nn.Layer):
+    def __init__(
+        self,
+        img_size=[32, 128],
+        patch_size=[4, 4],
+        in_channels=3,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.1,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-6,
+        act="nn.GELU",
+        prenorm=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.out_channels = embed_dim
+        self.prenorm = prenorm
+        self.patch_embed = nn.Conv2D(
+            in_channels, embed_dim, patch_size, patch_size, padding=(0, 0)
+        )
+        self.pos_embed = self.create_parameter(
+            shape=[1, 257, embed_dim], default_initializer=zeros_
+        )
+        self.add_parameter("pos_embed", self.pos_embed)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = np.linspace(0, drop_path_rate, depth)
+        self.blocks1 = nn.LayerList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    act_layer=eval(act),
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    epsilon=epsilon,
+                    prenorm=prenorm,
+                )
+                for i in range(depth)
+            ]
+        )
+        if not prenorm:
+            self.norm = eval(norm_layer)(embed_dim, epsilon=epsilon)
+
+        self.avg_pool = nn.AdaptiveAvgPool2D([1, 25])
+        self.last_conv = nn.Conv2D(
+            in_channels=embed_dim,
+            out_channels=self.out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias_attr=False,
+        )
+        self.hardswish = nn.Hardswish()
+        self.dropout = nn.Dropout(p=0.1, mode="downscale_in_infer")
+
+        trunc_normal_(self.pos_embed)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward(self, x):
+        x = self.patch_embed(x).flatten(2).transpose((0, 2, 1))
+        x = x + self.pos_embed[:, 1:, :]  # [:, :x.shape[1], :]
+        x = self.pos_drop(x)
+        for blk in self.blocks1:
+            x = blk(x)
+        if not self.prenorm:
+            x = self.norm(x)
+
+        x = self.avg_pool(x.transpose([0, 2, 1]).reshape([0, self.embed_dim, -1, 25]))
+        x = self.last_conv(x)
+        x = self.hardswish(x)
+        x = self.dropout(x)
+        return x

ppocr/modeling/backbones/rec_vit_parseq.py

@@ -0,0 +1,348 @@
+# copyright (c) 2021 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/PaddlePaddle/PaddleClas/blob/release%2F2.5/ppcls/arch/backbone/model_zoo/vision_transformer.py
+"""
+
+from collections.abc import Callable
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+
+trunc_normal_ = TruncatedNormal(std=0.02)
+normal_ = Normal
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+def to_2tuple(x):
+    return tuple([x] * 2)
+
+
+def drop_path(x, drop_prob=0.0, training=False):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ...
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = paddle.to_tensor(1 - drop_prob, dtype=x.dtype)
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + paddle.rand(shape).astype(x.dtype)
+    random_tensor = paddle.floor(random_tensor)  # binarize
+    output = x.divide(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Layer):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Identity(nn.Layer):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Layer):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        # B= x.shape[0]
+        N, C = x.shape[1:]
+        qkv = (
+            self.qkv(x)
+            .reshape((-1, N, 3, self.num_heads, C // self.num_heads))
+            .transpose((2, 0, 3, 1, 4))
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
+        attn = nn.functional.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((-1, N, C))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-5,
+    ):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
+        elif isinstance(norm_layer, Callable):
+            self.norm1 = norm_layer(dim)
+        else:
+            raise TypeError("The norm_layer must be str or paddle.nn.layer.Layer class")
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        if isinstance(norm_layer, str):
+            self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
+        elif isinstance(norm_layer, Callable):
+            self.norm2 = norm_layer(dim)
+        else:
+            raise TypeError("The norm_layer must be str or paddle.nn.layer.Layer class")
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Layer):
+    """Image to Patch Embedding"""
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        if isinstance(img_size, int):
+            img_size = to_2tuple(img_size)
+        if isinstance(patch_size, int):
+            patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2D(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+
+        x = self.proj(x).flatten(2).transpose((0, 2, 1))
+        return x
+
+
+class VisionTransformer(nn.Layer):
+    """Vision Transformer with support for patch input"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_channels=3,
+        class_num=1000,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-5,
+        **kwargs,
+    ):
+        super().__init__()
+        self.class_num = class_num
+
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_channels,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.pos_embed = self.create_parameter(
+            shape=(1, num_patches, embed_dim), default_initializer=zeros_
+        )
+        self.add_parameter("pos_embed", self.pos_embed)
+        self.cls_token = self.create_parameter(
+            shape=(1, 1, embed_dim), default_initializer=zeros_
+        )
+        self.add_parameter("cls_token", self.cls_token)
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = np.linspace(0, drop_path_rate, depth)
+
+        self.blocks = nn.LayerList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    epsilon=epsilon,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.norm = eval(norm_layer)(embed_dim, epsilon=epsilon)
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, class_num) if class_num > 0 else Identity()
+
+        trunc_normal_(self.pos_embed)
+        self.out_channels = embed_dim
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+class ViTParseQ(VisionTransformer):
+    def __init__(
+        self,
+        img_size=[224, 224],
+        patch_size=[16, 16],
+        in_channels=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+    ):
+        super().__init__(
+            img_size,
+            patch_size,
+            in_channels,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            drop_rate=drop_rate,
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate,
+            class_num=0,
+        )
+
+    def forward(self, x):
+        return self.forward_features(x)

ppocr/modeling/backbones/rec_vitstr.py

@@ -0,0 +1,133 @@
+# copyright (c) 2021 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/roatienza/deep-text-recognition-benchmark/blob/master/modules/vitstr.py
+"""
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+from ppocr.modeling.backbones.rec_svtrnet import (
+    Block,
+    PatchEmbed,
+    zeros_,
+    trunc_normal_,
+    ones_,
+)
+
+scale_dim_heads = {"tiny": [192, 3], "small": [384, 6], "base": [768, 12]}
+
+
+class ViTSTR(nn.Layer):
+    def __init__(
+        self,
+        img_size=[224, 224],
+        in_channels=1,
+        scale="tiny",
+        seqlen=27,
+        patch_size=[16, 16],
+        embed_dim=None,
+        depth=12,
+        num_heads=None,
+        mlp_ratio=4,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_path_rate=0.0,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        norm_layer="nn.LayerNorm",
+        act_layer="nn.GELU",
+        epsilon=1e-6,
+        out_channels=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.seqlen = seqlen
+        embed_dim = embed_dim if embed_dim is not None else scale_dim_heads[scale][0]
+        num_heads = num_heads if num_heads is not None else scale_dim_heads[scale][1]
+        out_channels = out_channels if out_channels is not None else embed_dim
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            in_channels=in_channels,
+            embed_dim=embed_dim,
+            patch_size=patch_size,
+            mode="linear",
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.pos_embed = self.create_parameter(
+            shape=[1, num_patches + 1, embed_dim], default_initializer=zeros_
+        )
+        self.add_parameter("pos_embed", self.pos_embed)
+        self.cls_token = self.create_parameter(
+            shape=[1, 1, embed_dim], default_initializer=zeros_
+        )
+        self.add_parameter("cls_token", self.cls_token)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = np.linspace(0, drop_path_rate, depth)
+        self.blocks = nn.LayerList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    act_layer=eval(act_layer),
+                    epsilon=epsilon,
+                    prenorm=False,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = eval(norm_layer)(embed_dim, epsilon=epsilon)
+
+        self.out_channels = out_channels
+
+        trunc_normal_(self.pos_embed)
+        trunc_normal_(self.cls_token)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+        cls_tokens = paddle.tile(self.cls_token, repeat_times=[B, 1, 1])
+        x = paddle.concat((cls_tokens, x), axis=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = x[:, : self.seqlen]
+        return x.transpose([0, 2, 1]).unsqueeze(2)

ppocr/modeling/backbones/table_master_resnet.py

@@ -0,0 +1,365 @@
+# Copyright (c) 2022 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.
+"""
+This code is refer from:
+https://github.com/JiaquanYe/TableMASTER-mmocr/blob/master/mmocr/models/textrecog/backbones/table_resnet_extra.py
+"""
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class BasicBlock(nn.Layer):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, gcb_config=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2D(
+            inplanes, planes, kernel_size=3, stride=stride, padding=1, bias_attr=False
+        )
+        self.bn1 = nn.BatchNorm2D(planes, momentum=0.9)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2D(
+            planes, planes, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn2 = nn.BatchNorm2D(planes, momentum=0.9)
+        self.downsample = downsample
+        self.stride = stride
+        self.gcb_config = gcb_config
+
+        if self.gcb_config is not None:
+            gcb_ratio = gcb_config["ratio"]
+            gcb_headers = gcb_config["headers"]
+            att_scale = gcb_config["att_scale"]
+            fusion_type = gcb_config["fusion_type"]
+            self.context_block = MultiAspectGCAttention(
+                inplanes=planes,
+                ratio=gcb_ratio,
+                headers=gcb_headers,
+                att_scale=att_scale,
+                fusion_type=fusion_type,
+            )
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.gcb_config is not None:
+            out = self.context_block(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+def get_gcb_config(gcb_config, layer):
+    if gcb_config is None or not gcb_config["layers"][layer]:
+        return None
+    else:
+        return gcb_config
+
+
+class TableResNetExtra(nn.Layer):
+    def __init__(self, layers, in_channels=3, gcb_config=None):
+        assert len(layers) >= 4
+
+        super(TableResNetExtra, self).__init__()
+        self.inplanes = 128
+        self.conv1 = nn.Conv2D(
+            in_channels, 64, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn1 = nn.BatchNorm2D(64)
+        self.relu1 = nn.ReLU()
+
+        self.conv2 = nn.Conv2D(
+            64, 128, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn2 = nn.BatchNorm2D(128)
+        self.relu2 = nn.ReLU()
+
+        self.maxpool1 = nn.MaxPool2D(kernel_size=2, stride=2)
+
+        self.layer1 = self._make_layer(
+            BasicBlock,
+            256,
+            layers[0],
+            stride=1,
+            gcb_config=get_gcb_config(gcb_config, 0),
+        )
+
+        self.conv3 = nn.Conv2D(
+            256, 256, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn3 = nn.BatchNorm2D(256)
+        self.relu3 = nn.ReLU()
+
+        self.maxpool2 = nn.MaxPool2D(kernel_size=2, stride=2)
+
+        self.layer2 = self._make_layer(
+            BasicBlock,
+            256,
+            layers[1],
+            stride=1,
+            gcb_config=get_gcb_config(gcb_config, 1),
+        )
+
+        self.conv4 = nn.Conv2D(
+            256, 256, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn4 = nn.BatchNorm2D(256)
+        self.relu4 = nn.ReLU()
+
+        self.maxpool3 = nn.MaxPool2D(kernel_size=2, stride=2)
+
+        self.layer3 = self._make_layer(
+            BasicBlock,
+            512,
+            layers[2],
+            stride=1,
+            gcb_config=get_gcb_config(gcb_config, 2),
+        )
+
+        self.conv5 = nn.Conv2D(
+            512, 512, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn5 = nn.BatchNorm2D(512)
+        self.relu5 = nn.ReLU()
+
+        self.layer4 = self._make_layer(
+            BasicBlock,
+            512,
+            layers[3],
+            stride=1,
+            gcb_config=get_gcb_config(gcb_config, 3),
+        )
+
+        self.conv6 = nn.Conv2D(
+            512, 512, kernel_size=3, stride=1, padding=1, bias_attr=False
+        )
+        self.bn6 = nn.BatchNorm2D(512)
+        self.relu6 = nn.ReLU()
+
+        self.out_channels = [256, 256, 512]
+
+    def _make_layer(self, block, planes, blocks, stride=1, gcb_config=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2D(
+                    self.inplanes,
+                    planes * block.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                    bias_attr=False,
+                ),
+                nn.BatchNorm2D(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, gcb_config=gcb_config)
+        )
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        f = []
+        x = self.conv1(x)
+
+        x = self.bn1(x)
+        x = self.relu1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu2(x)
+
+        x = self.maxpool1(x)
+        x = self.layer1(x)
+
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu3(x)
+        f.append(x)
+
+        x = self.maxpool2(x)
+        x = self.layer2(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu4(x)
+        f.append(x)
+
+        x = self.maxpool3(x)
+
+        x = self.layer3(x)
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu5(x)
+
+        x = self.layer4(x)
+        x = self.conv6(x)
+        x = self.bn6(x)
+        x = self.relu6(x)
+        f.append(x)
+        return f
+
+
+class MultiAspectGCAttention(nn.Layer):
+    def __init__(
+        self,
+        inplanes,
+        ratio,
+        headers,
+        pooling_type="att",
+        att_scale=False,
+        fusion_type="channel_add",
+    ):
+        super(MultiAspectGCAttention, self).__init__()
+        assert pooling_type in ["avg", "att"]
+
+        assert fusion_type in ["channel_add", "channel_mul", "channel_concat"]
+        assert (
+            inplanes % headers == 0 and inplanes >= 8
+        )  # inplanes must be divided by headers evenly
+
+        self.headers = headers
+        self.inplanes = inplanes
+        self.ratio = ratio
+        self.planes = int(inplanes * ratio)
+        self.pooling_type = pooling_type
+        self.fusion_type = fusion_type
+        self.att_scale = False
+
+        self.single_header_inplanes = int(inplanes / headers)
+
+        if pooling_type == "att":
+            self.conv_mask = nn.Conv2D(self.single_header_inplanes, 1, kernel_size=1)
+            self.softmax = nn.Softmax(axis=2)
+        else:
+            self.avg_pool = nn.AdaptiveAvgPool2D(1)
+
+        if fusion_type == "channel_add":
+            self.channel_add_conv = nn.Sequential(
+                nn.Conv2D(self.inplanes, self.planes, kernel_size=1),
+                nn.LayerNorm([self.planes, 1, 1]),
+                nn.ReLU(),
+                nn.Conv2D(self.planes, self.inplanes, kernel_size=1),
+            )
+        elif fusion_type == "channel_concat":
+            self.channel_concat_conv = nn.Sequential(
+                nn.Conv2D(self.inplanes, self.planes, kernel_size=1),
+                nn.LayerNorm([self.planes, 1, 1]),
+                nn.ReLU(),
+                nn.Conv2D(self.planes, self.inplanes, kernel_size=1),
+            )
+            # for concat
+            self.cat_conv = nn.Conv2D(2 * self.inplanes, self.inplanes, kernel_size=1)
+        elif fusion_type == "channel_mul":
+            self.channel_mul_conv = nn.Sequential(
+                nn.Conv2D(self.inplanes, self.planes, kernel_size=1),
+                nn.LayerNorm([self.planes, 1, 1]),
+                nn.ReLU(),
+                nn.Conv2D(self.planes, self.inplanes, kernel_size=1),
+            )
+
+    def spatial_pool(self, x):
+        batch, channel, height, width = x.shape
+        if self.pooling_type == "att":
+            # [N*headers, C', H , W] C = headers * C'
+            x = x.reshape(
+                [batch * self.headers, self.single_header_inplanes, height, width]
+            )
+            input_x = x
+
+            # [N*headers, C', H * W] C = headers * C'
+            # input_x = input_x.view(batch, channel, height * width)
+            input_x = input_x.reshape(
+                [batch * self.headers, self.single_header_inplanes, height * width]
+            )
+
+            # [N*headers, 1, C', H * W]
+            input_x = input_x.unsqueeze(1)
+            # [N*headers, 1, H, W]
+            context_mask = self.conv_mask(x)
+            # [N*headers, 1, H * W]
+            context_mask = context_mask.reshape(
+                [batch * self.headers, 1, height * width]
+            )
+
+            # scale variance
+            if self.att_scale and self.headers > 1:
+                context_mask = context_mask / paddle.sqrt(self.single_header_inplanes)
+
+            # [N*headers, 1, H * W]
+            context_mask = self.softmax(context_mask)
+
+            # [N*headers, 1, H * W, 1]
+            context_mask = context_mask.unsqueeze(-1)
+            # [N*headers, 1, C', 1] = [N*headers, 1, C', H * W] * [N*headers, 1, H * W, 1]
+            context = paddle.matmul(input_x, context_mask)
+
+            # [N, headers * C', 1, 1]
+            context = context.reshape(
+                [batch, self.headers * self.single_header_inplanes, 1, 1]
+            )
+        else:
+            # [N, C, 1, 1]
+            context = self.avg_pool(x)
+
+        return context
+
+    def forward(self, x):
+        # [N, C, 1, 1]
+        context = self.spatial_pool(x)
+
+        out = x
+
+        if self.fusion_type == "channel_mul":
+            # [N, C, 1, 1]
+            channel_mul_term = F.sigmoid(self.channel_mul_conv(context))
+            out = out * channel_mul_term
+        elif self.fusion_type == "channel_add":
+            # [N, C, 1, 1]
+            channel_add_term = self.channel_add_conv(context)
+            out = out + channel_add_term
+        else:
+            # [N, C, 1, 1]
+            channel_concat_term = self.channel_concat_conv(context)
+
+            # use concat
+            _, C1, _, _ = channel_concat_term.shape
+            N, C2, H, W = out.shape
+
+            out = paddle.concat(
+                [out, channel_concat_term.expand([-1, -1, H, W])], axis=1
+            )
+            out = self.cat_conv(out)
+            out = F.layer_norm(out, [self.inplanes, H, W])
+            out = F.relu(out)
+
+        return out

ppocr/modeling/backbones/vqa_layoutlm.py

@@ -0,0 +1,260 @@
+# copyright (c) 2021 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 os
+from paddle import nn
+
+from paddlenlp.transformers import (
+    LayoutXLMModel,
+    LayoutXLMForTokenClassification,
+    LayoutXLMForRelationExtraction,
+)
+from paddlenlp.transformers import LayoutLMModel, LayoutLMForTokenClassification
+from paddlenlp.transformers import (
+    LayoutLMv2Model,
+    LayoutLMv2ForTokenClassification,
+    LayoutLMv2ForRelationExtraction,
+)
+from paddlenlp.transformers import AutoModel
+
+__all__ = ["LayoutXLMForSer", "LayoutLMForSer"]
+
+pretrained_model_dict = {
+    LayoutXLMModel: {
+        "base": "layoutxlm-base-uncased",
+        "vi": "vi-layoutxlm-base-uncased",
+    },
+    LayoutLMModel: {
+        "base": "layoutlm-base-uncased",
+    },
+    LayoutLMv2Model: {
+        "base": "layoutlmv2-base-uncased",
+        "vi": "vi-layoutlmv2-base-uncased",
+    },
+}
+
+
+class NLPBaseModel(nn.Layer):
+    def __init__(
+        self,
+        base_model_class,
+        model_class,
+        mode="base",
+        type="ser",
+        pretrained=True,
+        checkpoints=None,
+        **kwargs,
+    ):
+        super(NLPBaseModel, self).__init__()
+        if checkpoints is not None:  # load the trained model
+            self.model = model_class.from_pretrained(checkpoints)
+        else:  # load the pretrained-model
+            pretrained_model_name = pretrained_model_dict[base_model_class][mode]
+            if type == "ser":
+                self.model = model_class.from_pretrained(
+                    pretrained_model_name, num_classes=kwargs["num_classes"], dropout=0
+                )
+            else:
+                self.model = model_class.from_pretrained(
+                    pretrained_model_name, dropout=0
+                )
+        self.out_channels = 1
+        self.use_visual_backbone = True
+
+
+class LayoutLMForSer(NLPBaseModel):
+    def __init__(
+        self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
+    ):
+        super(LayoutLMForSer, self).__init__(
+            LayoutLMModel,
+            LayoutLMForTokenClassification,
+            mode,
+            "ser",
+            pretrained,
+            checkpoints,
+            num_classes=num_classes,
+        )
+        self.use_visual_backbone = False
+
+    def forward(self, x):
+        x = self.model(
+            input_ids=x[0],
+            bbox=x[1],
+            attention_mask=x[2],
+            token_type_ids=x[3],
+            position_ids=None,
+            output_hidden_states=False,
+        )
+        return x
+
+
+class LayoutLMv2ForSer(NLPBaseModel):
+    def __init__(
+        self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
+    ):
+        super(LayoutLMv2ForSer, self).__init__(
+            LayoutLMv2Model,
+            LayoutLMv2ForTokenClassification,
+            mode,
+            "ser",
+            pretrained,
+            checkpoints,
+            num_classes=num_classes,
+        )
+        if (
+            hasattr(self.model.layoutlmv2, "use_visual_backbone")
+            and self.model.layoutlmv2.use_visual_backbone is False
+        ):
+            self.use_visual_backbone = False
+
+    def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+        else:
+            image = None
+        x = self.model(
+            input_ids=x[0],
+            bbox=x[1],
+            attention_mask=x[2],
+            token_type_ids=x[3],
+            image=image,
+            position_ids=None,
+            head_mask=None,
+            labels=None,
+        )
+        if self.training:
+            res = {"backbone_out": x[0]}
+            res.update(x[1])
+            return res
+        else:
+            return x
+
+
+class LayoutXLMForSer(NLPBaseModel):
+    def __init__(
+        self, num_classes, pretrained=True, checkpoints=None, mode="base", **kwargs
+    ):
+        super(LayoutXLMForSer, self).__init__(
+            LayoutXLMModel,
+            LayoutXLMForTokenClassification,
+            mode,
+            "ser",
+            pretrained,
+            checkpoints,
+            num_classes=num_classes,
+        )
+        if (
+            hasattr(self.model.layoutxlm, "use_visual_backbone")
+            and self.model.layoutxlm.use_visual_backbone is False
+        ):
+            self.use_visual_backbone = False
+
+    def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+        else:
+            image = None
+        x = self.model(
+            input_ids=x[0],
+            bbox=x[1],
+            attention_mask=x[2],
+            token_type_ids=x[3],
+            image=image,
+            position_ids=None,
+            head_mask=None,
+            labels=None,
+        )
+        if self.training:
+            res = {"backbone_out": x[0]}
+            res.update(x[1])
+            return res
+        else:
+            return x
+
+
+class LayoutLMv2ForRe(NLPBaseModel):
+    def __init__(self, pretrained=True, checkpoints=None, mode="base", **kwargs):
+        super(LayoutLMv2ForRe, self).__init__(
+            LayoutLMv2Model,
+            LayoutLMv2ForRelationExtraction,
+            mode,
+            "re",
+            pretrained,
+            checkpoints,
+        )
+        if (
+            hasattr(self.model.layoutlmv2, "use_visual_backbone")
+            and self.model.layoutlmv2.use_visual_backbone is False
+        ):
+            self.use_visual_backbone = False
+
+    def forward(self, x):
+        x = self.model(
+            input_ids=x[0],
+            bbox=x[1],
+            attention_mask=x[2],
+            token_type_ids=x[3],
+            image=x[4],
+            position_ids=None,
+            head_mask=None,
+            labels=None,
+            entities=x[5],
+            relations=x[6],
+        )
+        return x
+
+
+class LayoutXLMForRe(NLPBaseModel):
+    def __init__(self, pretrained=True, checkpoints=None, mode="base", **kwargs):
+        super(LayoutXLMForRe, self).__init__(
+            LayoutXLMModel,
+            LayoutXLMForRelationExtraction,
+            mode,
+            "re",
+            pretrained,
+            checkpoints,
+        )
+        if (
+            hasattr(self.model.layoutxlm, "use_visual_backbone")
+            and self.model.layoutxlm.use_visual_backbone is False
+        ):
+            self.use_visual_backbone = False
+
+    def forward(self, x):
+        if self.use_visual_backbone is True:
+            image = x[4]
+            entities = x[5]
+            relations = x[6]
+        else:
+            image = None
+            entities = x[4]
+            relations = x[5]
+        x = self.model(
+            input_ids=x[0],
+            bbox=x[1],
+            attention_mask=x[2],
+            token_type_ids=x[3],
+            image=image,
+            position_ids=None,
+            head_mask=None,
+            labels=None,
+            entities=entities,
+            relations=relations,
+        )
+        return x

ppocr/modeling/heads/__init__.py

@@ -0,0 +1,108 @@
+# 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_head"]
+
+
+def build_head(config):
+    # det head
+    from .det_db_head import DBHead, PFHeadLocal
+    from .det_east_head import EASTHead
+    from .det_sast_head import SASTHead
+    from .det_pse_head import PSEHead
+    from .det_fce_head import FCEHead
+    from .e2e_pg_head import PGHead
+    from .det_ct_head import CT_Head
+
+    # rec head
+    from .rec_ctc_head import CTCHead
+    from .rec_att_head import AttentionHead
+    from .rec_srn_head import SRNHead
+    from .rec_nrtr_head import Transformer
+    from .rec_sar_head import SARHead
+    from .rec_aster_head import AsterHead
+    from .rec_pren_head import PRENHead
+    from .rec_multi_head import MultiHead
+    from .rec_spin_att_head import SPINAttentionHead
+    from .rec_abinet_head import ABINetHead
+    from .rec_robustscanner_head import RobustScannerHead
+    from .rec_visionlan_head import VLHead
+    from .rec_rfl_head import RFLHead
+    from .rec_can_head import CANHead
+    from .rec_latexocr_head import LaTeXOCRHead
+    from .rec_satrn_head import SATRNHead
+    from .rec_parseq_head import ParseQHead
+    from .rec_cppd_head import CPPDHead
+    from .rec_unimernet_head import UniMERNetHead
+    from .rec_ppformulanet_head import PPFormulaNet_Head
+
+    # cls head
+    from .cls_head import ClsHead
+
+    # kie head
+    from .kie_sdmgr_head import SDMGRHead
+
+    from .table_att_head import TableAttentionHead, SLAHead
+    from .table_master_head import TableMasterHead
+
+    support_dict = [
+        "DBHead",
+        "PSEHead",
+        "FCEHead",
+        "EASTHead",
+        "SASTHead",
+        "CTCHead",
+        "ClsHead",
+        "AttentionHead",
+        "SRNHead",
+        "PGHead",
+        "Transformer",
+        "TableAttentionHead",
+        "SARHead",
+        "AsterHead",
+        "SDMGRHead",
+        "PRENHead",
+        "MultiHead",
+        "ABINetHead",
+        "TableMasterHead",
+        "SPINAttentionHead",
+        "VLHead",
+        "SLAHead",
+        "RobustScannerHead",
+        "CT_Head",
+        "RFLHead",
+        "DRRGHead",
+        "CANHead",
+        "LaTeXOCRHead",
+        "SATRNHead",
+        "PFHeadLocal",
+        "ParseQHead",
+        "CPPDHead",
+        "UniMERNetHead",
+        "PPFormulaNet_Head",
+    ]
+
+    if config["name"] == "DRRGHead":
+        from .det_drrg_head import DRRGHead
+
+        support_dict.append("DRRGHead")
+
+    # table head
+
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "head only support {}".format(support_dict)
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/modeling/heads/cls_head.py

@@ -0,0 +1,53 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn, ParamAttr
+import paddle.nn.functional as F
+
+
+class ClsHead(nn.Layer):
+    """
+    Class orientation
+
+    Args:
+
+        params(dict): super parameters for build Class network
+    """
+
+    def __init__(self, in_channels, class_dim, **kwargs):
+        super(ClsHead, self).__init__()
+        self.pool = nn.AdaptiveAvgPool2D(1)
+        stdv = 1.0 / math.sqrt(in_channels * 1.0)
+        self.fc = nn.Linear(
+            in_channels,
+            class_dim,
+            weight_attr=ParamAttr(
+                name="fc_0.w_0", initializer=nn.initializer.Uniform(-stdv, stdv)
+            ),
+            bias_attr=ParamAttr(name="fc_0.b_0"),
+        )
+
+    def forward(self, x, targets=None):
+        x = self.pool(x)
+        x = paddle.reshape(x, shape=[x.shape[0], x.shape[1]])
+        x = self.fc(x)
+        if not self.training:
+            x = F.softmax(x, axis=1)
+        return x

ppocr/modeling/heads/det_ct_head.py

@@ -0,0 +1,69 @@
+# copyright (c) 2019 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
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+import math
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+ones_ = Constant(value=1.0)
+zeros_ = Constant(value=0.0)
+
+
+class CT_Head(nn.Layer):
+    def __init__(
+        self, in_channels, hidden_dim, num_classes, loss_kernel=None, loss_loc=None
+    ):
+        super(CT_Head, self).__init__()
+        self.conv1 = nn.Conv2D(
+            in_channels, hidden_dim, kernel_size=3, stride=1, padding=1
+        )
+        self.bn1 = nn.BatchNorm2D(hidden_dim)
+        self.relu1 = nn.ReLU()
+
+        self.conv2 = nn.Conv2D(
+            hidden_dim, num_classes, kernel_size=1, stride=1, padding=0
+        )
+
+        for m in self.sublayers():
+            if isinstance(m, nn.Conv2D):
+                n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels
+                normal_ = Normal(mean=0.0, std=math.sqrt(2.0 / n))
+                normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2D):
+                zeros_(m.bias)
+                ones_(m.weight)
+
+    def _upsample(self, x, scale=1):
+        return F.upsample(x, scale_factor=scale, mode="bilinear")
+
+    def forward(self, f, targets=None):
+        out = self.conv1(f)
+        out = self.relu1(self.bn1(out))
+        out = self.conv2(out)
+
+        if self.training:
+            out = self._upsample(out, scale=4)
+            return {"maps": out}
+        else:
+            score = F.sigmoid(out[:, 0, :, :])
+            return {"maps": out, "score": score}

ppocr/modeling/heads/det_db_head.py

@@ -0,0 +1,159 @@
+# copyright (c) 2019 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
+import paddle.nn.functional as F
+from paddle import ParamAttr
+from ppocr.modeling.backbones.det_mobilenet_v3 import ConvBNLayer
+
+
+def get_bias_attr(k):
+    stdv = 1.0 / math.sqrt(k * 1.0)
+    initializer = paddle.nn.initializer.Uniform(-stdv, stdv)
+    bias_attr = ParamAttr(initializer=initializer)
+    return bias_attr
+
+
+class Head(nn.Layer):
+    def __init__(self, in_channels, kernel_list=[3, 2, 2], fix_nan=False, **kwargs):
+        super(Head, self).__init__()
+
+        self.conv1 = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=in_channels // 4,
+            kernel_size=kernel_list[0],
+            padding=int(kernel_list[0] // 2),
+            weight_attr=ParamAttr(),
+            bias_attr=False,
+        )
+        self.conv_bn1 = nn.BatchNorm(
+            num_channels=in_channels // 4,
+            param_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1.0)),
+            bias_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1e-4)),
+            act="relu",
+        )
+
+        self.conv2 = nn.Conv2DTranspose(
+            in_channels=in_channels // 4,
+            out_channels=in_channels // 4,
+            kernel_size=kernel_list[1],
+            stride=2,
+            weight_attr=ParamAttr(initializer=paddle.nn.initializer.KaimingUniform()),
+            bias_attr=get_bias_attr(in_channels // 4),
+        )
+        self.conv_bn2 = nn.BatchNorm(
+            num_channels=in_channels // 4,
+            param_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1.0)),
+            bias_attr=ParamAttr(initializer=paddle.nn.initializer.Constant(value=1e-4)),
+            act="relu",
+        )
+        self.conv3 = nn.Conv2DTranspose(
+            in_channels=in_channels // 4,
+            out_channels=1,
+            kernel_size=kernel_list[2],
+            stride=2,
+            weight_attr=ParamAttr(initializer=paddle.nn.initializer.KaimingUniform()),
+            bias_attr=get_bias_attr(in_channels // 4),
+        )
+
+        self.fix_nan = fix_nan
+
+    def forward(self, x, return_f=False):
+        x = self.conv1(x)
+        x = self.conv_bn1(x)
+        if self.fix_nan and self.training:
+            x = paddle.where(paddle.isnan(x), paddle.zeros_like(x), x)
+        x = self.conv2(x)
+        x = self.conv_bn2(x)
+        if self.fix_nan and self.training:
+            x = paddle.where(paddle.isnan(x), paddle.zeros_like(x), x)
+        if return_f is True:
+            f = x
+        x = self.conv3(x)
+        x = F.sigmoid(x)
+        if return_f is True:
+            return x, f
+        return x
+
+
+class DBHead(nn.Layer):
+    """
+    Differentiable Binarization (DB) for text detection:
+        see https://arxiv.org/abs/1911.08947
+    args:
+        params(dict): super parameters for build DB network
+    """
+
+    def __init__(self, in_channels, k=50, **kwargs):
+        super(DBHead, self).__init__()
+        self.k = k
+        self.binarize = Head(in_channels, **kwargs)
+        self.thresh = Head(in_channels, **kwargs)
+
+    def step_function(self, x, y):
+        return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))
+
+    def forward(self, x, targets=None):
+        shrink_maps = self.binarize(x)
+        if not self.training:
+            return {"maps": shrink_maps}
+
+        threshold_maps = self.thresh(x)
+        binary_maps = self.step_function(shrink_maps, threshold_maps)
+        y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1)
+        return {"maps": y}
+
+
+class LocalModule(nn.Layer):
+    def __init__(self, in_c, mid_c, use_distance=True):
+        super(self.__class__, self).__init__()
+        self.last_3 = ConvBNLayer(in_c + 1, mid_c, 3, 1, 1, act="relu")
+        self.last_1 = nn.Conv2D(mid_c, 1, 1, 1, 0)
+
+    def forward(self, x, init_map, distance_map):
+        outf = paddle.concat([init_map, x], axis=1)
+        # last Conv
+        out = self.last_1(self.last_3(outf))
+        return out
+
+
+class PFHeadLocal(DBHead):
+    def __init__(self, in_channels, k=50, mode="small", **kwargs):
+        super(PFHeadLocal, self).__init__(in_channels, k, **kwargs)
+        self.mode = mode
+
+        self.up_conv = nn.Upsample(scale_factor=2, mode="nearest", align_mode=1)
+        if self.mode == "large":
+            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 4)
+        elif self.mode == "small":
+            self.cbn_layer = LocalModule(in_channels // 4, in_channels // 8)
+
+    def forward(self, x, targets=None):
+        shrink_maps, f = self.binarize(x, return_f=True)
+        base_maps = shrink_maps
+        cbn_maps = self.cbn_layer(self.up_conv(f), shrink_maps, None)
+        cbn_maps = F.sigmoid(cbn_maps)
+        if not self.training:
+            return {"maps": 0.5 * (base_maps + cbn_maps)}
+
+        threshold_maps = self.thresh(x)
+        binary_maps = self.step_function(shrink_maps, threshold_maps)
+        y = paddle.concat([cbn_maps, threshold_maps, binary_maps], axis=1)
+        return {"maps": y, "distance_maps": cbn_maps, "cbn_maps": binary_maps}

ppocr/modeling/heads/det_drrg_head.py

@@ -0,0 +1,216 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/dense_heads/drrg_head.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import warnings
+import cv2
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from .gcn import GCN
+from .local_graph import LocalGraphs
+from .proposal_local_graph import ProposalLocalGraphs
+
+
+class DRRGHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        k_at_hops=(8, 4),
+        num_adjacent_linkages=3,
+        node_geo_feat_len=120,
+        pooling_scale=1.0,
+        pooling_output_size=(4, 3),
+        nms_thr=0.3,
+        min_width=8.0,
+        max_width=24.0,
+        comp_shrink_ratio=1.03,
+        comp_ratio=0.4,
+        comp_score_thr=0.3,
+        text_region_thr=0.2,
+        center_region_thr=0.2,
+        center_region_area_thr=50,
+        local_graph_thr=0.7,
+        **kwargs,
+    ):
+        super().__init__()
+
+        assert isinstance(in_channels, int)
+        assert isinstance(k_at_hops, tuple)
+        assert isinstance(num_adjacent_linkages, int)
+        assert isinstance(node_geo_feat_len, int)
+        assert isinstance(pooling_scale, float)
+        assert isinstance(pooling_output_size, tuple)
+        assert isinstance(comp_shrink_ratio, float)
+        assert isinstance(nms_thr, float)
+        assert isinstance(min_width, float)
+        assert isinstance(max_width, float)
+        assert isinstance(comp_ratio, float)
+        assert isinstance(comp_score_thr, float)
+        assert isinstance(text_region_thr, float)
+        assert isinstance(center_region_thr, float)
+        assert isinstance(center_region_area_thr, int)
+        assert isinstance(local_graph_thr, float)
+
+        self.in_channels = in_channels
+        self.out_channels = 6
+        self.downsample_ratio = 1.0
+        self.k_at_hops = k_at_hops
+        self.num_adjacent_linkages = num_adjacent_linkages
+        self.node_geo_feat_len = node_geo_feat_len
+        self.pooling_scale = pooling_scale
+        self.pooling_output_size = pooling_output_size
+        self.comp_shrink_ratio = comp_shrink_ratio
+        self.nms_thr = nms_thr
+        self.min_width = min_width
+        self.max_width = max_width
+        self.comp_ratio = comp_ratio
+        self.comp_score_thr = comp_score_thr
+        self.text_region_thr = text_region_thr
+        self.center_region_thr = center_region_thr
+        self.center_region_area_thr = center_region_area_thr
+        self.local_graph_thr = local_graph_thr
+
+        self.out_conv = nn.Conv2D(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+        self.graph_train = LocalGraphs(
+            self.k_at_hops,
+            self.num_adjacent_linkages,
+            self.node_geo_feat_len,
+            self.pooling_scale,
+            self.pooling_output_size,
+            self.local_graph_thr,
+        )
+
+        self.graph_test = ProposalLocalGraphs(
+            self.k_at_hops,
+            self.num_adjacent_linkages,
+            self.node_geo_feat_len,
+            self.pooling_scale,
+            self.pooling_output_size,
+            self.nms_thr,
+            self.min_width,
+            self.max_width,
+            self.comp_shrink_ratio,
+            self.comp_ratio,
+            self.comp_score_thr,
+            self.text_region_thr,
+            self.center_region_thr,
+            self.center_region_area_thr,
+        )
+
+        pool_w, pool_h = self.pooling_output_size
+        node_feat_len = (pool_w * pool_h) * (
+            self.in_channels + self.out_channels
+        ) + self.node_geo_feat_len
+        self.gcn = GCN(node_feat_len)
+
+    def forward(self, inputs, targets=None):
+        """
+        Args:
+            inputs (Tensor): Shape of :math:`(N, C, H, W)`.
+            gt_comp_attribs (list[ndarray]): The padded text component
+                attributes. Shape: (num_component, 8).
+
+        Returns:
+            tuple: Returns (pred_maps, (gcn_pred, gt_labels)).
+
+                - | pred_maps (Tensor): Prediction map with shape
+                    :math:`(N, C_{out}, H, W)`.
+                - | gcn_pred (Tensor): Prediction from GCN module, with
+                    shape :math:`(N, 2)`.
+                - | gt_labels (Tensor): Ground-truth label with shape
+                    :math:`(N, 8)`.
+        """
+        if self.training:
+            assert targets is not None
+            gt_comp_attribs = targets[7]
+            pred_maps = self.out_conv(inputs)
+            feat_maps = paddle.concat([inputs, pred_maps], axis=1)
+            node_feats, adjacent_matrices, knn_inds, gt_labels = self.graph_train(
+                feat_maps, np.stack(gt_comp_attribs)
+            )
+
+            gcn_pred = self.gcn(node_feats, adjacent_matrices, knn_inds)
+
+            return pred_maps, (gcn_pred, gt_labels)
+        else:
+            return self.single_test(inputs)
+
+    def single_test(self, feat_maps):
+        r"""
+        Args:
+            feat_maps (Tensor): Shape of :math:`(N, C, H, W)`.
+
+        Returns:
+            tuple: Returns (edge, score, text_comps).
+
+                - | edge (ndarray): The edge array of shape :math:`(N, 2)`
+                    where each row is a pair of text component indices
+                    that makes up an edge in graph.
+                - | score (ndarray): The score array of shape :math:`(N,)`,
+                    corresponding to the edge above.
+                - | text_comps (ndarray): The text components of shape
+                    :math:`(N, 9)` where each row corresponds to one box and
+                    its score: (x1, y1, x2, y2, x3, y3, x4, y4, score).
+        """
+        pred_maps = self.out_conv(feat_maps)
+        feat_maps = paddle.concat([feat_maps, pred_maps], axis=1)
+
+        none_flag, graph_data = self.graph_test(pred_maps, feat_maps)
+
+        (
+            local_graphs_node_feat,
+            adjacent_matrices,
+            pivots_knn_inds,
+            pivot_local_graphs,
+            text_comps,
+        ) = graph_data
+
+        if none_flag:
+            return None, None, None
+        gcn_pred = self.gcn(local_graphs_node_feat, adjacent_matrices, pivots_knn_inds)
+        pred_labels = F.softmax(gcn_pred, axis=1)
+
+        edges = []
+        scores = []
+        pivot_local_graphs = pivot_local_graphs.squeeze().numpy()
+
+        for pivot_ind, pivot_local_graph in enumerate(pivot_local_graphs):
+            pivot = pivot_local_graph[0]
+            for k_ind, neighbor_ind in enumerate(pivots_knn_inds[pivot_ind]):
+                neighbor = pivot_local_graph[neighbor_ind.item()]
+                edges.append([pivot, neighbor])
+                scores.append(
+                    pred_labels[pivot_ind * pivots_knn_inds.shape[1] + k_ind, 1].item()
+                )
+
+        edges = np.asarray(edges)
+        scores = np.asarray(scores)
+
+        return edges, scores, text_comps

ppocr/modeling/heads/det_east_head.py

@@ -0,0 +1,129 @@
+# copyright (c) 2019 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
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="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 EASTHead(nn.Layer):
+    """ """
+
+    def __init__(self, in_channels, model_name, **kwargs):
+        super(EASTHead, self).__init__()
+        self.model_name = model_name
+        if self.model_name == "large":
+            num_outputs = [128, 64, 1, 8]
+        else:
+            num_outputs = [64, 32, 1, 8]
+
+        self.det_conv1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=num_outputs[0],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="det_head1",
+        )
+        self.det_conv2 = ConvBNLayer(
+            in_channels=num_outputs[0],
+            out_channels=num_outputs[1],
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="det_head2",
+        )
+        self.score_conv = ConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[2],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=False,
+            act=None,
+            name="f_score",
+        )
+        self.geo_conv = ConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[3],
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            if_act=False,
+            act=None,
+            name="f_geo",
+        )
+
+    def forward(self, x, targets=None):
+        f_det = self.det_conv1(x)
+        f_det = self.det_conv2(f_det)
+        f_score = self.score_conv(f_det)
+        f_score = F.sigmoid(f_score)
+        f_geo = self.geo_conv(f_det)
+        f_geo = (F.sigmoid(f_geo) - 0.5) * 2 * 800
+
+        pred = {"f_score": f_score, "f_geo": f_geo}
+        return pred

ppocr/modeling/heads/det_fce_head.py

@@ -0,0 +1,100 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/dense_heads/fce_head.py
+"""
+
+from paddle import nn
+from paddle import ParamAttr
+import paddle.nn.functional as F
+from paddle.nn.initializer import Normal
+import paddle
+from functools import partial
+
+
+def multi_apply(func, *args, **kwargs):
+    pfunc = partial(func, **kwargs) if kwargs else func
+    map_results = map(pfunc, *args)
+    return tuple(map(list, zip(*map_results)))
+
+
+class FCEHead(nn.Layer):
+    """The class for implementing FCENet head.
+    FCENet(CVPR2021): Fourier Contour Embedding for Arbitrary-shaped Text
+    Detection.
+
+    [https://arxiv.org/abs/2104.10442]
+
+    Args:
+        in_channels (int): The number of input channels.
+        scales (list[int]) : The scale of each layer.
+        fourier_degree (int) : The maximum Fourier transform degree k.
+    """
+
+    def __init__(self, in_channels, fourier_degree=5):
+        super().__init__()
+        assert isinstance(in_channels, int)
+
+        self.downsample_ratio = 1.0
+        self.in_channels = in_channels
+        self.fourier_degree = fourier_degree
+        self.out_channels_cls = 4
+        self.out_channels_reg = (2 * self.fourier_degree + 1) * 2
+
+        self.out_conv_cls = nn.Conv2D(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels_cls,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(
+                name="cls_weights", initializer=Normal(mean=0.0, std=0.01)
+            ),
+            bias_attr=True,
+        )
+        self.out_conv_reg = nn.Conv2D(
+            in_channels=self.in_channels,
+            out_channels=self.out_channels_reg,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(
+                name="reg_weights", initializer=Normal(mean=0.0, std=0.01)
+            ),
+            bias_attr=True,
+        )
+
+    def forward(self, feats, targets=None):
+        cls_res, reg_res = multi_apply(self.forward_single, feats)
+        level_num = len(cls_res)
+        outs = {}
+        if not self.training:
+            for i in range(level_num):
+                tr_pred = F.softmax(cls_res[i][:, 0:2, :, :], axis=1)
+                tcl_pred = F.softmax(cls_res[i][:, 2:, :, :], axis=1)
+                outs["level_{}".format(i)] = paddle.concat(
+                    [tr_pred, tcl_pred, reg_res[i]], axis=1
+                )
+        else:
+            preds = [[cls_res[i], reg_res[i]] for i in range(level_num)]
+            outs["levels"] = preds
+        return outs
+
+    def forward_single(self, x):
+        cls_predict = self.out_conv_cls(x)
+        reg_predict = self.out_conv_reg(x)
+        return cls_predict, reg_predict

ppocr/modeling/heads/det_pse_head.py

@@ -0,0 +1,39 @@
+# copyright (c) 2021 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/whai362/PSENet/blob/python3/models/head/psenet_head.py
+"""
+
+from paddle import nn
+
+
+class PSEHead(nn.Layer):
+    def __init__(self, in_channels, hidden_dim=256, out_channels=7, **kwargs):
+        super(PSEHead, self).__init__()
+        self.conv1 = nn.Conv2D(
+            in_channels, hidden_dim, kernel_size=3, stride=1, padding=1
+        )
+        self.bn1 = nn.BatchNorm2D(hidden_dim)
+        self.relu1 = nn.ReLU()
+
+        self.conv2 = nn.Conv2D(
+            hidden_dim, out_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x, **kwargs):
+        out = self.conv1(x)
+        out = self.relu1(self.bn1(out))
+        out = self.conv2(out)
+        return {"maps": out}

ppocr/modeling/heads/det_sast_head.py

@@ -0,0 +1,152 @@
+# copyright (c) 2019 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
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        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,
+        )
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="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 SAST_Header1(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super(SAST_Header1, self).__init__()
+        out_channels = [64, 64, 128]
+        self.score_conv = nn.Sequential(
+            ConvBNLayer(
+                in_channels, out_channels[0], 1, 1, act="relu", name="f_score1"
+            ),
+            ConvBNLayer(
+                out_channels[0], out_channels[1], 3, 1, act="relu", name="f_score2"
+            ),
+            ConvBNLayer(
+                out_channels[1], out_channels[2], 1, 1, act="relu", name="f_score3"
+            ),
+            ConvBNLayer(out_channels[2], 1, 3, 1, act=None, name="f_score4"),
+        )
+        self.border_conv = nn.Sequential(
+            ConvBNLayer(
+                in_channels, out_channels[0], 1, 1, act="relu", name="f_border1"
+            ),
+            ConvBNLayer(
+                out_channels[0], out_channels[1], 3, 1, act="relu", name="f_border2"
+            ),
+            ConvBNLayer(
+                out_channels[1], out_channels[2], 1, 1, act="relu", name="f_border3"
+            ),
+            ConvBNLayer(out_channels[2], 4, 3, 1, act=None, name="f_border4"),
+        )
+
+    def forward(self, x):
+        f_score = self.score_conv(x)
+        f_score = F.sigmoid(f_score)
+        f_border = self.border_conv(x)
+        return f_score, f_border
+
+
+class SAST_Header2(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super(SAST_Header2, self).__init__()
+        out_channels = [64, 64, 128]
+        self.tvo_conv = nn.Sequential(
+            ConvBNLayer(in_channels, out_channels[0], 1, 1, act="relu", name="f_tvo1"),
+            ConvBNLayer(
+                out_channels[0], out_channels[1], 3, 1, act="relu", name="f_tvo2"
+            ),
+            ConvBNLayer(
+                out_channels[1], out_channels[2], 1, 1, act="relu", name="f_tvo3"
+            ),
+            ConvBNLayer(out_channels[2], 8, 3, 1, act=None, name="f_tvo4"),
+        )
+        self.tco_conv = nn.Sequential(
+            ConvBNLayer(in_channels, out_channels[0], 1, 1, act="relu", name="f_tco1"),
+            ConvBNLayer(
+                out_channels[0], out_channels[1], 3, 1, act="relu", name="f_tco2"
+            ),
+            ConvBNLayer(
+                out_channels[1], out_channels[2], 1, 1, act="relu", name="f_tco3"
+            ),
+            ConvBNLayer(out_channels[2], 2, 3, 1, act=None, name="f_tco4"),
+        )
+
+    def forward(self, x):
+        f_tvo = self.tvo_conv(x)
+        f_tco = self.tco_conv(x)
+        return f_tvo, f_tco
+
+
+class SASTHead(nn.Layer):
+    """ """
+
+    def __init__(self, in_channels, **kwargs):
+        super(SASTHead, self).__init__()
+
+        self.head1 = SAST_Header1(in_channels)
+        self.head2 = SAST_Header2(in_channels)
+
+    def forward(self, x, targets=None):
+        f_score, f_border = self.head1(x)
+        f_tvo, f_tco = self.head2(x)
+
+        predicts = {}
+        predicts["f_score"] = f_score
+        predicts["f_border"] = f_border
+        predicts["f_tvo"] = f_tvo
+        predicts["f_tco"] = f_tco
+        return predicts

ppocr/modeling/heads/e2e_pg_head.py

@@ -0,0 +1,282 @@
+# copyright (c) 2021 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
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+            use_global_stats=False,
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class PGHead(nn.Layer):
+    """ """
+
+    def __init__(
+        self, in_channels, character_dict_path="ppocr/utils/ic15_dict.txt", **kwargs
+    ):
+        super(PGHead, self).__init__()
+
+        # get character_length
+        with open(character_dict_path, "rb") as fin:
+            lines = fin.readlines()
+            character_length = len(lines) + 1
+
+        self.conv_f_score1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_score{}".format(1),
+        )
+        self.conv_f_score2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act="relu",
+            name="conv_f_score{}".format(2),
+        )
+        self.conv_f_score3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_score{}".format(3),
+        )
+
+        self.conv1 = nn.Conv2D(
+            in_channels=128,
+            out_channels=1,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_score{}".format(4)),
+            bias_attr=False,
+        )
+
+        self.conv_f_boder1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_boder{}".format(1),
+        )
+        self.conv_f_boder2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act="relu",
+            name="conv_f_boder{}".format(2),
+        )
+        self.conv_f_boder3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_boder{}".format(3),
+        )
+        self.conv2 = nn.Conv2D(
+            in_channels=128,
+            out_channels=4,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_boder{}".format(4)),
+            bias_attr=False,
+        )
+        self.conv_f_char1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_char{}".format(1),
+        )
+        self.conv_f_char2 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act="relu",
+            name="conv_f_char{}".format(2),
+        )
+        self.conv_f_char3 = ConvBNLayer(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_char{}".format(3),
+        )
+        self.conv_f_char4 = ConvBNLayer(
+            in_channels=256,
+            out_channels=256,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act="relu",
+            name="conv_f_char{}".format(4),
+        )
+        self.conv_f_char5 = ConvBNLayer(
+            in_channels=256,
+            out_channels=256,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_char{}".format(5),
+        )
+        self.conv3 = nn.Conv2D(
+            in_channels=256,
+            out_channels=character_length,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_char{}".format(6)),
+            bias_attr=False,
+        )
+
+        self.conv_f_direc1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_direc{}".format(1),
+        )
+        self.conv_f_direc2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act="relu",
+            name="conv_f_direc{}".format(2),
+        )
+        self.conv_f_direc3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act="relu",
+            name="conv_f_direc{}".format(3),
+        )
+        self.conv4 = nn.Conv2D(
+            in_channels=128,
+            out_channels=2,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_direc{}".format(4)),
+            bias_attr=False,
+        )
+
+    def forward(self, x, targets=None):
+        f_score = self.conv_f_score1(x)
+        f_score = self.conv_f_score2(f_score)
+        f_score = self.conv_f_score3(f_score)
+        f_score = self.conv1(f_score)
+        f_score = F.sigmoid(f_score)
+
+        # f_border
+        f_border = self.conv_f_boder1(x)
+        f_border = self.conv_f_boder2(f_border)
+        f_border = self.conv_f_boder3(f_border)
+        f_border = self.conv2(f_border)
+
+        f_char = self.conv_f_char1(x)
+        f_char = self.conv_f_char2(f_char)
+        f_char = self.conv_f_char3(f_char)
+        f_char = self.conv_f_char4(f_char)
+        f_char = self.conv_f_char5(f_char)
+        f_char = self.conv3(f_char)
+
+        f_direction = self.conv_f_direc1(x)
+        f_direction = self.conv_f_direc2(f_direction)
+        f_direction = self.conv_f_direc3(f_direction)
+        f_direction = self.conv4(f_direction)
+
+        predicts = {}
+        predicts["f_score"] = f_score
+        predicts["f_border"] = f_border
+        predicts["f_char"] = f_char
+        predicts["f_direction"] = f_direction
+        return predicts

ppocr/modeling/heads/gcn.py

@@ -0,0 +1,118 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/modules/gcn.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class BatchNorm1D(nn.BatchNorm1D):
+    def __init__(
+        self,
+        num_features,
+        eps=1e-05,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+    ):
+        momentum = 1 - momentum
+        weight_attr = None
+        bias_attr = None
+        if not affine:
+            weight_attr = paddle.ParamAttr(learning_rate=0.0)
+            bias_attr = paddle.ParamAttr(learning_rate=0.0)
+        super().__init__(
+            num_features,
+            momentum=momentum,
+            epsilon=eps,
+            weight_attr=weight_attr,
+            bias_attr=bias_attr,
+            use_global_stats=track_running_stats,
+        )
+
+
+class MeanAggregator(nn.Layer):
+    def forward(self, features, A):
+        x = paddle.bmm(A, features)
+        return x
+
+
+class GraphConv(nn.Layer):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.weight = self.create_parameter(
+            [in_dim * 2, out_dim], default_initializer=nn.initializer.XavierUniform()
+        )
+        self.bias = self.create_parameter(
+            [out_dim],
+            is_bias=True,
+            default_initializer=nn.initializer.Assign([0] * out_dim),
+        )
+
+        self.aggregator = MeanAggregator()
+
+    def forward(self, features, A):
+        b, n, d = features.shape
+        assert d == self.in_dim
+        agg_feats = self.aggregator(features, A)
+        cat_feats = paddle.concat([features, agg_feats], axis=2)
+        out = paddle.einsum("bnd,df->bnf", cat_feats, self.weight)
+        out = F.relu(out + self.bias)
+        return out
+
+
+class GCN(nn.Layer):
+    def __init__(self, feat_len):
+        super(GCN, self).__init__()
+        self.bn0 = BatchNorm1D(feat_len, affine=False)
+        self.conv1 = GraphConv(feat_len, 512)
+        self.conv2 = GraphConv(512, 256)
+        self.conv3 = GraphConv(256, 128)
+        self.conv4 = GraphConv(128, 64)
+        self.classifier = nn.Sequential(
+            nn.Linear(64, 32), nn.PReLU(32), nn.Linear(32, 2)
+        )
+
+    def forward(self, x, A, knn_inds):
+        num_local_graphs, num_max_nodes, feat_len = x.shape
+
+        x = x.reshape([-1, feat_len])
+        x = self.bn0(x)
+        x = x.reshape([num_local_graphs, num_max_nodes, feat_len])
+
+        x = self.conv1(x, A)
+        x = self.conv2(x, A)
+        x = self.conv3(x, A)
+        x = self.conv4(x, A)
+        k = knn_inds.shape[-1]
+        mid_feat_len = x.shape[-1]
+        edge_feat = paddle.zeros([num_local_graphs, k, mid_feat_len])
+        for graph_ind in range(num_local_graphs):
+            edge_feat[graph_ind, :, :] = x[graph_ind][
+                paddle.to_tensor(knn_inds[graph_ind])
+            ]
+        edge_feat = edge_feat.reshape([-1, mid_feat_len])
+        pred = self.classifier(edge_feat)
+
+        return pred

ppocr/modeling/heads/kie_sdmgr_head.py

@@ -0,0 +1,223 @@
+# 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.
+# reference from : https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/kie/heads/sdmgr_head.py
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class SDMGRHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        num_chars=92,
+        visual_dim=16,
+        fusion_dim=1024,
+        node_input=32,
+        node_embed=256,
+        edge_input=5,
+        edge_embed=256,
+        num_gnn=2,
+        num_classes=26,
+        bidirectional=False,
+    ):
+        super().__init__()
+
+        self.fusion = Block([visual_dim, node_embed], node_embed, fusion_dim)
+        self.node_embed = nn.Embedding(num_chars, node_input, 0)
+        hidden = node_embed // 2 if bidirectional else node_embed
+        self.rnn = nn.LSTM(input_size=node_input, hidden_size=hidden, num_layers=1)
+        self.edge_embed = nn.Linear(edge_input, edge_embed)
+        self.gnn_layers = nn.LayerList(
+            [GNNLayer(node_embed, edge_embed) for _ in range(num_gnn)]
+        )
+        self.node_cls = nn.Linear(node_embed, num_classes)
+        self.edge_cls = nn.Linear(edge_embed, 2)
+
+    def forward(self, input, targets):
+        relations, texts, x = input
+        node_nums, char_nums = [], []
+        for text in texts:
+            node_nums.append(text.shape[0])
+            char_nums.append(paddle.sum((text > -1).astype(int), axis=-1))
+
+        max_num = max([char_num.max() for char_num in char_nums])
+        all_nodes = paddle.concat(
+            [
+                paddle.concat(
+                    [text, paddle.zeros((text.shape[0], max_num - text.shape[1]))], -1
+                )
+                for text in texts
+            ]
+        )
+        temp = paddle.clip(all_nodes, min=0).astype(int)
+        embed_nodes = self.node_embed(temp)
+        rnn_nodes, _ = self.rnn(embed_nodes)
+
+        b, h, w = rnn_nodes.shape
+        nodes = paddle.zeros([b, w])
+        all_nums = paddle.concat(char_nums)
+        valid = paddle.nonzero((all_nums > 0).astype(int))
+        temp_all_nums = (paddle.gather(all_nums, valid) - 1).unsqueeze(-1).unsqueeze(-1)
+        temp_all_nums = paddle.expand(
+            temp_all_nums,
+            [temp_all_nums.shape[0], temp_all_nums.shape[1], rnn_nodes.shape[-1]],
+        )
+        temp_all_nodes = paddle.gather(rnn_nodes, valid)
+        N, C, A = temp_all_nodes.shape
+        one_hot = F.one_hot(temp_all_nums[:, 0, :], num_classes=C).transpose([0, 2, 1])
+        one_hot = paddle.multiply(temp_all_nodes, one_hot.astype("float32")).sum(
+            axis=1, keepdim=True
+        )
+        t = one_hot.expand([N, 1, A]).squeeze(1)
+        nodes = paddle.scatter(nodes, valid.squeeze(1), t)
+
+        if x is not None:
+            nodes = self.fusion([x, nodes])
+
+        all_edges = paddle.concat(
+            [rel.reshape([-1, rel.shape[-1]]) for rel in relations]
+        )
+        embed_edges = self.edge_embed(all_edges.astype("float32"))
+        embed_edges = F.normalize(embed_edges)
+
+        for gnn_layer in self.gnn_layers:
+            nodes, cat_nodes = gnn_layer(nodes, embed_edges, node_nums)
+
+        node_cls, edge_cls = self.node_cls(nodes), self.edge_cls(cat_nodes)
+        return node_cls, edge_cls
+
+
+class GNNLayer(nn.Layer):
+    def __init__(self, node_dim=256, edge_dim=256):
+        super().__init__()
+        self.in_fc = nn.Linear(node_dim * 2 + edge_dim, node_dim)
+        self.coef_fc = nn.Linear(node_dim, 1)
+        self.out_fc = nn.Linear(node_dim, node_dim)
+        self.relu = nn.ReLU()
+
+    def forward(self, nodes, edges, nums):
+        start, cat_nodes = 0, []
+        for num in nums:
+            sample_nodes = nodes[start : start + num]
+            cat_nodes.append(
+                paddle.concat(
+                    [
+                        paddle.expand(sample_nodes.unsqueeze(1), [-1, num, -1]),
+                        paddle.expand(sample_nodes.unsqueeze(0), [num, -1, -1]),
+                    ],
+                    -1,
+                ).reshape([num**2, -1])
+            )
+            start += num
+        cat_nodes = paddle.concat([paddle.concat(cat_nodes), edges], -1)
+        cat_nodes = self.relu(self.in_fc(cat_nodes))
+        coefs = self.coef_fc(cat_nodes)
+
+        start, residuals = 0, []
+        for num in nums:
+            residual = F.softmax(
+                -paddle.eye(num).unsqueeze(-1) * 1e9
+                + coefs[start : start + num**2].reshape([num, num, -1]),
+                1,
+            )
+            residuals.append(
+                (
+                    residual * cat_nodes[start : start + num**2].reshape([num, num, -1])
+                ).sum(1)
+            )
+            start += num**2
+
+        nodes += self.relu(self.out_fc(paddle.concat(residuals)))
+        return [nodes, cat_nodes]
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        input_dims,
+        output_dim,
+        mm_dim=1600,
+        chunks=20,
+        rank=15,
+        shared=False,
+        dropout_input=0.0,
+        dropout_pre_lin=0.0,
+        dropout_output=0.0,
+        pos_norm="before_cat",
+    ):
+        super().__init__()
+        self.rank = rank
+        self.dropout_input = dropout_input
+        self.dropout_pre_lin = dropout_pre_lin
+        self.dropout_output = dropout_output
+        assert pos_norm in ["before_cat", "after_cat"]
+        self.pos_norm = pos_norm
+        # Modules
+        self.linear0 = nn.Linear(input_dims[0], mm_dim)
+        self.linear1 = self.linear0 if shared else nn.Linear(input_dims[1], mm_dim)
+        self.merge_linears0 = nn.LayerList()
+        self.merge_linears1 = nn.LayerList()
+        self.chunks = self.chunk_sizes(mm_dim, chunks)
+        for size in self.chunks:
+            ml0 = nn.Linear(size, size * rank)
+            self.merge_linears0.append(ml0)
+            ml1 = ml0 if shared else nn.Linear(size, size * rank)
+            self.merge_linears1.append(ml1)
+        self.linear_out = nn.Linear(mm_dim, output_dim)
+
+    def forward(self, x):
+        x0 = self.linear0(x[0])
+        x1 = self.linear1(x[1])
+        bs = x1.shape[0]
+        if self.dropout_input > 0:
+            x0 = F.dropout(x0, p=self.dropout_input, training=self.training)
+            x1 = F.dropout(x1, p=self.dropout_input, training=self.training)
+        x0_chunks = paddle.split(x0, self.chunks, -1)
+        x1_chunks = paddle.split(x1, self.chunks, -1)
+        zs = []
+        for x0_c, x1_c, m0, m1 in zip(
+            x0_chunks, x1_chunks, self.merge_linears0, self.merge_linears1
+        ):
+            m = m0(x0_c) * m1(x1_c)  # bs x split_size*rank
+            m = m.reshape([bs, self.rank, -1])
+            z = paddle.sum(m, 1)
+            if self.pos_norm == "before_cat":
+                z = paddle.sqrt(F.relu(z)) - paddle.sqrt(F.relu(-z))
+                z = F.normalize(z)
+            zs.append(z)
+        z = paddle.concat(zs, 1)
+        if self.pos_norm == "after_cat":
+            z = paddle.sqrt(F.relu(z)) - paddle.sqrt(F.relu(-z))
+            z = F.normalize(z)
+
+        if self.dropout_pre_lin > 0:
+            z = F.dropout(z, p=self.dropout_pre_lin, training=self.training)
+        z = self.linear_out(z)
+        if self.dropout_output > 0:
+            z = F.dropout(z, p=self.dropout_output, training=self.training)
+        return z
+
+    def chunk_sizes(self, dim, chunks):
+        split_size = (dim + chunks - 1) // chunks
+        sizes_list = [split_size] * chunks
+        sizes_list[-1] = sizes_list[-1] - (sum(sizes_list) - dim)
+        return sizes_list

ppocr/modeling/heads/local_graph.py

@@ -0,0 +1,425 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/modules/local_graph.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+from ppocr.ext_op import RoIAlignRotated
+
+
+def normalize_adjacent_matrix(A):
+    assert A.ndim == 2
+    assert A.shape[0] == A.shape[1]
+
+    A = A + np.eye(A.shape[0])
+    d = np.sum(A, axis=0)
+    d = np.clip(d, 0, None)
+    d_inv = np.power(d, -0.5).flatten()
+    d_inv[np.isinf(d_inv)] = 0.0
+    d_inv = np.diag(d_inv)
+    G = A.dot(d_inv).transpose().dot(d_inv)
+    return G
+
+
+def euclidean_distance_matrix(A, B):
+    """Calculate the Euclidean distance matrix.
+
+    Args:
+        A (ndarray): The point sequence.
+        B (ndarray): The point sequence with the same dimensions as A.
+
+    returns:
+        D (ndarray): The Euclidean distance matrix.
+    """
+    assert A.ndim == 2
+    assert B.ndim == 2
+    assert A.shape[1] == B.shape[1]
+
+    m = A.shape[0]
+    n = B.shape[0]
+
+    A_dots = (A * A).sum(axis=1).reshape((m, 1)) * np.ones(shape=(1, n))
+    B_dots = (B * B).sum(axis=1) * np.ones(shape=(m, 1))
+    D_squared = A_dots + B_dots - 2 * A.dot(B.T)
+
+    zero_mask = np.less(D_squared, 0.0)
+    D_squared[zero_mask] = 0.0
+    D = np.sqrt(D_squared)
+    return D
+
+
+def feature_embedding(input_feats, out_feat_len):
+    """Embed features. This code was partially adapted from
+    https://github.com/GXYM/DRRG licensed under the MIT license.
+
+    Args:
+        input_feats (ndarray): The input features of shape (N, d), where N is
+            the number of nodes in graph, d is the input feature vector length.
+        out_feat_len (int): The length of output feature vector.
+
+    Returns:
+        embedded_feats (ndarray): The embedded features.
+    """
+    assert input_feats.ndim == 2
+    assert isinstance(out_feat_len, int)
+    assert out_feat_len >= input_feats.shape[1]
+
+    num_nodes = input_feats.shape[0]
+    feat_dim = input_feats.shape[1]
+    feat_repeat_times = out_feat_len // feat_dim
+    residue_dim = out_feat_len % feat_dim
+
+    if residue_dim > 0:
+        embed_wave = np.array(
+            [
+                np.power(1000, 2.0 * (j // 2) / feat_repeat_times + 1)
+                for j in range(feat_repeat_times + 1)
+            ]
+        ).reshape((feat_repeat_times + 1, 1, 1))
+        repeat_feats = np.repeat(
+            np.expand_dims(input_feats, axis=0), feat_repeat_times, axis=0
+        )
+        residue_feats = np.hstack(
+            [
+                input_feats[:, 0:residue_dim],
+                np.zeros((num_nodes, feat_dim - residue_dim)),
+            ]
+        )
+        residue_feats = np.expand_dims(residue_feats, axis=0)
+        repeat_feats = np.concatenate([repeat_feats, residue_feats], axis=0)
+        embedded_feats = repeat_feats / embed_wave
+        embedded_feats[:, 0::2] = np.sin(embedded_feats[:, 0::2])
+        embedded_feats[:, 1::2] = np.cos(embedded_feats[:, 1::2])
+        embedded_feats = np.transpose(embedded_feats, (1, 0, 2)).reshape(
+            (num_nodes, -1)
+        )[:, 0:out_feat_len]
+    else:
+        embed_wave = np.array(
+            [
+                np.power(1000, 2.0 * (j // 2) / feat_repeat_times)
+                for j in range(feat_repeat_times)
+            ]
+        ).reshape((feat_repeat_times, 1, 1))
+        repeat_feats = np.repeat(
+            np.expand_dims(input_feats, axis=0), feat_repeat_times, axis=0
+        )
+        embedded_feats = repeat_feats / embed_wave
+        embedded_feats[:, 0::2] = np.sin(embedded_feats[:, 0::2])
+        embedded_feats[:, 1::2] = np.cos(embedded_feats[:, 1::2])
+        embedded_feats = (
+            np.transpose(embedded_feats, (1, 0, 2))
+            .reshape((num_nodes, -1))
+            .astype(np.float32)
+        )
+
+    return embedded_feats
+
+
+class LocalGraphs:
+    def __init__(
+        self,
+        k_at_hops,
+        num_adjacent_linkages,
+        node_geo_feat_len,
+        pooling_scale,
+        pooling_output_size,
+        local_graph_thr,
+    ):
+        assert len(k_at_hops) == 2
+        assert all(isinstance(n, int) for n in k_at_hops)
+        assert isinstance(num_adjacent_linkages, int)
+        assert isinstance(node_geo_feat_len, int)
+        assert isinstance(pooling_scale, float)
+        assert all(isinstance(n, int) for n in pooling_output_size)
+        assert isinstance(local_graph_thr, float)
+
+        self.k_at_hops = k_at_hops
+        self.num_adjacent_linkages = num_adjacent_linkages
+        self.node_geo_feat_dim = node_geo_feat_len
+        self.pooling = RoIAlignRotated(pooling_output_size, pooling_scale)
+        self.local_graph_thr = local_graph_thr
+
+    def generate_local_graphs(self, sorted_dist_inds, gt_comp_labels):
+        """Generate local graphs for GCN to predict which instance a text
+        component belongs to.
+
+        Args:
+            sorted_dist_inds (ndarray): The complete graph node indices, which
+                is sorted according to the Euclidean distance.
+            gt_comp_labels(ndarray): The ground truth labels define the
+                instance to which the text components (nodes in graphs) belong.
+
+        Returns:
+            pivot_local_graphs(list[list[int]]): The list of local graph
+                neighbor indices of pivots.
+            pivot_knns(list[list[int]]): The list of k-nearest neighbor indices
+                of pivots.
+        """
+
+        assert sorted_dist_inds.ndim == 2
+        assert (
+            sorted_dist_inds.shape[0]
+            == sorted_dist_inds.shape[1]
+            == gt_comp_labels.shape[0]
+        )
+
+        knn_graph = sorted_dist_inds[:, 1 : self.k_at_hops[0] + 1]
+        pivot_local_graphs = []
+        pivot_knns = []
+        for pivot_ind, knn in enumerate(knn_graph):
+            local_graph_neighbors = set(knn)
+
+            for neighbor_ind in knn:
+                local_graph_neighbors.update(
+                    set(sorted_dist_inds[neighbor_ind, 1 : self.k_at_hops[1] + 1])
+                )
+
+            local_graph_neighbors.discard(pivot_ind)
+            pivot_local_graph = list(local_graph_neighbors)
+            pivot_local_graph.insert(0, pivot_ind)
+            pivot_knn = [pivot_ind] + list(knn)
+
+            if pivot_ind < 1:
+                pivot_local_graphs.append(pivot_local_graph)
+                pivot_knns.append(pivot_knn)
+            else:
+                add_flag = True
+                for graph_ind, added_knn in enumerate(pivot_knns):
+                    added_pivot_ind = added_knn[0]
+                    added_local_graph = pivot_local_graphs[graph_ind]
+
+                    union = len(
+                        set(pivot_local_graph[1:]).union(set(added_local_graph[1:]))
+                    )
+                    intersect = len(
+                        set(pivot_local_graph[1:]).intersection(
+                            set(added_local_graph[1:])
+                        )
+                    )
+                    local_graph_iou = intersect / (union + 1e-8)
+
+                    if (
+                        local_graph_iou > self.local_graph_thr
+                        and pivot_ind in added_knn
+                        and gt_comp_labels[added_pivot_ind] == gt_comp_labels[pivot_ind]
+                        and gt_comp_labels[pivot_ind] != 0
+                    ):
+                        add_flag = False
+                        break
+                if add_flag:
+                    pivot_local_graphs.append(pivot_local_graph)
+                    pivot_knns.append(pivot_knn)
+
+        return pivot_local_graphs, pivot_knns
+
+    def generate_gcn_input(
+        self,
+        node_feat_batch,
+        node_label_batch,
+        local_graph_batch,
+        knn_batch,
+        sorted_dist_ind_batch,
+    ):
+        """Generate graph convolution network input data.
+
+        Args:
+            node_feat_batch (List[Tensor]): The batched graph node features.
+            node_label_batch (List[ndarray]): The batched text component
+                labels.
+            local_graph_batch (List[List[list[int]]]): The local graph node
+                indices of image batch.
+            knn_batch (List[List[list[int]]]): The knn graph node indices of
+                image batch.
+            sorted_dist_ind_batch (list[ndarray]): The node indices sorted
+                according to the Euclidean distance.
+
+        Returns:
+            local_graphs_node_feat (Tensor): The node features of graph.
+            adjacent_matrices (Tensor): The adjacent matrices of local graphs.
+            pivots_knn_inds (Tensor): The k-nearest neighbor indices in
+                local graph.
+            gt_linkage (Tensor): The surpervision signal of GCN for linkage
+                prediction.
+        """
+        assert isinstance(node_feat_batch, list)
+        assert isinstance(node_label_batch, list)
+        assert isinstance(local_graph_batch, list)
+        assert isinstance(knn_batch, list)
+        assert isinstance(sorted_dist_ind_batch, list)
+
+        num_max_nodes = max(
+            [
+                len(pivot_local_graph)
+                for pivot_local_graphs in local_graph_batch
+                for pivot_local_graph in pivot_local_graphs
+            ]
+        )
+
+        local_graphs_node_feat = []
+        adjacent_matrices = []
+        pivots_knn_inds = []
+        pivots_gt_linkage = []
+
+        for batch_ind, sorted_dist_inds in enumerate(sorted_dist_ind_batch):
+            node_feats = node_feat_batch[batch_ind]
+            pivot_local_graphs = local_graph_batch[batch_ind]
+            pivot_knns = knn_batch[batch_ind]
+            node_labels = node_label_batch[batch_ind]
+
+            for graph_ind, pivot_knn in enumerate(pivot_knns):
+                pivot_local_graph = pivot_local_graphs[graph_ind]
+                num_nodes = len(pivot_local_graph)
+                pivot_ind = pivot_local_graph[0]
+                node2ind_map = {j: i for i, j in enumerate(pivot_local_graph)}
+
+                knn_inds = paddle.to_tensor([node2ind_map[i] for i in pivot_knn[1:]])
+                pivot_feats = node_feats[pivot_ind]
+                normalized_feats = (
+                    node_feats[paddle.to_tensor(pivot_local_graph)] - pivot_feats
+                )
+
+                adjacent_matrix = np.zeros((num_nodes, num_nodes), dtype=np.float32)
+                for node in pivot_local_graph:
+                    neighbors = sorted_dist_inds[
+                        node, 1 : self.num_adjacent_linkages + 1
+                    ]
+                    for neighbor in neighbors:
+                        if neighbor in pivot_local_graph:
+                            adjacent_matrix[
+                                node2ind_map[node], node2ind_map[neighbor]
+                            ] = 1
+                            adjacent_matrix[
+                                node2ind_map[neighbor], node2ind_map[node]
+                            ] = 1
+
+                adjacent_matrix = normalize_adjacent_matrix(adjacent_matrix)
+                pad_adjacent_matrix = paddle.zeros((num_max_nodes, num_max_nodes))
+                pad_adjacent_matrix[:num_nodes, :num_nodes] = paddle.cast(
+                    paddle.to_tensor(adjacent_matrix), "float32"
+                )
+
+                pad_normalized_feats = paddle.concat(
+                    [
+                        normalized_feats,
+                        paddle.zeros(
+                            (num_max_nodes - num_nodes, normalized_feats.shape[1])
+                        ),
+                    ],
+                    axis=0,
+                )
+                local_graph_labels = node_labels[pivot_local_graph]
+                knn_labels = local_graph_labels[knn_inds.numpy()]
+                link_labels = (
+                    (node_labels[pivot_ind] == knn_labels)
+                    & (node_labels[pivot_ind] > 0)
+                ).astype(np.int64)
+                link_labels = paddle.to_tensor(link_labels)
+
+                local_graphs_node_feat.append(pad_normalized_feats)
+                adjacent_matrices.append(pad_adjacent_matrix)
+                pivots_knn_inds.append(knn_inds)
+                pivots_gt_linkage.append(link_labels)
+
+        local_graphs_node_feat = paddle.stack(local_graphs_node_feat, 0)
+        adjacent_matrices = paddle.stack(adjacent_matrices, 0)
+        pivots_knn_inds = paddle.stack(pivots_knn_inds, 0)
+        pivots_gt_linkage = paddle.stack(pivots_gt_linkage, 0)
+
+        return (
+            local_graphs_node_feat,
+            adjacent_matrices,
+            pivots_knn_inds,
+            pivots_gt_linkage,
+        )
+
+    def __call__(self, feat_maps, comp_attribs):
+        """Generate local graphs as GCN input.
+
+        Args:
+            feat_maps (Tensor): The feature maps to extract the content
+                features of text components.
+            comp_attribs (ndarray): The text component attributes.
+
+        Returns:
+            local_graphs_node_feat (Tensor): The node features of graph.
+            adjacent_matrices (Tensor): The adjacent matrices of local graphs.
+            pivots_knn_inds (Tensor): The k-nearest neighbor indices in local
+                graph.
+            gt_linkage (Tensor): The surpervision signal of GCN for linkage
+                prediction.
+        """
+
+        assert isinstance(feat_maps, paddle.Tensor)
+        assert comp_attribs.ndim == 3
+        assert comp_attribs.shape[2] == 8
+
+        sorted_dist_inds_batch = []
+        local_graph_batch = []
+        knn_batch = []
+        node_feat_batch = []
+        node_label_batch = []
+
+        for batch_ind in range(comp_attribs.shape[0]):
+            num_comps = int(comp_attribs[batch_ind, 0, 0])
+            comp_geo_attribs = comp_attribs[batch_ind, :num_comps, 1:7]
+            node_labels = comp_attribs[batch_ind, :num_comps, 7].astype(np.int32)
+
+            comp_centers = comp_geo_attribs[:, 0:2]
+            distance_matrix = euclidean_distance_matrix(comp_centers, comp_centers)
+
+            batch_id = (
+                np.zeros((comp_geo_attribs.shape[0], 1), dtype=np.float32) * batch_ind
+            )
+            comp_geo_attribs[:, -2] = np.clip(comp_geo_attribs[:, -2], -1, 1)
+            angle = np.arccos(comp_geo_attribs[:, -2]) * np.sign(
+                comp_geo_attribs[:, -1]
+            )
+            angle = angle.reshape((-1, 1))
+            rotated_rois = np.hstack([batch_id, comp_geo_attribs[:, :-2], angle])
+            rois = paddle.to_tensor(rotated_rois)
+            content_feats = self.pooling(feat_maps[batch_ind].unsqueeze(0), rois)
+
+            content_feats = content_feats.reshape([content_feats.shape[0], -1])
+            geo_feats = feature_embedding(comp_geo_attribs, self.node_geo_feat_dim)
+            geo_feats = paddle.to_tensor(geo_feats)
+            node_feats = paddle.concat([content_feats, geo_feats], axis=-1)
+
+            sorted_dist_inds = np.argsort(distance_matrix, axis=1)
+            pivot_local_graphs, pivot_knns = self.generate_local_graphs(
+                sorted_dist_inds, node_labels
+            )
+
+            node_feat_batch.append(node_feats)
+            node_label_batch.append(node_labels)
+            local_graph_batch.append(pivot_local_graphs)
+            knn_batch.append(pivot_knns)
+            sorted_dist_inds_batch.append(sorted_dist_inds)
+
+        (node_feats, adjacent_matrices, knn_inds, gt_linkage) = self.generate_gcn_input(
+            node_feat_batch,
+            node_label_batch,
+            local_graph_batch,
+            knn_batch,
+            sorted_dist_inds_batch,
+        )
+
+        return node_feats, adjacent_matrices, knn_inds, gt_linkage

ppocr/modeling/heads/proposal_local_graph.py

@@ -0,0 +1,476 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/modules/proposal_local_graph.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import cv2
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from lanms import merge_quadrangle_n9 as la_nms
+
+from ppocr.ext_op import RoIAlignRotated
+from .local_graph import (
+    euclidean_distance_matrix,
+    feature_embedding,
+    normalize_adjacent_matrix,
+)
+
+
+def fill_hole(input_mask):
+    h, w = input_mask.shape
+    canvas = np.zeros((h + 2, w + 2), np.uint8)
+    canvas[1 : h + 1, 1 : w + 1] = input_mask.copy()
+
+    mask = np.zeros((h + 4, w + 4), np.uint8)
+
+    cv2.floodFill(canvas, mask, (0, 0), 1)
+    canvas = canvas[1 : h + 1, 1 : w + 1].astype(np.bool_)
+
+    return ~canvas | input_mask
+
+
+class ProposalLocalGraphs:
+    def __init__(
+        self,
+        k_at_hops,
+        num_adjacent_linkages,
+        node_geo_feat_len,
+        pooling_scale,
+        pooling_output_size,
+        nms_thr,
+        min_width,
+        max_width,
+        comp_shrink_ratio,
+        comp_w_h_ratio,
+        comp_score_thr,
+        text_region_thr,
+        center_region_thr,
+        center_region_area_thr,
+    ):
+        assert len(k_at_hops) == 2
+        assert isinstance(k_at_hops, tuple)
+        assert isinstance(num_adjacent_linkages, int)
+        assert isinstance(node_geo_feat_len, int)
+        assert isinstance(pooling_scale, float)
+        assert isinstance(pooling_output_size, tuple)
+        assert isinstance(nms_thr, float)
+        assert isinstance(min_width, float)
+        assert isinstance(max_width, float)
+        assert isinstance(comp_shrink_ratio, float)
+        assert isinstance(comp_w_h_ratio, float)
+        assert isinstance(comp_score_thr, float)
+        assert isinstance(text_region_thr, float)
+        assert isinstance(center_region_thr, float)
+        assert isinstance(center_region_area_thr, int)
+
+        self.k_at_hops = k_at_hops
+        self.active_connection = num_adjacent_linkages
+        self.local_graph_depth = len(self.k_at_hops)
+        self.node_geo_feat_dim = node_geo_feat_len
+        self.pooling = RoIAlignRotated(pooling_output_size, pooling_scale)
+        self.nms_thr = nms_thr
+        self.min_width = min_width
+        self.max_width = max_width
+        self.comp_shrink_ratio = comp_shrink_ratio
+        self.comp_w_h_ratio = comp_w_h_ratio
+        self.comp_score_thr = comp_score_thr
+        self.text_region_thr = text_region_thr
+        self.center_region_thr = center_region_thr
+        self.center_region_area_thr = center_region_area_thr
+
+    def propose_comps(
+        self,
+        score_map,
+        top_height_map,
+        bot_height_map,
+        sin_map,
+        cos_map,
+        comp_score_thr,
+        min_width,
+        max_width,
+        comp_shrink_ratio,
+        comp_w_h_ratio,
+    ):
+        """Propose text components.
+
+        Args:
+            score_map (ndarray): The score map for NMS.
+            top_height_map (ndarray): The predicted text height map from each
+                pixel in text center region to top sideline.
+            bot_height_map (ndarray): The predicted text height map from each
+                pixel in text center region to bottom sideline.
+            sin_map (ndarray): The predicted sin(theta) map.
+            cos_map (ndarray): The predicted cos(theta) map.
+            comp_score_thr (float): The score threshold of text component.
+            min_width (float): The minimum width of text components.
+            max_width (float): The maximum width of text components.
+            comp_shrink_ratio (float): The shrink ratio of text components.
+            comp_w_h_ratio (float): The width to height ratio of text
+                components.
+
+        Returns:
+            text_comps (ndarray): The text components.
+        """
+
+        comp_centers = np.argwhere(score_map > comp_score_thr)
+        comp_centers = comp_centers[np.argsort(comp_centers[:, 0])]
+        y = comp_centers[:, 0]
+        x = comp_centers[:, 1]
+
+        top_height = top_height_map[y, x].reshape((-1, 1)) * comp_shrink_ratio
+        bot_height = bot_height_map[y, x].reshape((-1, 1)) * comp_shrink_ratio
+        sin = sin_map[y, x].reshape((-1, 1))
+        cos = cos_map[y, x].reshape((-1, 1))
+
+        top_mid_pts = comp_centers + np.hstack([top_height * sin, top_height * cos])
+        bot_mid_pts = comp_centers - np.hstack([bot_height * sin, bot_height * cos])
+
+        width = (top_height + bot_height) * comp_w_h_ratio
+        width = np.clip(width, min_width, max_width)
+        r = width / 2
+
+        tl = top_mid_pts[:, ::-1] - np.hstack([-r * sin, r * cos])
+        tr = top_mid_pts[:, ::-1] + np.hstack([-r * sin, r * cos])
+        br = bot_mid_pts[:, ::-1] + np.hstack([-r * sin, r * cos])
+        bl = bot_mid_pts[:, ::-1] - np.hstack([-r * sin, r * cos])
+        text_comps = np.hstack([tl, tr, br, bl]).astype(np.float32)
+
+        score = score_map[y, x].reshape((-1, 1))
+        text_comps = np.hstack([text_comps, score])
+
+        return text_comps
+
+    def propose_comps_and_attribs(
+        self,
+        text_region_map,
+        center_region_map,
+        top_height_map,
+        bot_height_map,
+        sin_map,
+        cos_map,
+    ):
+        """Generate text components and attributes.
+
+        Args:
+            text_region_map (ndarray): The predicted text region probability
+                map.
+            center_region_map (ndarray): The predicted text center region
+                probability map.
+            top_height_map (ndarray): The predicted text height map from each
+                pixel in text center region to top sideline.
+            bot_height_map (ndarray): The predicted text height map from each
+                pixel in text center region to bottom sideline.
+            sin_map (ndarray): The predicted sin(theta) map.
+            cos_map (ndarray): The predicted cos(theta) map.
+
+        Returns:
+            comp_attribs (ndarray): The text component attributes.
+            text_comps (ndarray): The text components.
+        """
+
+        assert (
+            text_region_map.shape
+            == center_region_map.shape
+            == top_height_map.shape
+            == bot_height_map.shape
+            == sin_map.shape
+            == cos_map.shape
+        )
+        text_mask = text_region_map > self.text_region_thr
+        center_region_mask = (center_region_map > self.center_region_thr) * text_mask
+
+        scale = np.sqrt(1.0 / (sin_map**2 + cos_map**2 + 1e-8))
+        sin_map, cos_map = sin_map * scale, cos_map * scale
+
+        center_region_mask = fill_hole(center_region_mask)
+        center_region_contours, _ = cv2.findContours(
+            center_region_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+        )
+
+        mask_sz = center_region_map.shape
+        comp_list = []
+        for contour in center_region_contours:
+            current_center_mask = np.zeros(mask_sz)
+            cv2.drawContours(current_center_mask, [contour], -1, 1, -1)
+            if current_center_mask.sum() <= self.center_region_area_thr:
+                continue
+            score_map = text_region_map * current_center_mask
+
+            text_comps = self.propose_comps(
+                score_map,
+                top_height_map,
+                bot_height_map,
+                sin_map,
+                cos_map,
+                self.comp_score_thr,
+                self.min_width,
+                self.max_width,
+                self.comp_shrink_ratio,
+                self.comp_w_h_ratio,
+            )
+
+            text_comps = la_nms(text_comps, self.nms_thr)
+            text_comp_mask = np.zeros(mask_sz)
+            text_comp_boxes = text_comps[:, :8].reshape((-1, 4, 2)).astype(np.int32)
+
+            cv2.drawContours(text_comp_mask, text_comp_boxes, -1, 1, -1)
+            if (text_comp_mask * text_mask).sum() < text_comp_mask.sum() * 0.5:
+                continue
+            if text_comps.shape[-1] > 0:
+                comp_list.append(text_comps)
+
+        if len(comp_list) <= 0:
+            return None, None
+
+        text_comps = np.vstack(comp_list)
+        text_comp_boxes = text_comps[:, :8].reshape((-1, 4, 2))
+        centers = np.mean(text_comp_boxes, axis=1).astype(np.int32)
+        x = centers[:, 0]
+        y = centers[:, 1]
+
+        scores = []
+        for text_comp_box in text_comp_boxes:
+            text_comp_box[:, 0] = np.clip(text_comp_box[:, 0], 0, mask_sz[1] - 1)
+            text_comp_box[:, 1] = np.clip(text_comp_box[:, 1], 0, mask_sz[0] - 1)
+            min_coord = np.min(text_comp_box, axis=0).astype(np.int32)
+            max_coord = np.max(text_comp_box, axis=0).astype(np.int32)
+            text_comp_box = text_comp_box - min_coord
+            box_sz = max_coord - min_coord + 1
+            temp_comp_mask = np.zeros((box_sz[1], box_sz[0]), dtype=np.uint8)
+            cv2.fillPoly(temp_comp_mask, [text_comp_box.astype(np.int32)], 1)
+            temp_region_patch = text_region_map[
+                min_coord[1] : (max_coord[1] + 1), min_coord[0] : (max_coord[0] + 1)
+            ]
+            score = cv2.mean(temp_region_patch, temp_comp_mask)[0]
+            scores.append(score)
+        scores = np.array(scores).reshape((-1, 1))
+        text_comps = np.hstack([text_comps[:, :-1], scores])
+
+        h = top_height_map[y, x].reshape((-1, 1)) + bot_height_map[y, x].reshape(
+            (-1, 1)
+        )
+        w = np.clip(h * self.comp_w_h_ratio, self.min_width, self.max_width)
+        sin = sin_map[y, x].reshape((-1, 1))
+        cos = cos_map[y, x].reshape((-1, 1))
+
+        x = x.reshape((-1, 1))
+        y = y.reshape((-1, 1))
+        comp_attribs = np.hstack([x, y, h, w, cos, sin])
+
+        return comp_attribs, text_comps
+
+    def generate_local_graphs(self, sorted_dist_inds, node_feats):
+        """Generate local graphs and graph convolution network input data.
+
+        Args:
+            sorted_dist_inds (ndarray): The node indices sorted according to
+                the Euclidean distance.
+            node_feats (tensor): The features of nodes in graph.
+
+        Returns:
+            local_graphs_node_feats (tensor): The features of nodes in local
+                graphs.
+            adjacent_matrices (tensor): The adjacent matrices.
+            pivots_knn_inds (tensor): The k-nearest neighbor indices in
+                local graphs.
+            pivots_local_graphs (tensor): The indices of nodes in local
+                graphs.
+        """
+
+        assert sorted_dist_inds.ndim == 2
+        assert (
+            sorted_dist_inds.shape[0]
+            == sorted_dist_inds.shape[1]
+            == node_feats.shape[0]
+        )
+
+        knn_graph = sorted_dist_inds[:, 1 : self.k_at_hops[0] + 1]
+        pivot_local_graphs = []
+        pivot_knns = []
+
+        for pivot_ind, knn in enumerate(knn_graph):
+            local_graph_neighbors = set(knn)
+
+            for neighbor_ind in knn:
+                local_graph_neighbors.update(
+                    set(sorted_dist_inds[neighbor_ind, 1 : self.k_at_hops[1] + 1])
+                )
+
+            local_graph_neighbors.discard(pivot_ind)
+            pivot_local_graph = list(local_graph_neighbors)
+            pivot_local_graph.insert(0, pivot_ind)
+            pivot_knn = [pivot_ind] + list(knn)
+
+            pivot_local_graphs.append(pivot_local_graph)
+            pivot_knns.append(pivot_knn)
+
+        num_max_nodes = max(
+            [len(pivot_local_graph) for pivot_local_graph in pivot_local_graphs]
+        )
+
+        local_graphs_node_feat = []
+        adjacent_matrices = []
+        pivots_knn_inds = []
+        pivots_local_graphs = []
+
+        for graph_ind, pivot_knn in enumerate(pivot_knns):
+            pivot_local_graph = pivot_local_graphs[graph_ind]
+            num_nodes = len(pivot_local_graph)
+            pivot_ind = pivot_local_graph[0]
+            node2ind_map = {j: i for i, j in enumerate(pivot_local_graph)}
+
+            knn_inds = paddle.cast(
+                paddle.to_tensor([node2ind_map[i] for i in pivot_knn[1:]]), "int64"
+            )
+            pivot_feats = node_feats[pivot_ind]
+            normalized_feats = (
+                node_feats[paddle.to_tensor(pivot_local_graph)] - pivot_feats
+            )
+
+            adjacent_matrix = np.zeros((num_nodes, num_nodes), dtype=np.float32)
+            for node in pivot_local_graph:
+                neighbors = sorted_dist_inds[node, 1 : self.active_connection + 1]
+                for neighbor in neighbors:
+                    if neighbor in pivot_local_graph:
+                        adjacent_matrix[node2ind_map[node], node2ind_map[neighbor]] = 1
+                        adjacent_matrix[node2ind_map[neighbor], node2ind_map[node]] = 1
+
+            adjacent_matrix = normalize_adjacent_matrix(adjacent_matrix)
+            pad_adjacent_matrix = paddle.zeros(
+                (num_max_nodes, num_max_nodes),
+            )
+            pad_adjacent_matrix[:num_nodes, :num_nodes] = paddle.cast(
+                paddle.to_tensor(adjacent_matrix), "float32"
+            )
+
+            pad_normalized_feats = paddle.concat(
+                [
+                    normalized_feats,
+                    paddle.zeros(
+                        (num_max_nodes - num_nodes, normalized_feats.shape[1]),
+                    ),
+                ],
+                axis=0,
+            )
+
+            local_graph_nodes = paddle.to_tensor(pivot_local_graph)
+            local_graph_nodes = paddle.concat(
+                [
+                    local_graph_nodes,
+                    paddle.zeros([num_max_nodes - num_nodes], dtype="int64"),
+                ],
+                axis=-1,
+            )
+
+            local_graphs_node_feat.append(pad_normalized_feats)
+            adjacent_matrices.append(pad_adjacent_matrix)
+            pivots_knn_inds.append(knn_inds)
+            pivots_local_graphs.append(local_graph_nodes)
+
+        local_graphs_node_feat = paddle.stack(local_graphs_node_feat, 0)
+        adjacent_matrices = paddle.stack(adjacent_matrices, 0)
+        pivots_knn_inds = paddle.stack(pivots_knn_inds, 0)
+        pivots_local_graphs = paddle.stack(pivots_local_graphs, 0)
+
+        return (
+            local_graphs_node_feat,
+            adjacent_matrices,
+            pivots_knn_inds,
+            pivots_local_graphs,
+        )
+
+    def __call__(self, preds, feat_maps):
+        """Generate local graphs and graph convolutional network input data.
+
+        Args:
+            preds (tensor): The predicted maps.
+            feat_maps (tensor): The feature maps to extract content feature of
+                text components.
+
+        Returns:
+            none_flag (bool): The flag showing whether the number of proposed
+                text components is 0.
+            local_graphs_node_feats (tensor): The features of nodes in local
+                graphs.
+            adjacent_matrices (tensor): The adjacent matrices.
+            pivots_knn_inds (tensor): The k-nearest neighbor indices in
+                local graphs.
+            pivots_local_graphs (tensor): The indices of nodes in local
+                graphs.
+            text_comps (ndarray): The predicted text components.
+        """
+        if preds.ndim == 4:
+            assert preds.shape[0] == 1
+            preds = paddle.squeeze(preds)
+        pred_text_region = F.sigmoid(preds[0]).numpy()
+        pred_center_region = F.sigmoid(preds[1]).numpy()
+        pred_sin_map = preds[2].numpy()
+        pred_cos_map = preds[3].numpy()
+        pred_top_height_map = preds[4].numpy()
+        pred_bot_height_map = preds[5].numpy()
+
+        comp_attribs, text_comps = self.propose_comps_and_attribs(
+            pred_text_region,
+            pred_center_region,
+            pred_top_height_map,
+            pred_bot_height_map,
+            pred_sin_map,
+            pred_cos_map,
+        )
+
+        if comp_attribs is None or len(comp_attribs) < 2:
+            none_flag = True
+            return none_flag, (0, 0, 0, 0, 0)
+
+        comp_centers = comp_attribs[:, 0:2]
+        distance_matrix = euclidean_distance_matrix(comp_centers, comp_centers)
+
+        geo_feats = feature_embedding(comp_attribs, self.node_geo_feat_dim)
+        geo_feats = paddle.to_tensor(geo_feats)
+
+        batch_id = np.zeros((comp_attribs.shape[0], 1), dtype=np.float32)
+        comp_attribs = comp_attribs.astype(np.float32)
+        angle = np.arccos(comp_attribs[:, -2]) * np.sign(comp_attribs[:, -1])
+        angle = angle.reshape((-1, 1))
+        rotated_rois = np.hstack([batch_id, comp_attribs[:, :-2], angle])
+        rois = paddle.to_tensor(rotated_rois)
+
+        content_feats = self.pooling(feat_maps, rois)
+        content_feats = content_feats.reshape([content_feats.shape[0], -1])
+        node_feats = paddle.concat([content_feats, geo_feats], axis=-1)
+
+        sorted_dist_inds = np.argsort(distance_matrix, axis=1)
+        (
+            local_graphs_node_feat,
+            adjacent_matrices,
+            pivots_knn_inds,
+            pivots_local_graphs,
+        ) = self.generate_local_graphs(sorted_dist_inds, node_feats)
+
+        none_flag = False
+        return none_flag, (
+            local_graphs_node_feat,
+            adjacent_matrices,
+            pivots_knn_inds,
+            pivots_local_graphs,
+            text_comps,
+        )

ppocr/modeling/heads/rec_abinet_head.py

@@ -0,0 +1,299 @@
+# copyright (c) 2021 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/FangShancheng/ABINet/tree/main/modules
+"""
+
+import math
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle.nn import LayerList
+from ppocr.modeling.heads.rec_nrtr_head import TransformerBlock, PositionalEncoding
+
+
+class BCNLanguage(nn.Layer):
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        num_layers=4,
+        dim_feedforward=2048,
+        dropout=0.0,
+        max_length=25,
+        detach=True,
+        num_classes=37,
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.detach = detach
+        self.max_length = max_length + 1  # additional stop token
+        self.proj = nn.Linear(num_classes, d_model, bias_attr=False)
+        self.token_encoder = PositionalEncoding(
+            dropout=0.1, dim=d_model, max_len=self.max_length
+        )
+        self.pos_encoder = PositionalEncoding(
+            dropout=0, dim=d_model, max_len=self.max_length
+        )
+
+        self.decoder = nn.LayerList(
+            [
+                TransformerBlock(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    attention_dropout_rate=dropout,
+                    residual_dropout_rate=dropout,
+                    with_self_attn=False,
+                    with_cross_attn=True,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.cls = nn.Linear(d_model, num_classes)
+
+    def forward(self, tokens, lengths):
+        """
+        Args:
+            tokens: (B, N, C) where N is length, B is batch size and C is classes number
+            lengths: (B,)
+        """
+        if self.detach:
+            tokens = tokens.detach()
+        embed = self.proj(tokens)  # (B, N, C)
+        embed = self.token_encoder(embed)  # (B, N, C)
+        padding_mask = _get_mask(lengths, self.max_length)
+        zeros = paddle.zeros_like(embed)  # (B, N, C)
+        query = self.pos_encoder(zeros)
+        for decoder_layer in self.decoder:
+            query = decoder_layer(query, embed, cross_mask=padding_mask)
+        output = query  # (B, N, C)
+
+        logits = self.cls(output)  # (B, N, C)
+
+        return output, logits
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(
+        nn.Conv2D(in_c, out_c, k, s, p), nn.BatchNorm2D(out_c), nn.ReLU()
+    )
+
+
+def decoder_layer(
+    in_c, out_c, k=3, s=1, p=1, mode="nearest", scale_factor=None, size=None
+):
+    align_corners = False if mode == "nearest" else True
+    return nn.Sequential(
+        nn.Upsample(
+            size=size, scale_factor=scale_factor, mode=mode, align_corners=align_corners
+        ),
+        nn.Conv2D(in_c, out_c, k, s, p),
+        nn.BatchNorm2D(out_c),
+        nn.ReLU(),
+    )
+
+
+class PositionAttention(nn.Layer):
+    def __init__(
+        self,
+        max_length,
+        in_channels=512,
+        num_channels=64,
+        h=8,
+        w=32,
+        mode="nearest",
+        **kwargs,
+    ):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+        )
+        self.k_decoder = nn.Sequential(
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode),
+        )
+
+        self.pos_encoder = PositionalEncoding(
+            dropout=0, dim=in_channels, max_len=max_length
+        )
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        k, v = x, x
+
+        # calculate key vector
+        features = []
+        for i in range(0, len(self.k_encoder)):
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k)
+            # print(k.shape, features[len(self.k_decoder) - 2 - i].shape)
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k)
+
+        # calculate query vector
+        # TODO q=f(q,k)
+        zeros = paddle.zeros((B, self.max_length, C), dtype=x.dtype)  # (B, N, C)
+        q = self.pos_encoder(zeros)  # (B, N, C)
+        q = self.project(q)  # (B, N, C)
+
+        # calculate attention
+        attn_scores = q @ k.flatten(2)  # (B, N, (H*W))
+        attn_scores = attn_scores / (C**0.5)
+        attn_scores = F.softmax(attn_scores, axis=-1)
+
+        v = v.flatten(2).transpose([0, 2, 1])  # (B, (H*W), C)
+        attn_vecs = attn_scores @ v  # (B, N, C)
+
+        return attn_vecs, attn_scores.reshape([0, self.max_length, H, W])
+
+
+class ABINetHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        d_model=512,
+        nhead=8,
+        num_layers=3,
+        dim_feedforward=2048,
+        dropout=0.1,
+        max_length=25,
+        use_lang=False,
+        iter_size=1,
+        image_size=(32, 128),
+    ):
+        super().__init__()
+        self.max_length = max_length + 1
+        h, w = image_size[0] // 4, image_size[1] // 4
+        self.pos_encoder = PositionalEncoding(dropout=0.1, dim=d_model, max_len=h * w)
+        self.encoder = nn.LayerList(
+            [
+                TransformerBlock(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dim_feedforward=dim_feedforward,
+                    attention_dropout_rate=dropout,
+                    residual_dropout_rate=dropout,
+                    with_self_attn=True,
+                    with_cross_attn=False,
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.decoder = PositionAttention(
+            max_length=max_length + 1, mode="nearest", h=h, w=w  # additional stop token
+        )
+        self.out_channels = out_channels
+        self.cls = nn.Linear(d_model, self.out_channels)
+        self.use_lang = use_lang
+        if use_lang:
+            self.iter_size = iter_size
+            self.language = BCNLanguage(
+                d_model=d_model,
+                nhead=nhead,
+                num_layers=4,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                max_length=max_length,
+                num_classes=self.out_channels,
+            )
+            # alignment
+            self.w_att_align = nn.Linear(2 * d_model, d_model)
+            self.cls_align = nn.Linear(d_model, self.out_channels)
+
+    def forward(self, x, targets=None):
+        x = x.transpose([0, 2, 3, 1])
+        _, H, W, C = x.shape
+        feature = x.flatten(1, 2)
+        feature = self.pos_encoder(feature)
+        for encoder_layer in self.encoder:
+            feature = encoder_layer(feature)
+        feature = feature.reshape([0, H, W, C]).transpose([0, 3, 1, 2])
+        v_feature, attn_scores = self.decoder(feature)  # (B, N, C), (B, C, H, W)
+        vis_logits = self.cls(v_feature)  # (B, N, C)
+        logits = vis_logits
+        vis_lengths = _get_length(vis_logits)
+        if self.use_lang:
+            align_logits = vis_logits
+            align_lengths = vis_lengths
+            all_l_res, all_a_res = [], []
+            for i in range(self.iter_size):
+                tokens = F.softmax(align_logits, axis=-1)
+                lengths = align_lengths
+                lengths = paddle.clip(
+                    lengths, 2, self.max_length
+                )  # TODO:move to language model
+                l_feature, l_logits = self.language(tokens, lengths)
+
+                # alignment
+                all_l_res.append(l_logits)
+                fuse = paddle.concat((l_feature, v_feature), -1)
+                f_att = F.sigmoid(self.w_att_align(fuse))
+                output = f_att * v_feature + (1 - f_att) * l_feature
+                align_logits = self.cls_align(output)  # (B, N, C)
+
+                align_lengths = _get_length(align_logits)
+                all_a_res.append(align_logits)
+            if self.training:
+                return {"align": all_a_res, "lang": all_l_res, "vision": vis_logits}
+            else:
+                logits = align_logits
+        if self.training:
+            return logits
+        else:
+            return F.softmax(logits, -1)
+
+
+def _get_length(logit):
+    """Greed decoder to obtain length from logit"""
+    out = logit.argmax(-1) == 0
+    abn = out.any(-1)
+    out_int = out.cast("int32")
+    out = (out_int.cumsum(-1) == 1) & out
+    out = out.cast("int32")
+    out = out.argmax(-1)
+    out = out + 1
+    len_seq = paddle.zeros_like(out) + logit.shape[1]
+    out = paddle.where(abn, out, len_seq)
+    return out
+
+
+def _get_mask(length, max_length):
+    """Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+    Unmasked positions are filled with float(0.0).
+    """
+    length = length.unsqueeze(-1)
+    B = length.shape[0]
+    grid = paddle.arange(0, max_length).unsqueeze(0).tile([B, 1])
+    zero_mask = paddle.zeros([B, max_length], dtype="float32")
+    inf_mask = paddle.full([B, max_length], "-inf", dtype="float32")
+    diag_mask = paddle.diag(
+        paddle.full([max_length], "-inf", dtype=paddle.float32), offset=0, name=None
+    )
+    mask = paddle.where(grid >= length, inf_mask, zero_mask)
+    mask = mask.unsqueeze(1) + diag_mask
+    return mask.unsqueeze(1)

ppocr/modeling/heads/rec_aster_head.py

@@ -0,0 +1,404 @@
+# 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/attention_recognition_head.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+
+import paddle
+from paddle import nn
+from paddle.nn import functional as F
+
+
+class AsterHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        sDim,
+        attDim,
+        max_len_labels,
+        time_step=25,
+        beam_width=5,
+        **kwargs,
+    ):
+        super(AsterHead, self).__init__()
+        self.num_classes = out_channels
+        self.in_planes = in_channels
+        self.sDim = sDim
+        self.attDim = attDim
+        self.max_len_labels = max_len_labels
+        self.decoder = AttentionRecognitionHead(
+            in_channels, out_channels, sDim, attDim, max_len_labels
+        )
+        self.time_step = time_step
+        self.embedder = Embedding(self.time_step, in_channels)
+        self.beam_width = beam_width
+        self.eos = self.num_classes - 3
+
+    def forward(self, x, targets=None, embed=None):
+        return_dict = {}
+        embedding_vectors = self.embedder(x)
+
+        if self.training:
+            rec_targets, rec_lengths, _ = targets
+            rec_pred = self.decoder([x, rec_targets, rec_lengths], embedding_vectors)
+            return_dict["rec_pred"] = rec_pred
+            return_dict["embedding_vectors"] = embedding_vectors
+        else:
+            rec_pred, rec_pred_scores = self.decoder.beam_search(
+                x, self.beam_width, self.eos, embedding_vectors
+            )
+            return_dict["rec_pred"] = rec_pred
+            return_dict["rec_pred_scores"] = rec_pred_scores
+            return_dict["embedding_vectors"] = embedding_vectors
+
+        return return_dict
+
+
+class Embedding(nn.Layer):
+    def __init__(self, in_timestep, in_planes, mid_dim=4096, embed_dim=300):
+        super(Embedding, self).__init__()
+        self.in_timestep = in_timestep
+        self.in_planes = in_planes
+        self.embed_dim = embed_dim
+        self.mid_dim = mid_dim
+        self.eEmbed = nn.Linear(
+            in_timestep * in_planes, self.embed_dim
+        )  # Embed encoder output to a word-embedding like
+
+    def forward(self, x):
+        x = paddle.reshape(x, [x.shape[0], -1])
+        x = self.eEmbed(x)
+        return x
+
+
+class AttentionRecognitionHead(nn.Layer):
+    """
+    input: [b x 16 x 64 x in_planes]
+    output: probability sequence: [b x T x num_classes]
+    """
+
+    def __init__(self, in_channels, out_channels, sDim, attDim, max_len_labels):
+        super(AttentionRecognitionHead, self).__init__()
+        self.num_classes = (
+            out_channels  # this is the output classes. So it includes the <EOS>.
+        )
+        self.in_planes = in_channels
+        self.sDim = sDim
+        self.attDim = attDim
+        self.max_len_labels = max_len_labels
+
+        self.decoder = DecoderUnit(
+            sDim=sDim, xDim=in_channels, yDim=self.num_classes, attDim=attDim
+        )
+
+    def forward(self, x, embed):
+        x, targets, lengths = x
+        batch_size = x.shape[0]
+        # Decoder
+        state = self.decoder.get_initial_state(embed)
+        outputs = []
+        for i in range(max(lengths)):
+            if i == 0:
+                y_prev = paddle.full(shape=[batch_size], fill_value=self.num_classes)
+            else:
+                y_prev = targets[:, i - 1]
+            output, state = self.decoder(x, state, y_prev)
+            outputs.append(output)
+        outputs = paddle.concat([_.unsqueeze(1) for _ in outputs], 1)
+        return outputs
+
+    # inference stage.
+    def sample(self, x):
+        x, _, _ = x
+        batch_size = x.size(0)
+        # Decoder
+        state = paddle.zeros([1, batch_size, self.sDim])
+
+        predicted_ids, predicted_scores, predicted = [], [], None
+        for i in range(self.max_len_labels):
+            if i == 0:
+                y_prev = paddle.full(shape=[batch_size], fill_value=self.num_classes)
+            else:
+                y_prev = predicted
+
+            output, state = self.decoder(x, state, y_prev)
+            output = F.softmax(output, axis=1)
+            score, predicted = output.max(1)
+            predicted_ids.append(predicted.unsqueeze(1))
+            predicted_scores.append(score.unsqueeze(1))
+        predicted_ids = paddle.concat([predicted_ids, 1])
+        predicted_scores = paddle.concat([predicted_scores, 1])
+        # return predicted_ids.squeeze(), predicted_scores.squeeze()
+        return predicted_ids, predicted_scores
+
+    def beam_search(self, x, beam_width, eos, embed):
+        def _inflate(tensor, times, dim):
+            repeat_dims = [1] * tensor.dim()
+            repeat_dims[dim] = times
+            output = paddle.tile(tensor, repeat_dims)
+            return output
+
+        # https://github.com/IBM/pytorch-seq2seq/blob/fede87655ddce6c94b38886089e05321dc9802af/seq2seq/models/TopKDecoder.py
+        batch_size, l, d = x.shape
+        x = paddle.tile(
+            paddle.transpose(x.unsqueeze(1), perm=[1, 0, 2, 3]), [beam_width, 1, 1, 1]
+        )
+        inflated_encoder_feats = paddle.reshape(
+            paddle.transpose(x, perm=[1, 0, 2, 3]), [-1, l, d]
+        )
+
+        # Initialize the decoder
+        state = self.decoder.get_initial_state(embed, tile_times=beam_width)
+
+        pos_index = paddle.reshape(
+            paddle.arange(batch_size) * beam_width, shape=[-1, 1]
+        )
+
+        # Initialize the scores
+        sequence_scores = paddle.full(
+            shape=[batch_size * beam_width, 1], fill_value=-float("Inf")
+        )
+        index = [i * beam_width for i in range(0, batch_size)]
+        sequence_scores[index] = 0.0
+
+        # Initialize the input vector
+        y_prev = paddle.full(
+            shape=[batch_size * beam_width], fill_value=self.num_classes
+        )
+
+        # Store decisions for backtracking
+        stored_scores = list()
+        stored_predecessors = list()
+        stored_emitted_symbols = list()
+
+        for i in range(self.max_len_labels):
+            output, state = self.decoder(inflated_encoder_feats, state, y_prev)
+            state = paddle.unsqueeze(state, axis=0)
+            log_softmax_output = paddle.nn.functional.log_softmax(output, axis=1)
+
+            sequence_scores = _inflate(sequence_scores, self.num_classes, 1)
+            sequence_scores += log_softmax_output
+            scores, candidates = paddle.topk(
+                paddle.reshape(sequence_scores, [batch_size, -1]), beam_width, axis=1
+            )
+
+            # Reshape input = (bk, 1) and sequence_scores = (bk, 1)
+            y_prev = paddle.reshape(
+                candidates % self.num_classes, shape=[batch_size * beam_width]
+            )
+            sequence_scores = paddle.reshape(scores, shape=[batch_size * beam_width, 1])
+
+            # Update fields for next timestep
+            pos_index = paddle.expand_as(pos_index, candidates)
+            predecessors = paddle.cast(
+                candidates / self.num_classes + pos_index, dtype="int64"
+            )
+            predecessors = paddle.reshape(
+                predecessors, shape=[batch_size * beam_width, 1]
+            )
+            state = paddle.index_select(state, index=predecessors.squeeze(), axis=1)
+
+            # Update sequence scores and erase scores for <eos> symbol so that they aren't expanded
+            stored_scores.append(sequence_scores.clone())
+            y_prev = paddle.reshape(y_prev, shape=[-1, 1])
+            eos_prev = paddle.full_like(y_prev, fill_value=eos)
+            mask = eos_prev == y_prev
+            mask = paddle.nonzero(mask)
+            if mask.dim() > 0:
+                sequence_scores = sequence_scores.numpy()
+                mask = mask.numpy()
+                sequence_scores[mask] = -float("inf")
+                sequence_scores = paddle.to_tensor(sequence_scores)
+
+            # Cache results for backtracking
+            stored_predecessors.append(predecessors)
+            y_prev = paddle.squeeze(y_prev)
+            stored_emitted_symbols.append(y_prev)
+
+        # Do backtracking to return the optimal values
+        # ====== backtrak ======#
+        # Initialize return variables given different types
+        p = list()
+        l = [
+            [self.max_len_labels] * beam_width for _ in range(batch_size)
+        ]  # Placeholder for lengths of top-k sequences
+
+        # the last step output of the beams are not sorted
+        # thus they are sorted here
+        sorted_score, sorted_idx = paddle.topk(
+            paddle.reshape(stored_scores[-1], shape=[batch_size, beam_width]),
+            beam_width,
+        )
+
+        # initialize the sequence scores with the sorted last step beam scores
+        s = sorted_score.clone()
+
+        batch_eos_found = [0] * batch_size  # the number of EOS found
+        # in the backward loop below for each batch
+        t = self.max_len_labels - 1
+        # initialize the back pointer with the sorted order of the last step beams.
+        # add pos_index for indexing variable with b*k as the first dimension.
+        t_predecessors = paddle.reshape(
+            sorted_idx + pos_index.expand_as(sorted_idx),
+            shape=[batch_size * beam_width],
+        )
+        while t >= 0:
+            # Re-order the variables with the back pointer
+            current_symbol = paddle.index_select(
+                stored_emitted_symbols[t], index=t_predecessors, axis=0
+            )
+            t_predecessors = paddle.index_select(
+                stored_predecessors[t].squeeze(), index=t_predecessors, axis=0
+            )
+            eos_indices = stored_emitted_symbols[t] == eos
+            eos_indices = paddle.nonzero(eos_indices)
+
+            if eos_indices.dim() > 0:
+                for i in range(eos_indices.shape[0] - 1, -1, -1):
+                    # Indices of the EOS symbol for both variables
+                    # with b*k as the first dimension, and b, k for
+                    # the first two dimensions
+                    idx = eos_indices[i]
+                    b_idx = int(idx[0] / beam_width)
+                    # The indices of the replacing position
+                    # according to the replacement strategy noted above
+                    res_k_idx = beam_width - (batch_eos_found[b_idx] % beam_width) - 1
+                    batch_eos_found[b_idx] += 1
+                    res_idx = b_idx * beam_width + res_k_idx
+
+                    # Replace the old information in return variables
+                    # with the new ended sequence information
+                    t_predecessors[res_idx] = stored_predecessors[t][idx[0]]
+                    current_symbol[res_idx] = stored_emitted_symbols[t][idx[0]]
+                    s[b_idx, res_k_idx] = stored_scores[t][idx[0], 0]
+                    l[b_idx][res_k_idx] = t + 1
+
+            # record the back tracked results
+            p.append(current_symbol)
+            t -= 1
+
+        # Sort and re-order again as the added ended sequences may change
+        # the order (very unlikely)
+        s, re_sorted_idx = s.topk(beam_width)
+        for b_idx in range(batch_size):
+            l[b_idx] = [l[b_idx][k_idx.item()] for k_idx in re_sorted_idx[b_idx, :]]
+
+        re_sorted_idx = paddle.reshape(
+            re_sorted_idx + pos_index.expand_as(re_sorted_idx),
+            [batch_size * beam_width],
+        )
+
+        # Reverse the sequences and re-order at the same time
+        # It is reversed because the backtracking happens in reverse time order
+        p = [
+            paddle.reshape(
+                paddle.index_select(step, re_sorted_idx, 0),
+                shape=[batch_size, beam_width, -1],
+            )
+            for step in reversed(p)
+        ]
+        p = paddle.concat(p, -1)[:, 0, :]
+        return p, paddle.ones_like(p)
+
+
+class AttentionUnit(nn.Layer):
+    def __init__(self, sDim, xDim, attDim):
+        super(AttentionUnit, self).__init__()
+
+        self.sDim = sDim
+        self.xDim = xDim
+        self.attDim = attDim
+
+        self.sEmbed = nn.Linear(sDim, attDim)
+        self.xEmbed = nn.Linear(xDim, attDim)
+        self.wEmbed = nn.Linear(attDim, 1)
+
+    def forward(self, x, sPrev):
+        batch_size, T, _ = x.shape  # [b x T x xDim]
+        x = paddle.reshape(x, [-1, self.xDim])  # [(b x T) x xDim]
+        xProj = self.xEmbed(x)  # [(b x T) x attDim]
+        xProj = paddle.reshape(xProj, [batch_size, T, -1])  # [b x T x attDim]
+
+        sPrev = sPrev.squeeze(0)
+        sProj = self.sEmbed(sPrev)  # [b x attDim]
+        sProj = paddle.unsqueeze(sProj, 1)  # [b x 1 x attDim]
+        sProj = paddle.expand(sProj, [batch_size, T, self.attDim])  # [b x T x attDim]
+
+        sumTanh = paddle.tanh(sProj + xProj)
+        sumTanh = paddle.reshape(sumTanh, [-1, self.attDim])
+
+        vProj = self.wEmbed(sumTanh)  # [(b x T) x 1]
+        vProj = paddle.reshape(vProj, [batch_size, T])
+        alpha = F.softmax(
+            vProj, axis=1
+        )  # attention weights for each sample in the minibatch
+        return alpha
+
+
+class DecoderUnit(nn.Layer):
+    def __init__(self, sDim, xDim, yDim, attDim):
+        super(DecoderUnit, self).__init__()
+        self.sDim = sDim
+        self.xDim = xDim
+        self.yDim = yDim
+        self.attDim = attDim
+        self.emdDim = attDim
+
+        self.attention_unit = AttentionUnit(sDim, xDim, attDim)
+        self.tgt_embedding = nn.Embedding(
+            yDim + 1, self.emdDim, weight_attr=nn.initializer.Normal(std=0.01)
+        )  # the last is used for <BOS>
+        self.gru = nn.GRUCell(input_size=xDim + self.emdDim, hidden_size=sDim)
+        self.fc = nn.Linear(
+            sDim,
+            yDim,
+            weight_attr=nn.initializer.Normal(std=0.01),
+            bias_attr=nn.initializer.Constant(value=0),
+        )
+        self.embed_fc = nn.Linear(300, self.sDim)
+
+    def get_initial_state(self, embed, tile_times=1):
+        assert embed.shape[1] == 300
+        state = self.embed_fc(embed)  # N * sDim
+        if tile_times != 1:
+            state = state.unsqueeze(1)
+            trans_state = paddle.transpose(state, perm=[1, 0, 2])
+            state = paddle.tile(trans_state, repeat_times=[tile_times, 1, 1])
+            trans_state = paddle.transpose(state, perm=[1, 0, 2])
+            state = paddle.reshape(trans_state, shape=[-1, self.sDim])
+        state = state.unsqueeze(0)  # 1 * N * sDim
+        return state
+
+    def forward(self, x, sPrev, yPrev):
+        # x: feature sequence from the image decoder.
+        batch_size, T, _ = x.shape
+        alpha = self.attention_unit(x, sPrev)
+        context = paddle.squeeze(paddle.matmul(alpha.unsqueeze(1), x), axis=1)
+        yPrev = paddle.cast(yPrev, dtype="int64")
+        yProj = self.tgt_embedding(yPrev)
+
+        concat_context = paddle.concat([yProj, context], 1)
+        concat_context = paddle.squeeze(concat_context, 1)
+        sPrev = paddle.squeeze(sPrev, 0)
+        output, state = self.gru(concat_context, sPrev)
+        output = paddle.squeeze(output, axis=1)
+        output = self.fc(output)
+        return output, state

ppocr/modeling/heads/rec_att_head.py

@@ -0,0 +1,215 @@
+# copyright (c) 2021 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 paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+
+
+class AttentionHead(nn.Layer):
+    def __init__(self, in_channels, out_channels, hidden_size, **kwargs):
+        super(AttentionHead, self).__init__()
+        self.input_size = in_channels
+        self.hidden_size = hidden_size
+        self.num_classes = out_channels
+
+        self.attention_cell = AttentionGRUCell(
+            in_channels, hidden_size, out_channels, use_gru=False
+        )
+        self.generator = nn.Linear(hidden_size, out_channels)
+
+    def _char_to_onehot(self, input_char, onehot_dim):
+        input_ont_hot = F.one_hot(input_char, onehot_dim)
+        return input_ont_hot
+
+    def forward(self, inputs, targets=None, batch_max_length=25):
+        batch_size = inputs.shape[0]
+        num_steps = batch_max_length
+
+        hidden = paddle.zeros((batch_size, self.hidden_size))
+        output_hiddens = []
+
+        if targets is not None:
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(
+                    targets[:, i], onehot_dim=self.num_classes
+                )
+                (outputs, hidden), alpha = self.attention_cell(
+                    hidden, inputs, char_onehots
+                )
+                output_hiddens.append(paddle.unsqueeze(outputs, axis=1))
+            output = paddle.concat(output_hiddens, axis=1)
+            probs = self.generator(output)
+        else:
+            targets = paddle.zeros(shape=[batch_size], dtype="int32")
+            probs = None
+            char_onehots = None
+            outputs = None
+            alpha = None
+
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(
+                    targets, onehot_dim=self.num_classes
+                )
+                (outputs, hidden), alpha = self.attention_cell(
+                    hidden, inputs, char_onehots
+                )
+                probs_step = self.generator(outputs)
+                if probs is None:
+                    probs = paddle.unsqueeze(probs_step, axis=1)
+                else:
+                    probs = paddle.concat(
+                        [probs, paddle.unsqueeze(probs_step, axis=1)], axis=1
+                    )
+                next_input = probs_step.argmax(axis=1)
+                targets = next_input
+        if not self.training:
+            probs = paddle.nn.functional.softmax(probs, axis=2)
+        return probs
+
+
+class AttentionGRUCell(nn.Layer):
+    def __init__(self, input_size, hidden_size, num_embeddings, use_gru=False):
+        super(AttentionGRUCell, self).__init__()
+        self.i2h = nn.Linear(input_size, hidden_size, bias_attr=False)
+        self.h2h = nn.Linear(hidden_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias_attr=False)
+
+        self.rnn = nn.GRUCell(
+            input_size=input_size + num_embeddings, hidden_size=hidden_size
+        )
+
+        self.hidden_size = hidden_size
+
+    def forward(self, prev_hidden, batch_H, char_onehots):
+        batch_H_proj = self.i2h(batch_H)
+        prev_hidden_proj = paddle.unsqueeze(self.h2h(prev_hidden), axis=1)
+
+        res = paddle.add(batch_H_proj, prev_hidden_proj)
+        res = paddle.tanh(res)
+        e = self.score(res)
+
+        alpha = F.softmax(e, axis=1)
+        alpha = paddle.transpose(alpha, [0, 2, 1])
+        context = paddle.squeeze(paddle.mm(alpha, batch_H), axis=1)
+        concat_context = paddle.concat([context, char_onehots], 1)
+
+        cur_hidden = self.rnn(concat_context, prev_hidden)
+
+        return cur_hidden, alpha
+
+
+class AttentionLSTM(nn.Layer):
+    def __init__(self, in_channels, out_channels, hidden_size, **kwargs):
+        super(AttentionLSTM, self).__init__()
+        self.input_size = in_channels
+        self.hidden_size = hidden_size
+        self.num_classes = out_channels
+
+        self.attention_cell = AttentionLSTMCell(
+            in_channels, hidden_size, out_channels, use_gru=False
+        )
+        self.generator = nn.Linear(hidden_size, out_channels)
+
+    def _char_to_onehot(self, input_char, onehot_dim):
+        input_ont_hot = F.one_hot(input_char, onehot_dim)
+        return input_ont_hot
+
+    def forward(self, inputs, targets=None, batch_max_length=25):
+        batch_size = inputs.shape[0]
+        num_steps = batch_max_length
+
+        hidden = (
+            paddle.zeros((batch_size, self.hidden_size)),
+            paddle.zeros((batch_size, self.hidden_size)),
+        )
+        output_hiddens = []
+
+        if targets is not None:
+            for i in range(num_steps):
+                # one-hot vectors for a i-th char
+                char_onehots = self._char_to_onehot(
+                    targets[:, i], onehot_dim=self.num_classes
+                )
+                hidden, alpha = self.attention_cell(hidden, inputs, char_onehots)
+
+                hidden = (hidden[1][0], hidden[1][1])
+                output_hiddens.append(paddle.unsqueeze(hidden[0], axis=1))
+            output = paddle.concat(output_hiddens, axis=1)
+            probs = self.generator(output)
+
+        else:
+            targets = paddle.zeros(shape=[batch_size], dtype="int32")
+            probs = None
+            char_onehots = None
+            alpha = None
+
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(
+                    targets, onehot_dim=self.num_classes
+                )
+                hidden, alpha = self.attention_cell(hidden, inputs, char_onehots)
+                probs_step = self.generator(hidden[0])
+                hidden = (hidden[1][0], hidden[1][1])
+                if probs is None:
+                    probs = paddle.unsqueeze(probs_step, axis=1)
+                else:
+                    probs = paddle.concat(
+                        [probs, paddle.unsqueeze(probs_step, axis=1)], axis=1
+                    )
+
+                next_input = probs_step.argmax(axis=1)
+
+                targets = next_input
+        if not self.training:
+            probs = paddle.nn.functional.softmax(probs, axis=2)
+        return probs
+
+
+class AttentionLSTMCell(nn.Layer):
+    def __init__(self, input_size, hidden_size, num_embeddings, use_gru=False):
+        super(AttentionLSTMCell, self).__init__()
+        self.i2h = nn.Linear(input_size, hidden_size, bias_attr=False)
+        self.h2h = nn.Linear(hidden_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias_attr=False)
+        if not use_gru:
+            self.rnn = nn.LSTMCell(
+                input_size=input_size + num_embeddings, hidden_size=hidden_size
+            )
+        else:
+            self.rnn = nn.GRUCell(
+                input_size=input_size + num_embeddings, hidden_size=hidden_size
+            )
+
+        self.hidden_size = hidden_size
+
+    def forward(self, prev_hidden, batch_H, char_onehots):
+        batch_H_proj = self.i2h(batch_H)
+        prev_hidden_proj = paddle.unsqueeze(self.h2h(prev_hidden[0]), axis=1)
+        res = paddle.add(batch_H_proj, prev_hidden_proj)
+        res = paddle.tanh(res)
+        e = self.score(res)
+
+        alpha = F.softmax(e, axis=1)
+        alpha = paddle.transpose(alpha, [0, 2, 1])
+        context = paddle.squeeze(paddle.mm(alpha, batch_H), axis=1)
+        concat_context = paddle.concat([context, char_onehots], 1)
+        cur_hidden = self.rnn(concat_context, prev_hidden)
+
+        return cur_hidden, alpha

ppocr/modeling/heads/rec_can_head.py

@@ -0,0 +1,338 @@
+# copyright (c) 2019 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/LBH1024/CAN/models/can.py
+https://github.com/LBH1024/CAN/models/counting.py
+https://github.com/LBH1024/CAN/models/decoder.py
+https://github.com/LBH1024/CAN/models/attention.py
+
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle.nn as nn
+import paddle
+import math
+
+"""
+Counting Module
+"""
+
+
+class ChannelAtt(nn.Layer):
+    def __init__(self, channel, reduction):
+        super(ChannelAtt, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2D(1)
+
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.ReLU(),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.shape
+        y = paddle.reshape(self.avg_pool(x), [b, c])
+        y = paddle.reshape(self.fc(y), [b, c, 1, 1])
+        return x * y
+
+
+class CountingDecoder(nn.Layer):
+    def __init__(self, in_channel, out_channel, kernel_size):
+        super(CountingDecoder, self).__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+        self.trans_layer = nn.Sequential(
+            nn.Conv2D(
+                self.in_channel,
+                512,
+                kernel_size=kernel_size,
+                padding=kernel_size // 2,
+                bias_attr=False,
+            ),
+            nn.BatchNorm2D(512),
+        )
+
+        self.channel_att = ChannelAtt(512, 16)
+
+        self.pred_layer = nn.Sequential(
+            nn.Conv2D(512, self.out_channel, kernel_size=1, bias_attr=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x, mask):
+        b, _, h, w = x.shape
+        x = self.trans_layer(x)
+        x = self.channel_att(x)
+        x = self.pred_layer(x)
+
+        if mask is not None:
+            x = x * mask
+        x = paddle.reshape(x, [b, self.out_channel, -1])
+        x1 = paddle.sum(x, axis=-1)
+
+        return x1, paddle.reshape(x, [b, self.out_channel, h, w])
+
+
+"""
+Attention Decoder
+"""
+
+
+class PositionEmbeddingSine(nn.Layer):
+    def __init__(
+        self, num_pos_feats=64, temperature=10000, normalize=False, scale=None
+    ):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, x, mask):
+        y_embed = paddle.cumsum(mask, 1, dtype="float32")
+        x_embed = paddle.cumsum(mask, 2, dtype="float32")
+
+        if self.normalize:
+            eps = 1e-6
+            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
+            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
+        dim_t = paddle.arange(self.num_pos_feats, dtype="float32")
+        dim_d = paddle.expand(paddle.to_tensor(2), dim_t.shape)
+        dim_t = self.temperature ** (
+            2 * (dim_t / dim_d).astype("int64") / self.num_pos_feats
+        )
+
+        pos_x = paddle.unsqueeze(x_embed, [3]) / dim_t
+        pos_y = paddle.unsqueeze(y_embed, [3]) / dim_t
+
+        pos_x = paddle.flatten(
+            paddle.stack(
+                [paddle.sin(pos_x[:, :, :, 0::2]), paddle.cos(pos_x[:, :, :, 1::2])],
+                axis=4,
+            ),
+            3,
+        )
+        pos_y = paddle.flatten(
+            paddle.stack(
+                [paddle.sin(pos_y[:, :, :, 0::2]), paddle.cos(pos_y[:, :, :, 1::2])],
+                axis=4,
+            ),
+            3,
+        )
+
+        pos = paddle.transpose(paddle.concat([pos_y, pos_x], axis=3), [0, 3, 1, 2])
+
+        return pos
+
+
+class AttDecoder(nn.Layer):
+    def __init__(
+        self,
+        ratio,
+        is_train,
+        input_size,
+        hidden_size,
+        encoder_out_channel,
+        dropout,
+        dropout_ratio,
+        word_num,
+        counting_decoder_out_channel,
+        attention,
+    ):
+        super(AttDecoder, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.out_channel = encoder_out_channel
+        self.attention_dim = attention["attention_dim"]
+        self.dropout_prob = dropout
+        self.ratio = ratio
+        self.word_num = word_num
+
+        self.counting_num = counting_decoder_out_channel
+        self.is_train = is_train
+
+        self.init_weight = nn.Linear(self.out_channel, self.hidden_size)
+        self.embedding = nn.Embedding(self.word_num, self.input_size)
+        self.word_input_gru = nn.GRUCell(self.input_size, self.hidden_size)
+        self.word_attention = Attention(hidden_size, attention["attention_dim"])
+
+        self.encoder_feature_conv = nn.Conv2D(
+            self.out_channel,
+            self.attention_dim,
+            kernel_size=attention["word_conv_kernel"],
+            padding=attention["word_conv_kernel"] // 2,
+        )
+
+        self.word_state_weight = nn.Linear(self.hidden_size, self.hidden_size)
+        self.word_embedding_weight = nn.Linear(self.input_size, self.hidden_size)
+        self.word_context_weight = nn.Linear(self.out_channel, self.hidden_size)
+        self.counting_context_weight = nn.Linear(self.counting_num, self.hidden_size)
+        self.word_convert = nn.Linear(self.hidden_size, self.word_num)
+
+        if dropout:
+            self.dropout = nn.Dropout(dropout_ratio)
+
+    def forward(self, cnn_features, labels, counting_preds, images_mask):
+        if self.is_train:
+            _, num_steps = labels.shape
+        else:
+            num_steps = 36
+
+        batch_size, _, height, width = cnn_features.shape
+        images_mask = images_mask[:, :, :: self.ratio, :: self.ratio]
+
+        word_probs = paddle.zeros((batch_size, num_steps, self.word_num))
+        word_alpha_sum = paddle.zeros((batch_size, 1, height, width))
+
+        hidden = self.init_hidden(cnn_features, images_mask)
+        counting_context_weighted = self.counting_context_weight(counting_preds)
+        cnn_features_trans = self.encoder_feature_conv(cnn_features)
+
+        position_embedding = PositionEmbeddingSine(256, normalize=True)
+        pos = position_embedding(cnn_features_trans, images_mask[:, 0, :, :])
+
+        cnn_features_trans = cnn_features_trans + pos
+
+        word = paddle.ones([batch_size, 1], dtype="int64")  # init word as sos
+        word = word.squeeze(axis=1)
+        for i in range(num_steps):
+            word_embedding = self.embedding(word)
+            _, hidden = self.word_input_gru(word_embedding, hidden)
+            word_context_vec, _, word_alpha_sum = self.word_attention(
+                cnn_features, cnn_features_trans, hidden, word_alpha_sum, images_mask
+            )
+
+            current_state = self.word_state_weight(hidden)
+            word_weighted_embedding = self.word_embedding_weight(word_embedding)
+            word_context_weighted = self.word_context_weight(word_context_vec)
+
+            if self.dropout_prob:
+                word_out_state = self.dropout(
+                    current_state
+                    + word_weighted_embedding
+                    + word_context_weighted
+                    + counting_context_weighted
+                )
+            else:
+                word_out_state = (
+                    current_state
+                    + word_weighted_embedding
+                    + word_context_weighted
+                    + counting_context_weighted
+                )
+
+            word_prob = self.word_convert(word_out_state)
+            word_probs[:, i] = word_prob
+
+            if self.is_train:
+                word = labels[:, i]
+            else:
+                word = word_prob.argmax(1)
+                word = paddle.multiply(
+                    word, labels[:, i]
+                )  # labels are oneslike tensor in infer/predict mode
+
+        return word_probs
+
+    def init_hidden(self, features, feature_mask):
+        average = paddle.sum(
+            paddle.sum(features * feature_mask, axis=-1), axis=-1
+        ) / paddle.sum((paddle.sum(feature_mask, axis=-1)), axis=-1)
+        average = self.init_weight(average)
+        return paddle.tanh(average)
+
+
+"""
+Attention Module
+"""
+
+
+class Attention(nn.Layer):
+    def __init__(self, hidden_size, attention_dim):
+        super(Attention, self).__init__()
+        self.hidden = hidden_size
+        self.attention_dim = attention_dim
+        self.hidden_weight = nn.Linear(self.hidden, self.attention_dim)
+        self.attention_conv = nn.Conv2D(
+            1, 512, kernel_size=11, padding=5, bias_attr=False
+        )
+        self.attention_weight = nn.Linear(512, self.attention_dim, bias_attr=False)
+        self.alpha_convert = nn.Linear(self.attention_dim, 1)
+
+    def forward(
+        self, cnn_features, cnn_features_trans, hidden, alpha_sum, image_mask=None
+    ):
+        query = self.hidden_weight(hidden)
+        alpha_sum_trans = self.attention_conv(alpha_sum)
+        coverage_alpha = self.attention_weight(
+            paddle.transpose(alpha_sum_trans, [0, 2, 3, 1])
+        )
+        alpha_score = paddle.tanh(
+            paddle.unsqueeze(query, [1, 2])
+            + coverage_alpha
+            + paddle.transpose(cnn_features_trans, [0, 2, 3, 1])
+        )
+        energy = self.alpha_convert(alpha_score)
+        energy = energy - energy.max()
+        energy_exp = paddle.exp(paddle.squeeze(energy, -1))
+
+        if image_mask is not None:
+            energy_exp = energy_exp * paddle.squeeze(image_mask, 1)
+        alpha = energy_exp / (
+            paddle.unsqueeze(paddle.sum(paddle.sum(energy_exp, -1), -1), [1, 2]) + 1e-10
+        )
+        alpha_sum = paddle.unsqueeze(alpha, 1) + alpha_sum
+        context_vector = paddle.sum(
+            paddle.sum((paddle.unsqueeze(alpha, 1) * cnn_features), -1), -1
+        )
+
+        return context_vector, alpha, alpha_sum
+
+
+class CANHead(nn.Layer):
+    def __init__(self, in_channel, out_channel, ratio, attdecoder, **kwargs):
+        super(CANHead, self).__init__()
+
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+
+        self.counting_decoder1 = CountingDecoder(
+            self.in_channel, self.out_channel, 3
+        )  # mscm
+        self.counting_decoder2 = CountingDecoder(self.in_channel, self.out_channel, 5)
+
+        self.decoder = AttDecoder(ratio, **attdecoder)
+
+        self.ratio = ratio
+
+    def forward(self, inputs, targets=None):
+        cnn_features, images_mask, labels = inputs
+
+        counting_mask = images_mask[:, :, :: self.ratio, :: self.ratio]
+        counting_preds1, _ = self.counting_decoder1(cnn_features, counting_mask)
+        counting_preds2, _ = self.counting_decoder2(cnn_features, counting_mask)
+        counting_preds = (counting_preds1 + counting_preds2) / 2
+
+        word_probs = self.decoder(cnn_features, labels, counting_preds, images_mask)
+        return word_probs, counting_preds, counting_preds1, counting_preds2

ppocr/modeling/heads/rec_cppd_head.py

@@ -0,0 +1,387 @@
+# copyright (c) 2023 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
+
+try:
+    from collections import Callable
+except:
+    from collections.abc import Callable
+
+import numpy as np
+import paddle
+from paddle import nn
+from paddle.nn import functional as F
+from ppocr.modeling.heads.rec_nrtr_head import Embeddings
+from ppocr.modeling.backbones.rec_svtrnet import (
+    DropPath,
+    Identity,
+    trunc_normal_,
+    zeros_,
+    ones_,
+    Mlp,
+)
+
+
+class Attention(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.q = nn.Linear(dim, dim, bias_attr=qkv_bias)
+        self.kv = nn.Linear(dim, dim * 2, bias_attr=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, q, kv):
+        N, C = kv.shape[1:]
+        QN = q.shape[1]
+        q = (
+            self.q(q)
+            .reshape([-1, QN, self.num_heads, C // self.num_heads])
+            .transpose([0, 2, 1, 3])
+        )
+        k, v = (
+            self.kv(kv)
+            .reshape([-1, N, 2, self.num_heads, C // self.num_heads])
+            .transpose((2, 0, 3, 1, 4))
+        )
+        attn = q.matmul(k.transpose((0, 1, 3, 2))) * self.scale
+        attn = F.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((-1, QN, C))
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class EdgeDecoderLayer(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=[0.0, 0.0],
+        act_layer=nn.GELU,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-6,
+    ):
+        super().__init__()
+
+        self.head_dim = dim // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path1 = DropPath(drop_path[0]) if drop_path[0] > 0.0 else Identity()
+        self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
+        self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
+
+        self.p = nn.Linear(dim, dim)
+        self.cv = nn.Linear(dim, dim)
+        self.pv = nn.Linear(dim, dim)
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.p_proj = nn.Linear(dim, dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, p, cv, pv):
+        pN = p.shape[1]
+        vN = cv.shape[1]
+        p_shortcut = p
+
+        p1 = (
+            self.p(p)
+            .reshape([-1, pN, self.num_heads, self.dim // self.num_heads])
+            .transpose([0, 2, 1, 3])
+        )
+        cv1 = (
+            self.cv(cv)
+            .reshape([-1, vN, self.num_heads, self.dim // self.num_heads])
+            .transpose([0, 2, 1, 3])
+        )
+        pv1 = (
+            self.pv(pv)
+            .reshape([-1, vN, self.num_heads, self.dim // self.num_heads])
+            .transpose([0, 2, 1, 3])
+        )
+
+        edge = F.softmax(p1.matmul(pv1.transpose((0, 1, 3, 2))), -1)  # B h N N
+        p_c = (edge @ cv1).transpose((0, 2, 1, 3)).reshape((-1, pN, self.dim))
+
+        x1 = self.norm1(p_shortcut + self.drop_path1(self.p_proj(p_c)))
+
+        x = self.norm2(x1 + self.drop_path1(self.mlp(x1)))
+        return x
+
+
+class DecoderLayer(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="nn.LayerNorm",
+        epsilon=1e-6,
+    ):
+        super().__init__()
+        if isinstance(norm_layer, str):
+            self.norm1 = eval(norm_layer)(dim, epsilon=epsilon)
+            self.normkv = eval(norm_layer)(dim, epsilon=epsilon)
+        elif isinstance(norm_layer, Callable):
+            self.norm1 = norm_layer(dim)
+            self.normkv = norm_layer(dim)
+        else:
+            raise TypeError("The norm_layer must be str or paddle.nn.LayerNorm class")
+        self.mixer = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        if isinstance(norm_layer, str):
+            self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
+        elif isinstance(norm_layer, Callable):
+            self.norm2 = norm_layer(dim)
+        else:
+            raise TypeError("The norm_layer must be str or paddle.nn.layer.Layer class")
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_ratio = mlp_ratio
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, q, kv):
+        x1 = self.norm1(q + self.drop_path(self.mixer(q, kv)))
+        x = self.norm2(x1 + self.drop_path(self.mlp(x1)))
+        return x
+
+
+class CPPDHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        dim,
+        out_channels,
+        num_layer=2,
+        drop_path_rate=0.1,
+        max_len=25,
+        vis_seq=50,
+        ch=False,
+        **kwargs,
+    ):
+        super(CPPDHead, self).__init__()
+
+        self.out_channels = out_channels  # none + 26 + 10
+        self.dim = dim
+        self.ch = ch
+        self.max_len = max_len + 1  # max_len + eos
+        self.char_node_embed = Embeddings(
+            d_model=dim, vocab=self.out_channels, scale_embedding=True
+        )
+        self.pos_node_embed = Embeddings(
+            d_model=dim, vocab=self.max_len, scale_embedding=True
+        )
+        dpr = np.linspace(0, drop_path_rate, num_layer + 1)
+
+        self.char_node_decoder = nn.LayerList(
+            [
+                DecoderLayer(
+                    dim=dim,
+                    num_heads=dim // 32,
+                    mlp_ratio=4.0,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                )
+                for i in range(num_layer)
+            ]
+        )
+        self.pos_node_decoder = nn.LayerList(
+            [
+                DecoderLayer(
+                    dim=dim,
+                    num_heads=dim // 32,
+                    mlp_ratio=4.0,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                )
+                for i in range(num_layer)
+            ]
+        )
+
+        self.edge_decoder = EdgeDecoderLayer(
+            dim=dim,
+            num_heads=dim // 32,
+            mlp_ratio=4.0,
+            qkv_bias=True,
+            drop_path=dpr[num_layer : num_layer + 1],
+        )
+
+        self.char_pos_embed = self.create_parameter(
+            shape=[1, self.max_len, dim], default_initializer=zeros_
+        )
+        self.add_parameter("char_pos_embed", self.char_pos_embed)
+        self.vis_pos_embed = self.create_parameter(
+            shape=[1, vis_seq, dim], default_initializer=zeros_
+        )
+        self.add_parameter("vis_pos_embed", self.vis_pos_embed)
+
+        self.char_node_fc1 = nn.Linear(dim, max_len)
+        self.pos_node_fc1 = nn.Linear(dim, self.max_len)
+
+        self.edge_fc = nn.Linear(dim, self.out_channels)
+        trunc_normal_(self.char_pos_embed)
+        trunc_normal_(self.vis_pos_embed)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward(self, x, targets=None, epoch=0):
+        if self.training:
+            return self.forward_train(x, targets, epoch)
+        else:
+            return self.forward_test(x)
+
+    def forward_test(self, x):
+        visual_feats = x + self.vis_pos_embed
+        bs = visual_feats.shape[0]
+        pos_node_embed = (
+            self.pos_node_embed(paddle.arange(self.max_len)).unsqueeze(0)
+            + self.char_pos_embed
+        )
+        pos_node_embed = paddle.tile(pos_node_embed, [bs, 1, 1])
+        char_vis_node_query = visual_feats
+        pos_vis_node_query = paddle.concat([pos_node_embed, visual_feats], 1)
+
+        for char_decoder_layer, pos_decoder_layer in zip(
+            self.char_node_decoder, self.pos_node_decoder
+        ):
+            char_vis_node_query = char_decoder_layer(
+                char_vis_node_query, char_vis_node_query
+            )
+            pos_vis_node_query = pos_decoder_layer(
+                pos_vis_node_query, pos_vis_node_query[:, self.max_len :, :]
+            )
+        pos_node_query = pos_vis_node_query[:, : self.max_len, :]
+        char_vis_feats = char_vis_node_query
+
+        pos_node_feats = self.edge_decoder(
+            pos_node_query, char_vis_feats, char_vis_feats
+        )  # B, 26, dim
+        edge_feats = self.edge_fc(pos_node_feats)  # B, 26, 37
+        edge_logits = F.softmax(edge_feats, -1)
+
+        return edge_logits
+
+    def forward_train(self, x, targets=None, epoch=0):
+        visual_feats = x + self.vis_pos_embed
+        bs = visual_feats.shape[0]
+
+        if self.ch:
+            char_node_embed = self.char_node_embed(targets[-2])
+        else:
+            char_node_embed = self.char_node_embed(
+                paddle.arange(self.out_channels)
+            ).unsqueeze(0)
+            char_node_embed = paddle.tile(char_node_embed, [bs, 1, 1])
+        counting_char_num = char_node_embed.shape[1]
+        pos_node_embed = (
+            self.pos_node_embed(paddle.arange(self.max_len)).unsqueeze(0)
+            + self.char_pos_embed
+        )
+        pos_node_embed = paddle.tile(pos_node_embed, [bs, 1, 1])
+
+        node_feats = []
+
+        char_vis_node_query = paddle.concat([char_node_embed, visual_feats], 1)
+        pos_vis_node_query = paddle.concat([pos_node_embed, visual_feats], 1)
+
+        for char_decoder_layer, pos_decoder_layer in zip(
+            self.char_node_decoder, self.pos_node_decoder
+        ):
+            char_vis_node_query = char_decoder_layer(
+                char_vis_node_query, char_vis_node_query[:, counting_char_num:, :]
+            )
+            pos_vis_node_query = pos_decoder_layer(
+                pos_vis_node_query, pos_vis_node_query[:, self.max_len :, :]
+            )
+
+        char_node_query = char_vis_node_query[:, :counting_char_num, :]
+        pos_node_query = pos_vis_node_query[:, : self.max_len, :]
+
+        char_vis_feats = char_vis_node_query[:, counting_char_num:, :]
+        char_node_feats1 = self.char_node_fc1(char_node_query)
+
+        pos_node_feats1 = self.pos_node_fc1(pos_node_query)
+        diag_mask = (
+            paddle.eye(pos_node_feats1.shape[1])
+            .unsqueeze(0)
+            .tile([pos_node_feats1.shape[0], 1, 1])
+        )
+        pos_node_feats1 = (pos_node_feats1 * diag_mask).sum(-1)
+
+        node_feats.append(char_node_feats1)
+        node_feats.append(pos_node_feats1)
+
+        pos_node_feats = self.edge_decoder(
+            pos_node_query, char_vis_feats, char_vis_feats
+        )  # B, 26, dim
+        edge_feats = self.edge_fc(pos_node_feats)  # B, 26, 37
+
+        return node_feats, edge_feats

ppocr/modeling/heads/rec_ctc_head.py

@@ -0,0 +1,92 @@
+# copyright (c) 2019 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 ParamAttr, nn
+from paddle.nn import functional as F
+
+
+def get_para_bias_attr(l2_decay, k):
+    regularizer = paddle.regularizer.L2Decay(l2_decay)
+    stdv = 1.0 / math.sqrt(k * 1.0)
+    initializer = nn.initializer.Uniform(-stdv, stdv)
+    weight_attr = ParamAttr(regularizer=regularizer, initializer=initializer)
+    bias_attr = ParamAttr(regularizer=regularizer, initializer=initializer)
+    return [weight_attr, bias_attr]
+
+
+class CTCHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        fc_decay=0.0004,
+        mid_channels=None,
+        return_feats=False,
+        **kwargs,
+    ):
+        super(CTCHead, self).__init__()
+        if mid_channels is None:
+            weight_attr, bias_attr = get_para_bias_attr(
+                l2_decay=fc_decay, k=in_channels
+            )
+            self.fc = nn.Linear(
+                in_channels, out_channels, weight_attr=weight_attr, bias_attr=bias_attr
+            )
+        else:
+            weight_attr1, bias_attr1 = get_para_bias_attr(
+                l2_decay=fc_decay, k=in_channels
+            )
+            self.fc1 = nn.Linear(
+                in_channels,
+                mid_channels,
+                weight_attr=weight_attr1,
+                bias_attr=bias_attr1,
+            )
+
+            weight_attr2, bias_attr2 = get_para_bias_attr(
+                l2_decay=fc_decay, k=mid_channels
+            )
+            self.fc2 = nn.Linear(
+                mid_channels,
+                out_channels,
+                weight_attr=weight_attr2,
+                bias_attr=bias_attr2,
+            )
+        self.out_channels = out_channels
+        self.mid_channels = mid_channels
+        self.return_feats = return_feats
+
+    def forward(self, x, targets=None):
+        if self.mid_channels is None:
+            predicts = self.fc(x)
+        else:
+            x = self.fc1(x)
+            predicts = self.fc2(x)
+
+        if self.return_feats:
+            result = (x, predicts)
+        else:
+            result = predicts
+        if not self.training:
+            predicts = F.softmax(predicts, axis=2)
+            result = predicts
+
+        return result

ppocr/modeling/heads/rec_multi_head.py

@@ -0,0 +1,153 @@
+# 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 ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from ppocr.modeling.necks.rnn import (
+    Im2Seq,
+    EncoderWithRNN,
+    EncoderWithFC,
+    SequenceEncoder,
+    EncoderWithSVTR,
+    trunc_normal_,
+    zeros_,
+)
+from .rec_ctc_head import CTCHead
+from .rec_sar_head import SARHead
+from .rec_nrtr_head import Transformer
+
+
+class FCTranspose(nn.Layer):
+    def __init__(self, in_channels, out_channels, only_transpose=False):
+        super().__init__()
+        self.only_transpose = only_transpose
+        if not self.only_transpose:
+            self.fc = nn.Linear(in_channels, out_channels, bias_attr=False)
+
+    def forward(self, x):
+        if self.only_transpose:
+            return x.transpose([0, 2, 1])
+        else:
+            return self.fc(x.transpose([0, 2, 1]))
+
+
+class AddPos(nn.Layer):
+    def __init__(self, dim, w):
+        super().__init__()
+        self.dec_pos_embed = self.create_parameter(
+            shape=[1, w, dim], default_initializer=zeros_
+        )
+        self.add_parameter("dec_pos_embed", self.dec_pos_embed)
+        trunc_normal_(self.dec_pos_embed)
+
+    def forward(self, x):
+        x = x + self.dec_pos_embed[:, : x.shape[1], :]
+        return x
+
+
+class MultiHead(nn.Layer):
+    def __init__(self, in_channels, out_channels_list, **kwargs):
+        super().__init__()
+        self.head_list = kwargs.pop("head_list")
+        self.use_pool = kwargs.get("use_pool", False)
+        self.use_pos = kwargs.get("use_pos", False)
+        self.in_channels = in_channels
+        if self.use_pool:
+            self.pool = nn.AvgPool2D(kernel_size=[3, 2], stride=[3, 2], padding=0)
+        self.gtc_head = "sar"
+        assert len(self.head_list) >= 2
+        for idx, head_name in enumerate(self.head_list):
+            name = list(head_name)[0]
+            if name == "SARHead":
+                # sar head
+                sar_args = self.head_list[idx][name]
+                self.sar_head = eval(name)(
+                    in_channels=in_channels,
+                    out_channels=out_channels_list["SARLabelDecode"],
+                    **sar_args,
+                )
+            elif name == "NRTRHead":
+                gtc_args = self.head_list[idx][name]
+                max_text_length = gtc_args.get("max_text_length", 25)
+                nrtr_dim = gtc_args.get("nrtr_dim", 256)
+                num_decoder_layers = gtc_args.get("num_decoder_layers", 4)
+                if self.use_pos:
+                    self.before_gtc = nn.Sequential(
+                        nn.Flatten(2),
+                        FCTranspose(in_channels, nrtr_dim),
+                        AddPos(nrtr_dim, 80),
+                    )
+                else:
+                    self.before_gtc = nn.Sequential(
+                        nn.Flatten(2), FCTranspose(in_channels, nrtr_dim)
+                    )
+
+                self.gtc_head = Transformer(
+                    d_model=nrtr_dim,
+                    nhead=nrtr_dim // 32,
+                    num_encoder_layers=-1,
+                    beam_size=-1,
+                    num_decoder_layers=num_decoder_layers,
+                    max_len=max_text_length,
+                    dim_feedforward=nrtr_dim * 4,
+                    out_channels=out_channels_list["NRTRLabelDecode"],
+                )
+            elif name == "CTCHead":
+                # ctc neck
+                self.encoder_reshape = Im2Seq(in_channels)
+                neck_args = self.head_list[idx][name]["Neck"]
+                encoder_type = neck_args.pop("name")
+                self.ctc_encoder = SequenceEncoder(
+                    in_channels=in_channels, encoder_type=encoder_type, **neck_args
+                )
+                # ctc head
+                head_args = self.head_list[idx][name]["Head"]
+                self.ctc_head = eval(name)(
+                    in_channels=self.ctc_encoder.out_channels,
+                    out_channels=out_channels_list["CTCLabelDecode"],
+                    **head_args,
+                )
+            else:
+                raise NotImplementedError(
+                    "{} is not supported in MultiHead yet".format(name)
+                )
+
+    def forward(self, x, targets=None):
+        if self.use_pool:
+            x = self.pool(
+                x.reshape([0, 3, -1, self.in_channels]).transpose([0, 3, 1, 2])
+            )
+        ctc_encoder = self.ctc_encoder(x)
+        ctc_out = self.ctc_head(ctc_encoder, targets)
+        head_out = dict()
+        head_out["ctc"] = ctc_out
+        head_out["ctc_neck"] = ctc_encoder
+        # eval mode
+        if not self.training:
+            return ctc_out
+        if self.gtc_head == "sar":
+            sar_out = self.sar_head(x, targets[1:])
+            head_out["sar"] = sar_out
+        else:
+            gtc_out = self.gtc_head(self.before_gtc(x), targets[1:])
+            head_out["gtc"] = gtc_out
+        return head_out

ppocr/modeling/heads/rec_nrtr_head.py

@@ -0,0 +1,705 @@
+# copyright (c) 2021 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 math
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle.nn import Dropout, LayerNorm
+import numpy as np
+from ppocr.modeling.backbones.rec_svtrnet import Mlp, zeros_
+from paddle.nn.initializer import XavierNormal as xavier_normal_
+
+
+class Transformer(nn.Layer):
+    """A transformer model. User is able to modify the attributes as needed. The architecture
+    is based on the paper "Attention Is All You Need". Ashish Vaswani, Noam Shazeer,
+    Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and
+    Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information
+    Processing Systems, pages 6000-6010.
+
+    Args:
+        d_model: the number of expected features in the encoder/decoder inputs (default=512).
+        nhead: the number of heads in the multiheadattention models (default=8).
+        num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
+        num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+        custom_encoder: custom encoder (default=None).
+        custom_decoder: custom decoder (default=None).
+    """
+
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        num_encoder_layers=6,
+        beam_size=0,
+        num_decoder_layers=6,
+        max_len=25,
+        dim_feedforward=1024,
+        attention_dropout_rate=0.0,
+        residual_dropout_rate=0.1,
+        in_channels=0,
+        out_channels=0,
+        scale_embedding=True,
+    ):
+        super(Transformer, self).__init__()
+        self.out_channels = out_channels + 1
+        self.max_len = max_len
+        self.embedding = Embeddings(
+            d_model=d_model,
+            vocab=self.out_channels,
+            padding_idx=0,
+            scale_embedding=scale_embedding,
+        )
+        self.positional_encoding = PositionalEncoding(
+            dropout=residual_dropout_rate, dim=d_model
+        )
+
+        if num_encoder_layers > 0:
+            self.encoder = nn.LayerList(
+                [
+                    TransformerBlock(
+                        d_model,
+                        nhead,
+                        dim_feedforward,
+                        attention_dropout_rate,
+                        residual_dropout_rate,
+                        with_self_attn=True,
+                        with_cross_attn=False,
+                    )
+                    for i in range(num_encoder_layers)
+                ]
+            )
+        else:
+            self.encoder = None
+
+        self.decoder = nn.LayerList(
+            [
+                TransformerBlock(
+                    d_model,
+                    nhead,
+                    dim_feedforward,
+                    attention_dropout_rate,
+                    residual_dropout_rate,
+                    with_self_attn=True,
+                    with_cross_attn=True,
+                )
+                for i in range(num_decoder_layers)
+            ]
+        )
+
+        self.beam_size = beam_size
+        self.d_model = d_model
+        self.nhead = nhead
+        self.tgt_word_prj = nn.Linear(d_model, self.out_channels, bias_attr=False)
+        w0 = np.random.normal(
+            0.0, d_model**-0.5, (d_model, self.out_channels)
+        ).astype(np.float32)
+        self.tgt_word_prj.weight.set_value(w0)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            xavier_normal_(m.weight)
+            if m.bias is not None:
+                zeros_(m.bias)
+
+    def forward_train(self, src, tgt):
+        tgt = tgt[:, :-1]
+
+        tgt = self.embedding(tgt)
+        tgt = self.positional_encoding(tgt)
+        tgt_mask = self.generate_square_subsequent_mask(tgt.shape[1])
+
+        if self.encoder is not None:
+            src = self.positional_encoding(src)
+            for encoder_layer in self.encoder:
+                src = encoder_layer(src)
+            memory = src  # B N C
+        else:
+            memory = src  # B N C
+        for decoder_layer in self.decoder:
+            tgt = decoder_layer(tgt, memory, self_mask=tgt_mask)
+        output = tgt
+        logit = self.tgt_word_prj(output)
+        return logit
+
+    def forward(self, src, targets=None):
+        """Take in and process masked source/target sequences.
+        Args:
+            src: the sequence to the encoder (required).
+            tgt: the sequence to the decoder (required).
+        Shape:
+            - src: :math:`(B, sN, C)`.
+            - tgt: :math:`(B, tN, C)`.
+        Examples:
+            >>> output = transformer_model(src, tgt)
+        """
+
+        if self.training:
+            max_len = targets[1].max()
+            tgt = targets[0][:, : 2 + max_len]
+            return self.forward_train(src, tgt)
+        else:
+            if self.beam_size > 0:
+                return self.forward_beam(src)
+            else:
+                return self.forward_test(src)
+
+    def forward_test(self, src):
+        bs = src.shape[0]
+        if self.encoder is not None:
+            src = self.positional_encoding(src)
+            for encoder_layer in self.encoder:
+                src = encoder_layer(src)
+            memory = src  # B N C
+        else:
+            memory = src
+        dec_seq = paddle.full((bs, 1), 2, dtype=paddle.int64)
+        dec_prob = paddle.full((bs, 1), 1.0, dtype=paddle.float32)
+        for len_dec_seq in range(1, paddle.to_tensor(self.max_len)):
+            dec_seq_embed = self.embedding(dec_seq)
+            dec_seq_embed = self.positional_encoding(dec_seq_embed)
+            tgt_mask = self.generate_square_subsequent_mask(dec_seq_embed.shape[1])
+            tgt = dec_seq_embed
+            for decoder_layer in self.decoder:
+                tgt = decoder_layer(tgt, memory, self_mask=tgt_mask)
+            dec_output = tgt
+            dec_output = dec_output[:, -1, :]
+            word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=-1)
+            preds_idx = paddle.argmax(word_prob, axis=-1)
+            if paddle.equal_all(
+                preds_idx, paddle.full(preds_idx.shape, 3, dtype="int64")
+            ):
+                break
+            preds_prob = paddle.max(word_prob, axis=-1)
+            dec_seq = paddle.concat(
+                [dec_seq, paddle.reshape(preds_idx, [-1, 1])], axis=1
+            )
+            dec_prob = paddle.concat(
+                [dec_prob, paddle.reshape(preds_prob, [-1, 1])], axis=1
+            )
+        return [dec_seq, dec_prob]
+
+    def forward_beam(self, images):
+        """Translation work in one batch"""
+
+        def get_inst_idx_to_tensor_position_map(inst_idx_list):
+            """Indicate the position of an instance in a tensor."""
+            return {
+                inst_idx: tensor_position
+                for tensor_position, inst_idx in enumerate(inst_idx_list)
+            }
+
+        def collect_active_part(
+            beamed_tensor, curr_active_inst_idx, n_prev_active_inst, n_bm
+        ):
+            """Collect tensor parts associated to active instances."""
+
+            beamed_tensor_shape = beamed_tensor.shape
+            n_curr_active_inst = len(curr_active_inst_idx)
+            new_shape = (
+                n_curr_active_inst * n_bm,
+                beamed_tensor_shape[1],
+                beamed_tensor_shape[2],
+            )
+
+            beamed_tensor = beamed_tensor.reshape([n_prev_active_inst, -1])
+            beamed_tensor = beamed_tensor.index_select(curr_active_inst_idx, axis=0)
+            beamed_tensor = beamed_tensor.reshape(new_shape)
+
+            return beamed_tensor
+
+        def collate_active_info(
+            src_enc, inst_idx_to_position_map, active_inst_idx_list
+        ):
+            # Sentences which are still active are collected,
+            # so the decoder will not run on completed sentences.
+
+            n_prev_active_inst = len(inst_idx_to_position_map)
+            active_inst_idx = [
+                inst_idx_to_position_map[k] for k in active_inst_idx_list
+            ]
+            active_inst_idx = paddle.to_tensor(active_inst_idx, dtype="int64")
+            active_src_enc = collect_active_part(
+                src_enc.transpose([1, 0, 2]), active_inst_idx, n_prev_active_inst, n_bm
+            ).transpose([1, 0, 2])
+            active_inst_idx_to_position_map = get_inst_idx_to_tensor_position_map(
+                active_inst_idx_list
+            )
+            return active_src_enc, active_inst_idx_to_position_map
+
+        def beam_decode_step(
+            inst_dec_beams, len_dec_seq, enc_output, inst_idx_to_position_map, n_bm
+        ):
+            """Decode and update beam status, and then return active beam idx"""
+
+            def prepare_beam_dec_seq(inst_dec_beams, len_dec_seq):
+                dec_partial_seq = [
+                    b.get_current_state() for b in inst_dec_beams if not b.done
+                ]
+                dec_partial_seq = paddle.stack(dec_partial_seq)
+                dec_partial_seq = dec_partial_seq.reshape([-1, len_dec_seq])
+                return dec_partial_seq
+
+            def predict_word(dec_seq, enc_output, n_active_inst, n_bm):
+                dec_seq = self.embedding(dec_seq)
+                dec_seq = self.positional_encoding(dec_seq)
+                tgt_mask = self.generate_square_subsequent_mask(dec_seq.shape[1])
+                tgt = dec_seq
+                for decoder_layer in self.decoder:
+                    tgt = decoder_layer(tgt, enc_output, self_mask=tgt_mask)
+                dec_output = tgt
+                dec_output = dec_output[:, -1, :]  # Pick the last step: (bh * bm) * d_h
+                word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=1)
+                word_prob = paddle.reshape(word_prob, [n_active_inst, n_bm, -1])
+                return word_prob
+
+            def collect_active_inst_idx_list(
+                inst_beams, word_prob, inst_idx_to_position_map
+            ):
+                active_inst_idx_list = []
+                for inst_idx, inst_position in inst_idx_to_position_map.items():
+                    is_inst_complete = inst_beams[inst_idx].advance(
+                        word_prob[inst_position]
+                    )
+                    if not is_inst_complete:
+                        active_inst_idx_list += [inst_idx]
+
+                return active_inst_idx_list
+
+            n_active_inst = len(inst_idx_to_position_map)
+            dec_seq = prepare_beam_dec_seq(inst_dec_beams, len_dec_seq)
+            word_prob = predict_word(dec_seq, enc_output, n_active_inst, n_bm)
+            # Update the beam with predicted word prob information and collect incomplete instances
+            active_inst_idx_list = collect_active_inst_idx_list(
+                inst_dec_beams, word_prob, inst_idx_to_position_map
+            )
+            return active_inst_idx_list
+
+        def collect_hypothesis_and_scores(inst_dec_beams, n_best):
+            all_hyp, all_scores = [], []
+            for inst_idx in range(len(inst_dec_beams)):
+                scores, tail_idxs = inst_dec_beams[inst_idx].sort_scores()
+                all_scores += [scores[:n_best]]
+                hyps = [
+                    inst_dec_beams[inst_idx].get_hypothesis(i)
+                    for i in tail_idxs[:n_best]
+                ]
+                all_hyp += [hyps]
+            return all_hyp, all_scores
+
+        with paddle.no_grad():
+            # -- Encode
+            if self.encoder is not None:
+                src = self.positional_encoding(images)
+                src_enc = self.encoder(src)
+            else:
+                src_enc = images
+
+            n_bm = self.beam_size
+            src_shape = src_enc.shape
+            inst_dec_beams = [Beam(n_bm) for _ in range(1)]
+            active_inst_idx_list = list(range(1))
+            # Repeat data for beam search
+            src_enc = paddle.tile(src_enc, [1, n_bm, 1])
+            inst_idx_to_position_map = get_inst_idx_to_tensor_position_map(
+                active_inst_idx_list
+            )
+            # Decode
+            for len_dec_seq in range(1, paddle.to_tensor(self.max_len)):
+                src_enc_copy = src_enc.clone()
+                active_inst_idx_list = beam_decode_step(
+                    inst_dec_beams,
+                    len_dec_seq,
+                    src_enc_copy,
+                    inst_idx_to_position_map,
+                    n_bm,
+                )
+                if not active_inst_idx_list:
+                    break  # all instances have finished their path to <EOS>
+                src_enc, inst_idx_to_position_map = collate_active_info(
+                    src_enc_copy, inst_idx_to_position_map, active_inst_idx_list
+                )
+        batch_hyp, batch_scores = collect_hypothesis_and_scores(inst_dec_beams, 1)
+        result_hyp = []
+        hyp_scores = []
+        for bs_hyp, score in zip(batch_hyp, batch_scores):
+            l = len(bs_hyp[0])
+            bs_hyp_pad = bs_hyp[0] + [3] * (25 - l)
+            result_hyp.append(bs_hyp_pad)
+            score = float(score) / l
+            hyp_score = [score for _ in range(25)]
+            hyp_scores.append(hyp_score)
+        return [
+            paddle.to_tensor(np.array(result_hyp), dtype=paddle.int64),
+            paddle.to_tensor(hyp_scores),
+        ]
+
+    def generate_square_subsequent_mask(self, sz):
+        """Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+        Unmasked positions are filled with float(0.0).
+        """
+        mask = paddle.zeros([sz, sz], dtype="float32")
+        mask_inf = paddle.triu(
+            paddle.full(shape=[sz, sz], dtype="float32", fill_value=float("-inf")),
+            diagonal=1,
+        )
+        mask = mask + mask_inf
+        return mask.unsqueeze([0, 1])
+
+
+class MultiheadAttention(nn.Layer):
+    """Allows the model to jointly attend to information
+    from different representation subspaces.
+    See reference: Attention Is All You Need
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
+
+    Args:
+        embed_dim: total dimension of the model
+        num_heads: parallel attention layers, or heads
+
+    """
+
+    def __init__(self, embed_dim, num_heads, dropout=0.0, self_attn=False):
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        # self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scale = self.head_dim**-0.5
+        self.self_attn = self_attn
+        if self_attn:
+            self.qkv = nn.Linear(embed_dim, embed_dim * 3)
+        else:
+            self.q = nn.Linear(embed_dim, embed_dim)
+            self.kv = nn.Linear(embed_dim, embed_dim * 2)
+        self.attn_drop = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, query, key=None, attn_mask=None):
+        qN = query.shape[1]
+
+        if self.self_attn:
+            qkv = (
+                self.qkv(query)
+                .reshape((0, qN, 3, self.num_heads, self.head_dim))
+                .transpose((2, 0, 3, 1, 4))
+            )
+            q, k, v = qkv[0], qkv[1], qkv[2]
+        else:
+            kN = key.shape[1]
+            q = (
+                self.q(query)
+                .reshape([0, qN, self.num_heads, self.head_dim])
+                .transpose([0, 2, 1, 3])
+            )
+            kv = (
+                self.kv(key)
+                .reshape((0, kN, 2, self.num_heads, self.head_dim))
+                .transpose((2, 0, 3, 1, 4))
+            )
+            k, v = kv[0], kv[1]
+
+        attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
+
+        if attn_mask is not None:
+            attn += attn_mask
+
+        attn = F.softmax(attn, axis=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((0, qN, self.embed_dim))
+        x = self.out_proj(x)
+
+        return x
+
+
+class TransformerBlock(nn.Layer):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        attention_dropout_rate=0.0,
+        residual_dropout_rate=0.1,
+        with_self_attn=True,
+        with_cross_attn=False,
+        epsilon=1e-5,
+    ):
+        super(TransformerBlock, self).__init__()
+        self.with_self_attn = with_self_attn
+        if with_self_attn:
+            self.self_attn = MultiheadAttention(
+                d_model, nhead, dropout=attention_dropout_rate, self_attn=with_self_attn
+            )
+            self.norm1 = LayerNorm(d_model, epsilon=epsilon)
+            self.dropout1 = Dropout(residual_dropout_rate)
+        self.with_cross_attn = with_cross_attn
+        if with_cross_attn:
+            self.cross_attn = (
+                MultiheadAttention(  # for self_attn of encoder or cross_attn of decoder
+                    d_model, nhead, dropout=attention_dropout_rate
+                )
+            )
+            self.norm2 = LayerNorm(d_model, epsilon=epsilon)
+            self.dropout2 = Dropout(residual_dropout_rate)
+
+        self.mlp = Mlp(
+            in_features=d_model,
+            hidden_features=dim_feedforward,
+            act_layer=nn.ReLU,
+            drop=residual_dropout_rate,
+        )
+
+        self.norm3 = LayerNorm(d_model, epsilon=epsilon)
+
+        self.dropout3 = Dropout(residual_dropout_rate)
+
+    def forward(self, tgt, memory=None, self_mask=None, cross_mask=None):
+        if self.with_self_attn:
+            tgt1 = self.self_attn(tgt, attn_mask=self_mask)
+            tgt = self.norm1(tgt + self.dropout1(tgt1))
+
+        if self.with_cross_attn:
+            tgt2 = self.cross_attn(tgt, key=memory, attn_mask=cross_mask)
+            tgt = self.norm2(tgt + self.dropout2(tgt2))
+        tgt = self.norm3(tgt + self.dropout3(self.mlp(tgt)))
+        return tgt
+
+
+class PositionalEncoding(nn.Layer):
+    """Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)
+    Args:
+        d_model: the embed dim (required).
+        dropout: the dropout value (default=0.1).
+        max_len: the max. length of the incoming sequence (default=5000).
+    Examples:
+        >>> pos_encoder = PositionalEncoding(d_model)
+    """
+
+    def __init__(self, dropout, dim, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = paddle.zeros([max_len, dim])
+        position = paddle.arange(0, max_len, dtype=paddle.float32).unsqueeze(1)
+        div_term = paddle.exp(
+            paddle.arange(0, dim, 2).astype("float32") * (-math.log(10000.0) / dim)
+        )
+        pe[:, 0::2] = paddle.sin(position * div_term)
+        pe[:, 1::2] = paddle.cos(position * div_term)
+        pe = paddle.unsqueeze(pe, 0)
+        pe = paddle.transpose(pe, [1, 0, 2])
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """Inputs of forward function
+        Args:
+            x: the sequence fed to the positional encoder model (required).
+        Shape:
+            x: [sequence length, batch size, embed dim]
+            output: [sequence length, batch size, embed dim]
+        Examples:
+            >>> output = pos_encoder(x)
+        """
+        x = x.transpose([1, 0, 2])
+        x = x + self.pe[: x.shape[0], :]
+        return self.dropout(x).transpose([1, 0, 2])
+
+
+class PositionalEncoding_2d(nn.Layer):
+    """Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)
+    Args:
+        d_model: the embed dim (required).
+        dropout: the dropout value (default=0.1).
+        max_len: the max. length of the incoming sequence (default=5000).
+    Examples:
+        >>> pos_encoder = PositionalEncoding(d_model)
+    """
+
+    def __init__(self, dropout, dim, max_len=5000):
+        super(PositionalEncoding_2d, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = paddle.zeros([max_len, dim])
+        position = paddle.arange(0, max_len, dtype=paddle.float32).unsqueeze(1)
+        div_term = paddle.exp(
+            paddle.arange(0, dim, 2).astype("float32") * (-math.log(10000.0) / dim)
+        )
+        pe[:, 0::2] = paddle.sin(position * div_term)
+        pe[:, 1::2] = paddle.cos(position * div_term)
+        pe = paddle.transpose(paddle.unsqueeze(pe, 0), [1, 0, 2])
+        self.register_buffer("pe", pe)
+
+        self.avg_pool_1 = nn.AdaptiveAvgPool2D((1, 1))
+        self.linear1 = nn.Linear(dim, dim)
+        self.linear1.weight.data.fill_(1.0)
+        self.avg_pool_2 = nn.AdaptiveAvgPool2D((1, 1))
+        self.linear2 = nn.Linear(dim, dim)
+        self.linear2.weight.data.fill_(1.0)
+
+    def forward(self, x):
+        """Inputs of forward function
+        Args:
+            x: the sequence fed to the positional encoder model (required).
+        Shape:
+            x: [sequence length, batch size, embed dim]
+            output: [sequence length, batch size, embed dim]
+        Examples:
+            >>> output = pos_encoder(x)
+        """
+        w_pe = self.pe[: x.shape[-1], :]
+        w1 = self.linear1(self.avg_pool_1(x).squeeze()).unsqueeze(0)
+        w_pe = w_pe * w1
+        w_pe = paddle.transpose(w_pe, [1, 2, 0])
+        w_pe = paddle.unsqueeze(w_pe, 2)
+
+        h_pe = self.pe[: x.shape.shape[-2], :]
+        w2 = self.linear2(self.avg_pool_2(x).squeeze()).unsqueeze(0)
+        h_pe = h_pe * w2
+        h_pe = paddle.transpose(h_pe, [1, 2, 0])
+        h_pe = paddle.unsqueeze(h_pe, 3)
+
+        x = x + w_pe + h_pe
+        x = paddle.transpose(
+            paddle.reshape(x, [x.shape[0], x.shape[1], x.shape[2] * x.shape[3]]),
+            [2, 0, 1],
+        )
+
+        return self.dropout(x)
+
+
+class Embeddings(nn.Layer):
+    def __init__(self, d_model, vocab, padding_idx=None, scale_embedding=True):
+        super(Embeddings, self).__init__()
+        self.embedding = nn.Embedding(vocab, d_model, padding_idx=padding_idx)
+        w0 = np.random.normal(0.0, d_model**-0.5, (vocab, d_model)).astype(np.float32)
+        self.embedding.weight.set_value(w0)
+        self.d_model = d_model
+        self.scale_embedding = scale_embedding
+
+    def forward(self, x):
+        if self.scale_embedding:
+            x = self.embedding(x)
+            return x * math.sqrt(self.d_model)
+        return self.embedding(x)
+
+
+class Beam:
+    """Beam search"""
+
+    def __init__(self, size, device=False):
+        self.size = size
+        self._done = False
+        # The score for each translation on the beam.
+        self.scores = paddle.zeros((size,), dtype=paddle.float32)
+        self.all_scores = []
+        # The backpointers at each time-step.
+        self.prev_ks = []
+        # The outputs at each time-step.
+        self.next_ys = [paddle.full((size,), 0, dtype=paddle.int64)]
+        self.next_ys[0][0] = 2
+
+    def get_current_state(self):
+        "Get the outputs for the current timestep."
+        return self.get_tentative_hypothesis()
+
+    def get_current_origin(self):
+        "Get the backpointers for the current timestep."
+        return self.prev_ks[-1]
+
+    @property
+    def done(self):
+        return self._done
+
+    def advance(self, word_prob):
+        "Update beam status and check if finished or not."
+        num_words = word_prob.shape[1]
+
+        # Sum the previous scores.
+        if len(self.prev_ks) > 0:
+            beam_lk = word_prob + self.scores.unsqueeze(1).expand_as(word_prob)
+        else:
+            beam_lk = word_prob[0]
+
+        flat_beam_lk = beam_lk.reshape([-1])
+        best_scores, best_scores_id = flat_beam_lk.topk(
+            self.size, 0, True, True
+        )  # 1st sort
+        self.all_scores.append(self.scores)
+        self.scores = best_scores
+        # bestScoresId is flattened as a (beam x word) array,
+        # so we need to calculate which word and beam each score came from
+        prev_k = best_scores_id // num_words
+        self.prev_ks.append(prev_k)
+        self.next_ys.append(best_scores_id - prev_k * num_words)
+        # End condition is when top-of-beam is EOS.
+        if self.next_ys[-1][0] == 3:
+            self._done = True
+            self.all_scores.append(self.scores)
+
+        return self._done
+
+    def sort_scores(self):
+        "Sort the scores."
+        return self.scores, paddle.to_tensor(
+            [i for i in range(int(self.scores.shape[0]))], dtype="int32"
+        )
+
+    def get_the_best_score_and_idx(self):
+        "Get the score of the best in the beam."
+        scores, ids = self.sort_scores()
+        return scores[1], ids[1]
+
+    def get_tentative_hypothesis(self):
+        "Get the decoded sequence for the current timestep."
+        if len(self.next_ys) == 1:
+            dec_seq = self.next_ys[0].unsqueeze(1)
+        else:
+            _, keys = self.sort_scores()
+            hyps = [self.get_hypothesis(k) for k in keys]
+            hyps = [[2] + h for h in hyps]
+            dec_seq = paddle.to_tensor(hyps, dtype="int64")
+        return dec_seq
+
+    def get_hypothesis(self, k):
+        """Walk back to construct the full hypothesis."""
+        hyp = []
+        for j in range(len(self.prev_ks) - 1, -1, -1):
+            hyp.append(self.next_ys[j + 1][k])
+            k = self.prev_ks[j][k]
+        return list(map(lambda x: x.item(), hyp[::-1]))

ppocr/modeling/heads/rec_parseq_head.py

@@ -0,0 +1,504 @@
+# 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.
+
+# Code was based on https://github.com/baudm/parseq/blob/main/strhub/models/parseq/system.py
+# reference: https://arxiv.org/abs/2207.06966
+
+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 .self_attention import WrapEncoderForFeature
+from .self_attention import WrapEncoder
+from collections import OrderedDict
+from typing import Optional
+import copy
+from itertools import permutations
+
+
+class DecoderLayer(paddle.nn.Layer):
+    """A Transformer decoder layer supporting two-stream attention (XLNet)
+    This implements a pre-LN decoder, as opposed to the post-LN default in PyTorch."""
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation="gelu",
+        layer_norm_eps=1e-05,
+    ):
+        super().__init__()
+        self.self_attn = paddle.nn.MultiHeadAttention(
+            d_model, nhead, dropout=dropout, need_weights=True
+        )  # paddle.nn.MultiHeadAttention默认为batch_first模式
+        self.cross_attn = paddle.nn.MultiHeadAttention(
+            d_model, nhead, dropout=dropout, need_weights=True
+        )
+        self.linear1 = paddle.nn.Linear(
+            in_features=d_model, out_features=dim_feedforward
+        )
+        self.dropout = paddle.nn.Dropout(p=dropout)
+        self.linear2 = paddle.nn.Linear(
+            in_features=dim_feedforward, out_features=d_model
+        )
+        self.norm1 = paddle.nn.LayerNorm(
+            normalized_shape=d_model, epsilon=layer_norm_eps
+        )
+        self.norm2 = paddle.nn.LayerNorm(
+            normalized_shape=d_model, epsilon=layer_norm_eps
+        )
+        self.norm_q = paddle.nn.LayerNorm(
+            normalized_shape=d_model, epsilon=layer_norm_eps
+        )
+        self.norm_c = paddle.nn.LayerNorm(
+            normalized_shape=d_model, epsilon=layer_norm_eps
+        )
+        self.dropout1 = paddle.nn.Dropout(p=dropout)
+        self.dropout2 = paddle.nn.Dropout(p=dropout)
+        self.dropout3 = paddle.nn.Dropout(p=dropout)
+        if activation == "gelu":
+            self.activation = paddle.nn.GELU()
+
+    def __setstate__(self, state):
+        if "activation" not in state:
+            state["activation"] = paddle.nn.functional.gelu
+        super().__setstate__(state)
+
+    def forward_stream(
+        self, tgt, tgt_norm, tgt_kv, memory, tgt_mask, tgt_key_padding_mask
+    ):
+        """Forward pass for a single stream (i.e. content or query)
+        tgt_norm is just a LayerNorm'd tgt. Added as a separate parameter for efficiency.
+        Both tgt_kv and memory are expected to be LayerNorm'd too.
+        memory is LayerNorm'd by ViT.
+        """
+        if tgt_key_padding_mask is not None:
+            tgt_mask1 = (tgt_mask != float("-inf"))[None, None, :, :] & (
+                tgt_key_padding_mask[:, None, None, :] == False
+            )
+            tgt2, sa_weights = self.self_attn(
+                tgt_norm, tgt_kv, tgt_kv, attn_mask=tgt_mask1
+            )
+        else:
+            tgt2, sa_weights = self.self_attn(
+                tgt_norm, tgt_kv, tgt_kv, attn_mask=tgt_mask
+            )
+
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2, ca_weights = self.cross_attn(self.norm1(tgt), memory, memory)
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.linear2(
+            self.dropout(self.activation(self.linear1(self.norm2(tgt))))
+        )
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt, sa_weights, ca_weights
+
+    def forward(
+        self,
+        query,
+        content,
+        memory,
+        query_mask=None,
+        content_mask=None,
+        content_key_padding_mask=None,
+        update_content=True,
+    ):
+        query_norm = self.norm_q(query)
+        content_norm = self.norm_c(content)
+        query = self.forward_stream(
+            query,
+            query_norm,
+            content_norm,
+            memory,
+            query_mask,
+            content_key_padding_mask,
+        )[0]
+        if update_content:
+            content = self.forward_stream(
+                content,
+                content_norm,
+                content_norm,
+                memory,
+                content_mask,
+                content_key_padding_mask,
+            )[0]
+        return query, content
+
+
+def get_clones(module, N):
+    return paddle.nn.LayerList([copy.deepcopy(module) for i in range(N)])
+
+
+class Decoder(paddle.nn.Layer):
+    __constants__ = ["norm"]
+
+    def __init__(self, decoder_layer, num_layers, norm):
+        super().__init__()
+        self.layers = get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        query,
+        content,
+        memory,
+        query_mask: Optional[paddle.Tensor] = None,
+        content_mask: Optional[paddle.Tensor] = None,
+        content_key_padding_mask: Optional[paddle.Tensor] = None,
+    ):
+        for i, mod in enumerate(self.layers):
+            last = i == len(self.layers) - 1
+            query, content = mod(
+                query,
+                content,
+                memory,
+                query_mask,
+                content_mask,
+                content_key_padding_mask,
+                update_content=not last,
+            )
+        query = self.norm(query)
+        return query
+
+
+class TokenEmbedding(paddle.nn.Layer):
+    def __init__(self, charset_size: int, embed_dim: int):
+        super().__init__()
+        self.embedding = paddle.nn.Embedding(
+            num_embeddings=charset_size, embedding_dim=embed_dim
+        )
+        self.embed_dim = embed_dim
+
+    def forward(self, tokens: paddle.Tensor):
+        return math.sqrt(self.embed_dim) * self.embedding(tokens.astype(paddle.int64))
+
+
+def trunc_normal_init(param, **kwargs):
+    initializer = nn.initializer.TruncatedNormal(**kwargs)
+    initializer(param, param.block)
+
+
+def constant_init(param, **kwargs):
+    initializer = nn.initializer.Constant(**kwargs)
+    initializer(param, param.block)
+
+
+def kaiming_normal_init(param, **kwargs):
+    initializer = nn.initializer.KaimingNormal(**kwargs)
+    initializer(param, param.block)
+
+
+class ParseQHead(nn.Layer):
+    def __init__(
+        self,
+        out_channels,
+        max_text_length,
+        embed_dim,
+        dec_num_heads,
+        dec_mlp_ratio,
+        dec_depth,
+        perm_num,
+        perm_forward,
+        perm_mirrored,
+        decode_ar,
+        refine_iters,
+        dropout,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.bos_id = out_channels - 2
+        self.eos_id = 0
+        self.pad_id = out_channels - 1
+
+        self.max_label_length = max_text_length
+        self.decode_ar = decode_ar
+        self.refine_iters = refine_iters
+        decoder_layer = DecoderLayer(
+            embed_dim, dec_num_heads, embed_dim * dec_mlp_ratio, dropout
+        )
+        self.decoder = Decoder(
+            decoder_layer,
+            num_layers=dec_depth,
+            norm=paddle.nn.LayerNorm(normalized_shape=embed_dim),
+        )
+        self.rng = np.random.default_rng()
+        self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num
+        self.perm_forward = perm_forward
+        self.perm_mirrored = perm_mirrored
+        self.head = paddle.nn.Linear(
+            in_features=embed_dim, out_features=out_channels - 2
+        )
+        self.text_embed = TokenEmbedding(out_channels, embed_dim)
+        self.pos_queries = paddle.create_parameter(
+            shape=paddle.empty(shape=[1, max_text_length + 1, embed_dim]).shape,
+            dtype=paddle.empty(shape=[1, max_text_length + 1, embed_dim]).numpy().dtype,
+            default_initializer=paddle.nn.initializer.Assign(
+                paddle.empty(shape=[1, max_text_length + 1, embed_dim])
+            ),
+        )
+        self.pos_queries.stop_gradient = not True
+        self.dropout = paddle.nn.Dropout(p=dropout)
+        self._device = self.parameters()[0].place
+        trunc_normal_init(self.pos_queries, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, paddle.nn.Linear):
+            trunc_normal_init(m.weight, std=0.02)
+            if m.bias is not None:
+                constant_init(m.bias, value=0.0)
+        elif isinstance(m, paddle.nn.Embedding):
+            trunc_normal_init(m.weight, std=0.02)
+            if m._padding_idx is not None:
+                m.weight.data[m._padding_idx].zero_()
+        elif isinstance(m, paddle.nn.Conv2D):
+            kaiming_normal_init(m.weight, fan_in=None, nonlinearity="relu")
+            if m.bias is not None:
+                constant_init(m.bias, value=0.0)
+        elif isinstance(
+            m, (paddle.nn.LayerNorm, paddle.nn.BatchNorm2D, paddle.nn.GroupNorm)
+        ):
+            constant_init(m.weight, value=1.0)
+            constant_init(m.bias, value=0.0)
+
+    def no_weight_decay(self):
+        param_names = {"text_embed.embedding.weight", "pos_queries"}
+        enc_param_names = {("encoder." + n) for n in self.encoder.no_weight_decay()}
+        return param_names.union(enc_param_names)
+
+    def encode(self, img):
+        return self.encoder(img)
+
+    def decode(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        tgt_padding_mask=None,
+        tgt_query=None,
+        tgt_query_mask=None,
+    ):
+        N, L = tgt.shape
+        null_ctx = self.text_embed(tgt[:, :1])
+        if L != 1:
+            tgt_emb = self.pos_queries[:, : L - 1] + self.text_embed(tgt[:, 1:])
+            tgt_emb = self.dropout(paddle.concat(x=[null_ctx, tgt_emb], axis=1))
+        else:
+            tgt_emb = self.dropout(null_ctx)
+        if tgt_query is None:
+            tgt_query = self.pos_queries[:, :L].expand(shape=[N, -1, -1])
+        tgt_query = self.dropout(tgt_query)
+        return self.decoder(
+            tgt_query, tgt_emb, memory, tgt_query_mask, tgt_mask, tgt_padding_mask
+        )
+
+    def forward_test(self, memory, max_length=None):
+        testing = max_length is None
+        max_length = (
+            self.max_label_length
+            if max_length is None
+            else min(max_length, self.max_label_length)
+        )
+        bs = memory.shape[0]
+        num_steps = max_length + 1
+
+        pos_queries = self.pos_queries[:, :num_steps].expand(shape=[bs, -1, -1])
+        tgt_mask = query_mask = paddle.triu(
+            x=paddle.full(shape=(num_steps, num_steps), fill_value=float("-inf")),
+            diagonal=1,
+        )
+        if self.decode_ar:
+            tgt_in = paddle.full(shape=(bs, num_steps), fill_value=self.pad_id).astype(
+                "int64"
+            )
+            tgt_in[:, (0)] = self.bos_id
+
+            logits = []
+            for i in range(paddle.to_tensor(num_steps)):
+                j = i + 1
+                tgt_out = self.decode(
+                    tgt_in[:, :j],
+                    memory,
+                    tgt_mask[:j, :j],
+                    tgt_query=pos_queries[:, i:j],
+                    tgt_query_mask=query_mask[i:j, :j],
+                )
+                p_i = self.head(tgt_out)
+                logits.append(p_i)
+                if j < num_steps:
+                    tgt_in[:, (j)] = p_i.squeeze().argmax(axis=-1)
+                    if (
+                        testing
+                        and (tgt_in == self.eos_id)
+                        .astype("bool")
+                        .any(axis=-1)
+                        .astype("bool")
+                        .all()
+                    ):
+                        break
+            logits = paddle.concat(x=logits, axis=1)
+        else:
+            tgt_in = paddle.full(shape=(bs, 1), fill_value=self.bos_id).astype("int64")
+            tgt_out = self.decode(tgt_in, memory, tgt_query=pos_queries)
+            logits = self.head(tgt_out)
+        if self.refine_iters:
+            temp = paddle.triu(
+                x=paddle.ones(shape=[num_steps, num_steps], dtype="bool"), diagonal=2
+            )
+            posi = np.where(temp.cpu().numpy() == True)
+            query_mask[posi] = 0
+            bos = paddle.full(shape=(bs, 1), fill_value=self.bos_id).astype("int64")
+            for i in range(self.refine_iters):
+                tgt_in = paddle.concat(x=[bos, logits[:, :-1].argmax(axis=-1)], axis=1)
+                tgt_padding_mask = (tgt_in == self.eos_id).astype(dtype="int32")
+                tgt_padding_mask = tgt_padding_mask.cpu()
+                tgt_padding_mask = tgt_padding_mask.cumsum(axis=-1) > 0
+                tgt_padding_mask = (
+                    tgt_padding_mask.cuda().astype(dtype="float32") == 1.0
+                )
+                tgt_out = self.decode(
+                    tgt_in,
+                    memory,
+                    tgt_mask,
+                    tgt_padding_mask,
+                    tgt_query=pos_queries,
+                    tgt_query_mask=query_mask[:, : tgt_in.shape[1]],
+                )
+                logits = self.head(tgt_out)
+
+        # transfer to probability
+        logits = F.softmax(logits, axis=-1)
+
+        final_output = {"predict": logits}
+
+        return final_output
+
+    def gen_tgt_perms(self, tgt):
+        """Generate shared permutations for the whole batch.
+        This works because the same attention mask can be used for the shorter sequences
+        because of the padding mask.
+        """
+        max_num_chars = tgt.shape[1] - 2
+        if max_num_chars == 1:
+            return paddle.arange(end=3).unsqueeze(axis=0)
+        perms = [paddle.arange(end=max_num_chars)] if self.perm_forward else []
+        max_perms = math.factorial(max_num_chars)
+        if self.perm_mirrored:
+            max_perms //= 2
+        num_gen_perms = min(self.max_gen_perms, max_perms)
+        if max_num_chars < 5:
+            if max_num_chars == 4 and self.perm_mirrored:
+                selector = [0, 3, 4, 6, 9, 10, 12, 16, 17, 18, 19, 21]
+            else:
+                selector = list(range(max_perms))
+            perm_pool = paddle.to_tensor(
+                data=list(permutations(range(max_num_chars), max_num_chars)),
+                place=self._device,
+            )[selector]
+            if self.perm_forward:
+                perm_pool = perm_pool[1:]
+            perms = paddle.stack(x=perms)
+            if len(perm_pool):
+                i = self.rng.choice(
+                    len(perm_pool), size=num_gen_perms - len(perms), replace=False
+                )
+                perms = paddle.concat(x=[perms, perm_pool[i]])
+        else:
+            perms.extend(
+                [
+                    paddle.randperm(n=max_num_chars)
+                    for _ in range(num_gen_perms - len(perms))
+                ]
+            )
+            perms = paddle.stack(x=perms)
+        if self.perm_mirrored:
+            comp = perms.flip(axis=-1)
+            x = paddle.stack(x=[perms, comp])
+            perm_2 = list(range(x.ndim))
+            perm_2[0] = 1
+            perm_2[1] = 0
+            perms = x.transpose(perm=perm_2).reshape((-1, max_num_chars))
+        bos_idx = paddle.zeros(shape=(len(perms), 1), dtype=perms.dtype)
+        eos_idx = paddle.full(
+            shape=(len(perms), 1), fill_value=max_num_chars + 1, dtype=perms.dtype
+        )
+        perms = paddle.concat(x=[bos_idx, perms + 1, eos_idx], axis=1)
+        if len(perms) > 1:
+            perms[(1), 1:] = max_num_chars + 1 - paddle.arange(end=max_num_chars + 1)
+        return perms
+
+    def generate_attn_masks(self, perm):
+        """Generate attention masks given a sequence permutation (includes pos. for bos and eos tokens)
+        :param perm: the permutation sequence. i = 0 is always the BOS
+        :return: lookahead attention masks
+        """
+        sz = perm.shape[0]
+        mask = paddle.zeros(shape=(sz, sz))
+        for i in range(sz):
+            query_idx = perm[i].cpu().numpy().tolist()
+            masked_keys = perm[i + 1 :].cpu().numpy().tolist()
+            if len(masked_keys) == 0:
+                break
+            mask[query_idx, masked_keys] = float("-inf")
+        content_mask = mask[:-1, :-1].clone()
+        mask[paddle.eye(num_rows=sz).astype("bool")] = float("-inf")
+        query_mask = mask[1:, :-1]
+        return content_mask, query_mask
+
+    def forward_train(self, memory, tgt):
+        tgt_perms = self.gen_tgt_perms(tgt)
+        tgt_in = tgt[:, :-1]
+        tgt_padding_mask = (tgt_in == self.pad_id) | (tgt_in == self.eos_id)
+        logits_list = []
+        final_out = {}
+        for i, perm in enumerate(tgt_perms):
+            tgt_mask, query_mask = self.generate_attn_masks(perm)
+            out = self.decode(
+                tgt_in, memory, tgt_mask, tgt_padding_mask, tgt_query_mask=query_mask
+            )
+            logits = self.head(out)
+            if i == 0:
+                final_out["predict"] = logits
+            logits = logits.flatten(stop_axis=1)
+            logits_list.append(logits)
+
+        final_out["logits_list"] = logits_list
+        final_out["pad_id"] = self.pad_id
+        final_out["eos_id"] = self.eos_id
+
+        return final_out
+
+    def forward(self, feat, targets=None):
+        # feat : B, N, C
+        # targets : labels, labels_len
+
+        if self.training:
+            label = targets[0]  # label
+            label_len = targets[1]
+            max_step = paddle.max(label_len).cpu().numpy()[0] + 2
+            crop_label = label[:, :max_step]
+            final_out = self.forward_train(feat, crop_label)
+        else:
+            final_out = self.forward_test(feat)
+
+        return final_out

ppocr/modeling/heads/rec_pren_head.py

@@ -0,0 +1,34 @@
+# 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
+
+from paddle import nn
+from paddle.nn import functional as F
+
+
+class PRENHead(nn.Layer):
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(PRENHead, self).__init__()
+        self.linear = nn.Linear(in_channels, out_channels)
+
+    def forward(self, x, targets=None):
+        predicts = self.linear(x)
+
+        if not self.training:
+            predicts = F.softmax(predicts, axis=2)
+
+        return predicts

ppocr/modeling/heads/rec_rfl_head.py

@@ -0,0 +1,106 @@
+# 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/hikopensource/DAVAR-Lab-OCR/blob/main/davarocr/davar_rcg/models/sequence_heads/counting_head.py
+"""
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal, KaimingNormal
+
+from .rec_att_head import AttentionLSTM
+
+kaiming_init_ = KaimingNormal()
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+class CNTHead(nn.Layer):
+    def __init__(self, embed_size=512, encode_length=26, out_channels=38, **kwargs):
+        super(CNTHead, self).__init__()
+
+        self.out_channels = out_channels
+
+        self.Wv_fusion = nn.Linear(embed_size, embed_size, bias_attr=False)
+        self.Prediction_visual = nn.Linear(
+            encode_length * embed_size, self.out_channels
+        )
+
+    def forward(self, visual_feature):
+        b, c, h, w = visual_feature.shape
+        visual_feature = visual_feature.reshape([b, c, h * w]).transpose([0, 2, 1])
+        visual_feature_num = self.Wv_fusion(visual_feature)  # batch * 26 * 512
+        b, n, c = visual_feature_num.shape
+        # using visual feature directly calculate the text length
+        visual_feature_num = visual_feature_num.reshape([b, n * c])
+        prediction_visual = self.Prediction_visual(visual_feature_num)
+
+        return prediction_visual
+
+
+class RFLHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels=512,
+        hidden_size=256,
+        batch_max_legnth=25,
+        out_channels=38,
+        use_cnt=True,
+        use_seq=True,
+        **kwargs,
+    ):
+        super(RFLHead, self).__init__()
+        assert use_cnt or use_seq
+        self.use_cnt = use_cnt
+        self.use_seq = use_seq
+        if self.use_cnt:
+            self.cnt_head = CNTHead(
+                embed_size=in_channels,
+                encode_length=batch_max_legnth + 1,
+                out_channels=out_channels,
+                **kwargs,
+            )
+        if self.use_seq:
+            self.seq_head = AttentionLSTM(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                hidden_size=hidden_size,
+                **kwargs,
+            )
+        self.batch_max_legnth = batch_max_legnth
+        self.num_class = out_channels
+        self.apply(self.init_weights)
+
+    def init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            kaiming_init_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+
+    def forward(self, x, targets=None):
+        cnt_inputs, seq_inputs = x
+        if self.use_cnt:
+            cnt_outputs = self.cnt_head(cnt_inputs)
+        else:
+            cnt_outputs = None
+        if self.use_seq:
+            if self.training:
+                seq_outputs = self.seq_head(
+                    seq_inputs, targets[0], self.batch_max_legnth
+                )
+            else:
+                seq_outputs = self.seq_head(seq_inputs, None, self.batch_max_legnth)
+            return cnt_outputs, seq_outputs
+        else:
+            return cnt_outputs

ppocr/modeling/heads/rec_robustscanner_head.py

@@ -0,0 +1,748 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/encoders/channel_reduction_encoder.py
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/decoders/robust_scanner_decoder.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class BaseDecoder(nn.Layer):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def forward_train(self, feat, out_enc, targets, img_metas):
+        raise NotImplementedError
+
+    def forward_test(self, feat, out_enc, img_metas):
+        raise NotImplementedError
+
+    def forward(
+        self,
+        feat,
+        out_enc,
+        label=None,
+        valid_ratios=None,
+        word_positions=None,
+        train_mode=True,
+    ):
+        self.train_mode = train_mode
+
+        if train_mode:
+            return self.forward_train(
+                feat, out_enc, label, valid_ratios, word_positions
+            )
+        return self.forward_test(feat, out_enc, valid_ratios, word_positions)
+
+
+class ChannelReductionEncoder(nn.Layer):
+    """Change the channel number with a one by one convoluational layer.
+
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+    """
+
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(ChannelReductionEncoder, self).__init__()
+
+        self.layer = nn.Conv2D(
+            in_channels,
+            out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            weight_attr=nn.initializer.XavierNormal(),
+        )
+
+    def forward(self, feat):
+        """
+        Args:
+            feat (Tensor): Image features with the shape of
+                :math:`(N, C_{in}, H, W)`.
+
+        Returns:
+            Tensor: A tensor of shape :math:`(N, C_{out}, H, W)`.
+        """
+        return self.layer(feat)
+
+
+def masked_fill(x, mask, value):
+    y = paddle.full(x.shape, value, x.dtype)
+    return paddle.where(mask, y, x)
+
+
+class DotProductAttentionLayer(nn.Layer):
+    def __init__(self, dim_model=None):
+        super().__init__()
+
+        self.scale = dim_model**-0.5 if dim_model is not None else 1.0
+
+    def forward(self, query, key, value, h, w, valid_ratios=None):
+        query = paddle.transpose(query, (0, 2, 1))
+        logits = paddle.matmul(query, key) * self.scale
+        n, c, t = logits.shape
+        # reshape to (n, c, h, w)
+        logits = paddle.reshape(logits, [n, c, h, w])
+        if valid_ratios is not None:
+            # cal mask of attention weight
+            with paddle.base.framework._stride_in_no_check_dy2st_diff():
+                for i, valid_ratio in enumerate(valid_ratios):
+                    valid_width = min(w, int(w * valid_ratio + 0.5))
+                    if valid_width < w:
+                        logits[i, :, :, valid_width:] = float("-inf")
+
+        # reshape to (n, c, h, w)
+        logits = paddle.reshape(logits, [n, c, t])
+        weights = F.softmax(logits, axis=2)
+        value = paddle.transpose(value, (0, 2, 1))
+        glimpse = paddle.matmul(weights, value)
+        glimpse = paddle.transpose(glimpse, (0, 2, 1))
+        return glimpse
+
+
+class SequenceAttentionDecoder(BaseDecoder):
+    """Sequence attention decoder for RobustScanner.
+
+    RobustScanner: `RobustScanner: Dynamically Enhancing Positional Clues for
+    Robust Text Recognition <https://arxiv.org/abs/2007.07542>`_
+
+    Args:
+        num_classes (int): Number of output classes :math:`C`.
+        rnn_layers (int): Number of RNN layers.
+        dim_input (int): Dimension :math:`D_i` of input vector ``feat``.
+        dim_model (int): Dimension :math:`D_m` of the model. Should also be the
+            same as encoder output vector ``out_enc``.
+        max_seq_len (int): Maximum output sequence length :math:`T`.
+        start_idx (int): The index of `<SOS>`.
+        mask (bool): Whether to mask input features according to
+            ``img_meta['valid_ratio']``.
+        padding_idx (int): The index of `<PAD>`.
+        dropout (float): Dropout rate.
+        return_feature (bool): Return feature or logits as the result.
+        encode_value (bool): Whether to use the output of encoder ``out_enc``
+            as `value` of attention layer. If False, the original feature
+            ``feat`` will be used.
+
+    Warning:
+        This decoder will not predict the final class which is assumed to be
+        `<PAD>`. Therefore, its output size is always :math:`C - 1`. `<PAD>`
+        is also ignored by loss as specified in
+        :obj:`mmocr.models.textrecog.recognizer.EncodeDecodeRecognizer`.
+    """
+
+    def __init__(
+        self,
+        num_classes=None,
+        rnn_layers=2,
+        dim_input=512,
+        dim_model=128,
+        max_seq_len=40,
+        start_idx=0,
+        mask=True,
+        padding_idx=None,
+        dropout=0,
+        return_feature=False,
+        encode_value=False,
+    ):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.dim_input = dim_input
+        self.dim_model = dim_model
+        self.return_feature = return_feature
+        self.encode_value = encode_value
+        self.max_seq_len = max_seq_len
+        self.start_idx = start_idx
+        self.mask = mask
+
+        self.embedding = nn.Embedding(
+            self.num_classes, self.dim_model, padding_idx=padding_idx
+        )
+
+        self.sequence_layer = nn.LSTM(
+            input_size=dim_model,
+            hidden_size=dim_model,
+            num_layers=rnn_layers,
+            time_major=False,
+            dropout=dropout,
+        )
+
+        self.attention_layer = DotProductAttentionLayer()
+
+        self.prediction = None
+        if not self.return_feature:
+            pred_num_classes = num_classes - 1
+            self.prediction = nn.Linear(
+                dim_model if encode_value else dim_input, pred_num_classes
+            )
+
+    def forward_train(self, feat, out_enc, targets, valid_ratios):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            targets (Tensor): a tensor of shape :math:`(N, T)`. Each element is the index of a
+                character.
+            valid_ratios (Tensor): valid length ratio of img.
+        Returns:
+            Tensor: A raw logit tensor of shape :math:`(N, T, C-1)` if
+            ``return_feature=False``. Otherwise it would be the hidden feature
+            before the prediction projection layer, whose shape is
+            :math:`(N, T, D_m)`.
+        """
+
+        tgt_embedding = self.embedding(targets)
+
+        n, c_enc, h, w = out_enc.shape
+        assert c_enc == self.dim_model
+        _, c_feat, _, _ = feat.shape
+        assert c_feat == self.dim_input
+        _, len_q, c_q = tgt_embedding.shape
+        assert c_q == self.dim_model
+        assert len_q <= self.max_seq_len
+
+        query, _ = self.sequence_layer(tgt_embedding)
+        query = paddle.transpose(query, (0, 2, 1))
+        key = paddle.reshape(out_enc, [n, c_enc, h * w])
+        if self.encode_value:
+            value = key
+        else:
+            value = paddle.reshape(feat, [n, c_feat, h * w])
+
+        attn_out = self.attention_layer(query, key, value, h, w, valid_ratios)
+        attn_out = paddle.transpose(attn_out, (0, 2, 1))
+
+        if self.return_feature:
+            return attn_out
+
+        out = self.prediction(attn_out)
+
+        return out
+
+    def forward_test(self, feat, out_enc, valid_ratios):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            valid_ratios (Tensor): valid length ratio of img.
+
+        Returns:
+            Tensor: The output logit sequence tensor of shape
+            :math:`(N, T, C-1)`.
+        """
+        seq_len = self.max_seq_len
+        batch_size = feat.shape[0]
+
+        decode_sequence = (
+            paddle.ones((batch_size, seq_len), dtype="int64") * self.start_idx
+        )
+
+        outputs = []
+        for i in range(seq_len):
+            step_out = self.forward_test_step(
+                feat, out_enc, decode_sequence, i, valid_ratios
+            )
+            outputs.append(step_out)
+            max_idx = paddle.argmax(step_out, axis=1, keepdim=False)
+            if i < seq_len - 1:
+                decode_sequence[:, i + 1] = max_idx
+
+        outputs = paddle.stack(outputs, 1)
+
+        return outputs
+
+    def forward_test_step(
+        self, feat, out_enc, decode_sequence, current_step, valid_ratios
+    ):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            decode_sequence (Tensor): Shape :math:`(N, T)`. The tensor that
+                stores history decoding result.
+            current_step (int): Current decoding step.
+            valid_ratios (Tensor): valid length ratio of img
+
+        Returns:
+            Tensor: Shape :math:`(N, C-1)`. The logit tensor of predicted
+            tokens at current time step.
+        """
+
+        embed = self.embedding(decode_sequence)
+
+        n, c_enc, h, w = out_enc.shape
+        assert c_enc == self.dim_model
+        _, c_feat, _, _ = feat.shape
+        assert c_feat == self.dim_input
+        _, _, c_q = embed.shape
+        assert c_q == self.dim_model
+
+        query, _ = self.sequence_layer(embed)
+        query = paddle.transpose(query, (0, 2, 1))
+        key = paddle.reshape(out_enc, [n, c_enc, h * w])
+        if self.encode_value:
+            value = key
+        else:
+            value = paddle.reshape(feat, [n, c_feat, h * w])
+
+        # [n, c, l]
+        attn_out = self.attention_layer(query, key, value, h, w, valid_ratios)
+        out = attn_out[:, :, current_step]
+
+        if self.return_feature:
+            return out
+
+        out = self.prediction(out)
+        out = F.softmax(out, dim=-1)
+
+        return out
+
+
+class PositionAwareLayer(nn.Layer):
+    def __init__(self, dim_model, rnn_layers=2):
+        super().__init__()
+
+        self.dim_model = dim_model
+
+        self.rnn = nn.LSTM(
+            input_size=dim_model,
+            hidden_size=dim_model,
+            num_layers=rnn_layers,
+            time_major=False,
+        )
+
+        self.mixer = nn.Sequential(
+            nn.Conv2D(dim_model, dim_model, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2D(dim_model, dim_model, kernel_size=3, stride=1, padding=1),
+        )
+
+    def forward(self, img_feature):
+        n, c, h, w = img_feature.shape
+        rnn_input = paddle.transpose(img_feature, (0, 2, 3, 1))
+        rnn_input = paddle.reshape(rnn_input, (n * h, w, c))
+        rnn_output, _ = self.rnn(rnn_input)
+        rnn_output = paddle.reshape(rnn_output, (n, h, w, c))
+        rnn_output = paddle.transpose(rnn_output, (0, 3, 1, 2))
+        out = self.mixer(rnn_output)
+        return out
+
+
+class PositionAttentionDecoder(BaseDecoder):
+    """Position attention decoder for RobustScanner.
+
+    RobustScanner: `RobustScanner: Dynamically Enhancing Positional Clues for
+    Robust Text Recognition <https://arxiv.org/abs/2007.07542>`_
+
+    Args:
+        num_classes (int): Number of output classes :math:`C`.
+        rnn_layers (int): Number of RNN layers.
+        dim_input (int): Dimension :math:`D_i` of input vector ``feat``.
+        dim_model (int): Dimension :math:`D_m` of the model. Should also be the
+            same as encoder output vector ``out_enc``.
+        max_seq_len (int): Maximum output sequence length :math:`T`.
+        mask (bool): Whether to mask input features according to
+            ``img_meta['valid_ratio']``.
+        return_feature (bool): Return feature or logits as the result.
+        encode_value (bool): Whether to use the output of encoder ``out_enc``
+            as `value` of attention layer. If False, the original feature
+            ``feat`` will be used.
+
+    Warning:
+        This decoder will not predict the final class which is assumed to be
+        `<PAD>`. Therefore, its output size is always :math:`C - 1`. `<PAD>`
+        is also ignored by loss
+
+    """
+
+    def __init__(
+        self,
+        num_classes=None,
+        rnn_layers=2,
+        dim_input=512,
+        dim_model=128,
+        max_seq_len=40,
+        mask=True,
+        return_feature=False,
+        encode_value=False,
+    ):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.dim_input = dim_input
+        self.dim_model = dim_model
+        self.max_seq_len = max_seq_len
+        self.return_feature = return_feature
+        self.encode_value = encode_value
+        self.mask = mask
+
+        self.embedding = nn.Embedding(self.max_seq_len + 1, self.dim_model)
+
+        self.position_aware_module = PositionAwareLayer(self.dim_model, rnn_layers)
+
+        self.attention_layer = DotProductAttentionLayer()
+
+        self.prediction = None
+        if not self.return_feature:
+            pred_num_classes = num_classes - 1
+            self.prediction = nn.Linear(
+                dim_model if encode_value else dim_input, pred_num_classes
+            )
+
+    def _get_position_index(self, length, batch_size):
+        position_index_list = []
+        for i in range(batch_size):
+            position_index = paddle.arange(0, end=length, step=1, dtype="int64")
+            position_index_list.append(position_index)
+        batch_position_index = paddle.stack(position_index_list, axis=0)
+        return batch_position_index
+
+    def forward_train(self, feat, out_enc, targets, valid_ratios, position_index):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            targets (dict): A dict with the key ``padded_targets``, a
+                tensor of shape :math:`(N, T)`. Each element is the index of a
+                character.
+            valid_ratios (Tensor): valid length ratio of img.
+            position_index (Tensor): The position of each word.
+
+        Returns:
+            Tensor: A raw logit tensor of shape :math:`(N, T, C-1)` if
+            ``return_feature=False``. Otherwise it will be the hidden feature
+            before the prediction projection layer, whose shape is
+            :math:`(N, T, D_m)`.
+        """
+        n, c_enc, h, w = out_enc.shape
+        assert c_enc == self.dim_model
+        _, c_feat, _, _ = feat.shape
+        assert c_feat == self.dim_input
+        _, len_q = targets.shape
+        assert len_q <= self.max_seq_len
+
+        position_out_enc = self.position_aware_module(out_enc)
+
+        query = self.embedding(position_index)
+        query = paddle.transpose(query, (0, 2, 1))
+        key = paddle.reshape(position_out_enc, (n, c_enc, h * w))
+        if self.encode_value:
+            value = paddle.reshape(out_enc, (n, c_enc, h * w))
+        else:
+            value = paddle.reshape(feat, (n, c_feat, h * w))
+
+        attn_out = self.attention_layer(query, key, value, h, w, valid_ratios)
+        attn_out = paddle.transpose(attn_out, (0, 2, 1))  # [n, len_q, dim_v]
+
+        if self.return_feature:
+            return attn_out
+
+        return self.prediction(attn_out)
+
+    def forward_test(self, feat, out_enc, valid_ratios, position_index):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            valid_ratios (Tensor): valid length ratio of img
+            position_index (Tensor): The position of each word.
+
+        Returns:
+            Tensor: A raw logit tensor of shape :math:`(N, T, C-1)` if
+            ``return_feature=False``. Otherwise it would be the hidden feature
+            before the prediction projection layer, whose shape is
+            :math:`(N, T, D_m)`.
+        """
+        n, c_enc, h, w = out_enc.shape
+        assert c_enc == self.dim_model
+        _, c_feat, _, _ = feat.shape
+        assert c_feat == self.dim_input
+
+        position_out_enc = self.position_aware_module(out_enc)
+
+        query = self.embedding(position_index)
+        query = paddle.transpose(query, (0, 2, 1))
+        key = paddle.reshape(position_out_enc, (n, c_enc, h * w))
+        if self.encode_value:
+            value = paddle.reshape(out_enc, (n, c_enc, h * w))
+        else:
+            value = paddle.reshape(feat, (n, c_feat, h * w))
+
+        attn_out = self.attention_layer(query, key, value, h, w, valid_ratios)
+        attn_out = paddle.transpose(attn_out, (0, 2, 1))  # [n, len_q, dim_v]
+
+        if self.return_feature:
+            return attn_out
+
+        return self.prediction(attn_out)
+
+
+class RobustScannerFusionLayer(nn.Layer):
+    def __init__(self, dim_model, dim=-1):
+        super(RobustScannerFusionLayer, self).__init__()
+
+        self.dim_model = dim_model
+        self.dim = dim
+        self.linear_layer = nn.Linear(dim_model * 2, dim_model * 2)
+
+    def forward(self, x0, x1):
+        assert x0.shape == x1.shape
+        fusion_input = paddle.concat([x0, x1], self.dim)
+        output = self.linear_layer(fusion_input)
+        output = F.glu(output, self.dim)
+        return output
+
+
+class RobustScannerDecoder(BaseDecoder):
+    """Decoder for RobustScanner.
+
+    RobustScanner: `RobustScanner: Dynamically Enhancing Positional Clues for
+    Robust Text Recognition <https://arxiv.org/abs/2007.07542>`_
+
+    Args:
+        num_classes (int): Number of output classes :math:`C`.
+        dim_input (int): Dimension :math:`D_i` of input vector ``feat``.
+        dim_model (int): Dimension :math:`D_m` of the model. Should also be the
+            same as encoder output vector ``out_enc``.
+        max_seq_len (int): Maximum output sequence length :math:`T`.
+        start_idx (int): The index of `<SOS>`.
+        mask (bool): Whether to mask input features according to
+            ``img_meta['valid_ratio']``.
+        padding_idx (int): The index of `<PAD>`.
+        encode_value (bool): Whether to use the output of encoder ``out_enc``
+            as `value` of attention layer. If False, the original feature
+            ``feat`` will be used.
+
+    Warning:
+        This decoder will not predict the final class which is assumed to be
+        `<PAD>`. Therefore, its output size is always :math:`C - 1`. `<PAD>`
+        is also ignored by loss as specified in
+        :obj:`mmocr.models.textrecog.recognizer.EncodeDecodeRecognizer`.
+    """
+
+    def __init__(
+        self,
+        num_classes=None,
+        dim_input=512,
+        dim_model=128,
+        hybrid_decoder_rnn_layers=2,
+        hybrid_decoder_dropout=0,
+        position_decoder_rnn_layers=2,
+        max_seq_len=40,
+        start_idx=0,
+        mask=True,
+        padding_idx=None,
+        encode_value=False,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.dim_input = dim_input
+        self.dim_model = dim_model
+        self.max_seq_len = max_seq_len
+        self.encode_value = encode_value
+        self.start_idx = start_idx
+        self.padding_idx = padding_idx
+        self.mask = mask
+
+        # init hybrid decoder
+        self.hybrid_decoder = SequenceAttentionDecoder(
+            num_classes=num_classes,
+            rnn_layers=hybrid_decoder_rnn_layers,
+            dim_input=dim_input,
+            dim_model=dim_model,
+            max_seq_len=max_seq_len,
+            start_idx=start_idx,
+            mask=mask,
+            padding_idx=padding_idx,
+            dropout=hybrid_decoder_dropout,
+            encode_value=encode_value,
+            return_feature=True,
+        )
+
+        # init position decoder
+        self.position_decoder = PositionAttentionDecoder(
+            num_classes=num_classes,
+            rnn_layers=position_decoder_rnn_layers,
+            dim_input=dim_input,
+            dim_model=dim_model,
+            max_seq_len=max_seq_len,
+            mask=mask,
+            encode_value=encode_value,
+            return_feature=True,
+        )
+
+        self.fusion_module = RobustScannerFusionLayer(
+            self.dim_model if encode_value else dim_input
+        )
+
+        pred_num_classes = num_classes - 1
+        self.prediction = nn.Linear(
+            dim_model if encode_value else dim_input, pred_num_classes
+        )
+
+    def forward_train(self, feat, out_enc, target, valid_ratios, word_positions):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            target (dict): A dict with the key ``padded_targets``, a
+                tensor of shape :math:`(N, T)`. Each element is the index of a
+                character.
+            valid_ratios (Tensor):
+            word_positions (Tensor): The position of each word.
+
+        Returns:
+            Tensor: A raw logit tensor of shape :math:`(N, T, C-1)`.
+        """
+        hybrid_glimpse = self.hybrid_decoder.forward_train(
+            feat, out_enc, target, valid_ratios
+        )
+        position_glimpse = self.position_decoder.forward_train(
+            feat, out_enc, target, valid_ratios, word_positions
+        )
+
+        fusion_out = self.fusion_module(hybrid_glimpse, position_glimpse)
+
+        out = self.prediction(fusion_out)
+
+        return out
+
+    def forward_test(self, feat, out_enc, valid_ratios, word_positions):
+        """
+        Args:
+            feat (Tensor): Tensor of shape :math:`(N, D_i, H, W)`.
+            out_enc (Tensor): Encoder output of shape
+                :math:`(N, D_m, H, W)`.
+            valid_ratios (Tensor):
+            word_positions (Tensor): The position of each word.
+        Returns:
+            Tensor: The output logit sequence tensor of shape
+            :math:`(N, T, C-1)`.
+        """
+        seq_len = self.max_seq_len
+        batch_size = feat.shape[0]
+
+        decode_sequence = (
+            paddle.ones((batch_size, seq_len), dtype="int64") * self.start_idx
+        )
+
+        position_glimpse = self.position_decoder.forward_test(
+            feat, out_enc, valid_ratios, word_positions
+        )
+
+        outputs = []
+        for i in range(seq_len):
+            hybrid_glimpse_step = self.hybrid_decoder.forward_test_step(
+                feat, out_enc, decode_sequence, i, valid_ratios
+            )
+
+            fusion_out = self.fusion_module(
+                hybrid_glimpse_step, position_glimpse[:, i, :]
+            )
+
+            char_out = self.prediction(fusion_out)
+            char_out = F.softmax(char_out, -1)
+            outputs.append(char_out)
+            max_idx = paddle.argmax(char_out, axis=1, keepdim=False)
+            if i < seq_len - 1:
+                decode_sequence[:, i + 1] = max_idx
+
+        outputs = paddle.stack(outputs, 1)
+
+        return outputs
+
+
+class RobustScannerHead(nn.Layer):
+    def __init__(
+        self,
+        out_channels,  # 90 + unknown + start + padding
+        in_channels,
+        enc_outchannles=128,
+        hybrid_dec_rnn_layers=2,
+        hybrid_dec_dropout=0,
+        position_dec_rnn_layers=2,
+        start_idx=0,
+        max_text_length=40,
+        mask=True,
+        padding_idx=None,
+        encode_value=False,
+        **kwargs,
+    ):
+        super(RobustScannerHead, self).__init__()
+
+        # encoder module
+        self.encoder = ChannelReductionEncoder(
+            in_channels=in_channels, out_channels=enc_outchannles
+        )
+
+        # decoder module
+        self.decoder = RobustScannerDecoder(
+            num_classes=out_channels,
+            dim_input=in_channels,
+            dim_model=enc_outchannles,
+            hybrid_decoder_rnn_layers=hybrid_dec_rnn_layers,
+            hybrid_decoder_dropout=hybrid_dec_dropout,
+            position_decoder_rnn_layers=position_dec_rnn_layers,
+            max_seq_len=max_text_length,
+            start_idx=start_idx,
+            mask=mask,
+            padding_idx=padding_idx,
+            encode_value=encode_value,
+        )
+
+    def forward(self, inputs, targets=None):
+        """
+        targets: [label, valid_ratio, word_positions]
+        """
+        out_enc = self.encoder(inputs)
+        valid_ratios = None
+        word_positions = targets[-1]
+
+        if len(targets) > 1:
+            valid_ratios = targets[-2]
+
+        if self.training:
+            label = targets[0]  # label
+            label = paddle.to_tensor(label, dtype="int64")
+            final_out = self.decoder(
+                inputs, out_enc, label, valid_ratios, word_positions
+            )
+        if not self.training:
+            final_out = self.decoder(
+                inputs,
+                out_enc,
+                label=None,
+                valid_ratios=valid_ratios,
+                word_positions=word_positions,
+                train_mode=False,
+            )
+        return final_out

ppocr/modeling/heads/rec_sar_head.py

@@ -0,0 +1,407 @@
+# copyright (c) 2021 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/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/encoders/sar_encoder.py
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/decoders/sar_decoder.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class SAREncoder(nn.Layer):
+    """
+    Args:
+        enc_bi_rnn (bool): If True, use bidirectional RNN in encoder.
+        enc_drop_rnn (float): Dropout probability of RNN layer in encoder.
+        enc_gru (bool): If True, use GRU, else LSTM in encoder.
+        d_model (int): Dim of channels from backbone.
+        d_enc (int): Dim of encoder RNN layer.
+        mask (bool): If True, mask padding in RNN sequence.
+    """
+
+    def __init__(
+        self,
+        enc_bi_rnn=False,
+        enc_drop_rnn=0.1,
+        enc_gru=False,
+        d_model=512,
+        d_enc=512,
+        mask=True,
+        **kwargs,
+    ):
+        super().__init__()
+        assert isinstance(enc_bi_rnn, bool)
+        assert isinstance(enc_drop_rnn, (int, float))
+        assert 0 <= enc_drop_rnn < 1.0
+        assert isinstance(enc_gru, bool)
+        assert isinstance(d_model, int)
+        assert isinstance(d_enc, int)
+        assert isinstance(mask, bool)
+
+        self.enc_bi_rnn = enc_bi_rnn
+        self.enc_drop_rnn = enc_drop_rnn
+        self.mask = mask
+
+        # LSTM Encoder
+        if enc_bi_rnn:
+            direction = "bidirectional"
+        else:
+            direction = "forward"
+        kwargs = dict(
+            input_size=d_model,
+            hidden_size=d_enc,
+            num_layers=2,
+            time_major=False,
+            dropout=enc_drop_rnn,
+            direction=direction,
+        )
+        if enc_gru:
+            self.rnn_encoder = nn.GRU(**kwargs)
+        else:
+            self.rnn_encoder = nn.LSTM(**kwargs)
+
+        # global feature transformation
+        encoder_rnn_out_size = d_enc * (int(enc_bi_rnn) + 1)
+        self.linear = nn.Linear(encoder_rnn_out_size, encoder_rnn_out_size)
+
+    def forward(self, feat, img_metas=None):
+        if img_metas is not None:
+            assert len(img_metas[0]) == feat.shape[0]
+
+        valid_ratios = None
+        if img_metas is not None and self.mask:
+            valid_ratios = img_metas[-1]
+
+        h_feat = feat.shape[2]  # bsz c h w
+        feat_v = F.max_pool2d(feat, kernel_size=(h_feat, 1), stride=1, padding=0)
+        feat_v = feat_v.squeeze(2)  # bsz * C * W
+        feat_v = paddle.transpose(feat_v, perm=[0, 2, 1])  # bsz * W * C
+        holistic_feat = self.rnn_encoder(feat_v)[0]  # bsz * T * C
+
+        if valid_ratios is not None:
+            valid_hf = []
+            T = paddle.shape(holistic_feat)[1]
+            for i in range(valid_ratios.shape[0]):
+                valid_step = (
+                    paddle.minimum(T, paddle.ceil(valid_ratios[i] * T).astype(T.dtype))
+                    - 1
+                )
+                valid_hf.append(holistic_feat[i, valid_step, :])
+            valid_hf = paddle.stack(valid_hf, axis=0)
+        else:
+            valid_hf = holistic_feat[:, -1, :]  # bsz * C
+        holistic_feat = self.linear(valid_hf)  # bsz * C
+
+        return holistic_feat
+
+
+class BaseDecoder(nn.Layer):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+    def forward_train(self, feat, out_enc, targets, img_metas):
+        raise NotImplementedError
+
+    def forward_test(self, feat, out_enc, img_metas):
+        raise NotImplementedError
+
+    def forward(self, feat, out_enc, label=None, img_metas=None, train_mode=True):
+        self.train_mode = train_mode
+
+        if train_mode:
+            return self.forward_train(feat, out_enc, label, img_metas)
+        return self.forward_test(feat, out_enc, img_metas)
+
+
+class ParallelSARDecoder(BaseDecoder):
+    """
+    Args:
+        out_channels (int): Output class number.
+        enc_bi_rnn (bool): If True, use bidirectional RNN in encoder.
+        dec_bi_rnn (bool): If True, use bidirectional RNN in decoder.
+        dec_drop_rnn (float): Dropout of RNN layer in decoder.
+        dec_gru (bool): If True, use GRU, else LSTM in decoder.
+        d_model (int): Dim of channels from backbone.
+        d_enc (int): Dim of encoder RNN layer.
+        d_k (int): Dim of channels of attention module.
+        pred_dropout (float): Dropout probability of prediction layer.
+        max_seq_len (int): Maximum sequence length for decoding.
+        mask (bool): If True, mask padding in feature map.
+        start_idx (int): Index of start token.
+        padding_idx (int): Index of padding token.
+        pred_concat (bool): If True, concat glimpse feature from
+            attention with holistic feature and hidden state.
+    """
+
+    def __init__(
+        self,
+        out_channels,  # 90 + unknown + start + padding
+        enc_bi_rnn=False,
+        dec_bi_rnn=False,
+        dec_drop_rnn=0.0,
+        dec_gru=False,
+        d_model=512,
+        d_enc=512,
+        d_k=64,
+        pred_dropout=0.1,
+        max_text_length=30,
+        mask=True,
+        pred_concat=True,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.num_classes = out_channels
+        self.enc_bi_rnn = enc_bi_rnn
+        self.d_k = d_k
+        self.start_idx = out_channels - 2
+        self.padding_idx = out_channels - 1
+        self.max_seq_len = max_text_length
+        self.mask = mask
+        self.pred_concat = pred_concat
+
+        encoder_rnn_out_size = d_enc * (int(enc_bi_rnn) + 1)
+        decoder_rnn_out_size = encoder_rnn_out_size * (int(dec_bi_rnn) + 1)
+
+        # 2D attention layer
+        self.conv1x1_1 = nn.Linear(decoder_rnn_out_size, d_k)
+        self.conv3x3_1 = nn.Conv2D(d_model, d_k, kernel_size=3, stride=1, padding=1)
+        self.conv1x1_2 = nn.Linear(d_k, 1)
+
+        # Decoder RNN layer
+        if dec_bi_rnn:
+            direction = "bidirectional"
+        else:
+            direction = "forward"
+
+        kwargs = dict(
+            input_size=encoder_rnn_out_size,
+            hidden_size=encoder_rnn_out_size,
+            num_layers=2,
+            time_major=False,
+            dropout=dec_drop_rnn,
+            direction=direction,
+        )
+        if dec_gru:
+            self.rnn_decoder = nn.GRU(**kwargs)
+        else:
+            self.rnn_decoder = nn.LSTM(**kwargs)
+
+        # Decoder input embedding
+        self.embedding = nn.Embedding(
+            self.num_classes, encoder_rnn_out_size, padding_idx=self.padding_idx
+        )
+
+        # Prediction layer
+        self.pred_dropout = nn.Dropout(pred_dropout)
+        pred_num_classes = self.num_classes - 1
+        if pred_concat:
+            fc_in_channel = decoder_rnn_out_size + d_model + encoder_rnn_out_size
+        else:
+            fc_in_channel = d_model
+        self.prediction = nn.Linear(fc_in_channel, pred_num_classes)
+
+    def _2d_attention(self, decoder_input, feat, holistic_feat, valid_ratios=None):
+        y = self.rnn_decoder(decoder_input)[0]
+        # y: bsz * (seq_len + 1) * hidden_size
+
+        attn_query = self.conv1x1_1(y)  # bsz * (seq_len + 1) * attn_size
+        bsz, seq_len, attn_size = attn_query.shape
+        attn_query = paddle.unsqueeze(attn_query, axis=[3, 4])
+        # (bsz, seq_len + 1, attn_size, 1, 1)
+
+        attn_key = self.conv3x3_1(feat)
+        # bsz * attn_size * h * w
+        attn_key = attn_key.unsqueeze(1)
+        # bsz * 1 * attn_size * h * w
+
+        attn_weight = paddle.tanh(paddle.add(attn_key, attn_query))
+
+        # bsz * (seq_len + 1) * attn_size * h * w
+        attn_weight = paddle.transpose(attn_weight, perm=[0, 1, 3, 4, 2])
+        # bsz * (seq_len + 1) * h * w * attn_size
+        attn_weight = self.conv1x1_2(attn_weight)
+        # bsz * (seq_len + 1) * h * w * 1
+        bsz, T, h, w, c = paddle.shape(attn_weight)
+        assert c == 1
+
+        if valid_ratios is not None:
+            # cal mask of attention weight
+            for i in range(valid_ratios.shape[0]):
+                valid_width = paddle.minimum(
+                    w.astype("int64"), paddle.ceil(valid_ratios[i] * w).astype("int64")
+                )
+                if valid_width < w:
+                    attn_weight[i, :, :, valid_width:, :] = float("-inf")
+
+        attn_weight = paddle.reshape(attn_weight, [bsz, T, -1])
+        attn_weight = F.softmax(attn_weight, axis=-1)
+
+        attn_weight = paddle.reshape(attn_weight, [bsz, T, h, w, c])
+        attn_weight = paddle.transpose(attn_weight, perm=[0, 1, 4, 2, 3])
+        # attn_weight: bsz * T * c * h * w
+        # feat: bsz * c * h * w
+        attn_feat = paddle.sum(
+            paddle.multiply(feat.unsqueeze(1), attn_weight), (3, 4), keepdim=False
+        )
+        # bsz * (seq_len + 1) * C
+
+        # Linear transformation
+        if self.pred_concat:
+            hf_c = holistic_feat.shape[-1]
+            holistic_feat = paddle.expand(holistic_feat, shape=[bsz, seq_len, hf_c])
+            y = self.prediction(
+                paddle.concat(
+                    (y, attn_feat.astype(y.dtype), holistic_feat.astype(y.dtype)), 2
+                )
+            )
+        else:
+            y = self.prediction(attn_feat)
+        # bsz * (seq_len + 1) * num_classes
+        if self.train_mode:
+            y = self.pred_dropout(y)
+
+        return y
+
+    def forward_train(self, feat, out_enc, label, img_metas):
+        """
+        img_metas: [label, valid_ratio]
+        """
+        if img_metas is not None:
+            assert img_metas[0].shape[0] == feat.shape[0]
+
+        valid_ratios = None
+        if img_metas is not None and self.mask:
+            valid_ratios = img_metas[-1]
+
+        lab_embedding = self.embedding(label)
+        # bsz * seq_len * emb_dim
+        out_enc = out_enc.unsqueeze(1).astype(lab_embedding.dtype)
+        # bsz * 1 * emb_dim
+        in_dec = paddle.concat((out_enc, lab_embedding), axis=1)
+        # bsz * (seq_len + 1) * C
+        out_dec = self._2d_attention(in_dec, feat, out_enc, valid_ratios=valid_ratios)
+
+        return out_dec[:, 1:, :]  # bsz * seq_len * num_classes
+
+    def forward_test(self, feat, out_enc, img_metas):
+        if img_metas is not None:
+            assert len(img_metas[0]) == feat.shape[0]
+
+        valid_ratios = None
+        if img_metas is not None and self.mask:
+            valid_ratios = img_metas[-1]
+
+        seq_len = self.max_seq_len
+        bsz = feat.shape[0]
+        start_token = paddle.full((bsz,), fill_value=self.start_idx, dtype="int64")
+        # bsz
+        start_token = self.embedding(start_token)
+        # bsz * emb_dim
+        emb_dim = start_token.shape[1]
+        start_token = start_token.unsqueeze(1)
+        start_token = paddle.expand(start_token, shape=[bsz, seq_len, emb_dim])
+        # bsz * seq_len * emb_dim
+        out_enc = out_enc.unsqueeze(1)
+        # bsz * 1 * emb_dim
+        decoder_input = paddle.concat((out_enc, start_token), axis=1)
+        # bsz * (seq_len + 1) * emb_dim
+
+        outputs = []
+        for i in range(1, seq_len + 1):
+            decoder_output = self._2d_attention(
+                decoder_input, feat, out_enc, valid_ratios=valid_ratios
+            )
+            char_output = decoder_output[:, i, :]  # bsz * num_classes
+            char_output = F.softmax(char_output, -1)
+            outputs.append(char_output)
+            max_idx = paddle.argmax(char_output, axis=1, keepdim=False)
+            char_embedding = self.embedding(max_idx)  # bsz * emb_dim
+            if i < seq_len:
+                decoder_input[:, i + 1, :] = char_embedding
+
+        outputs = paddle.stack(outputs, 1)  # bsz * seq_len * num_classes
+
+        return outputs
+
+
+class SARHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        enc_dim=512,
+        max_text_length=30,
+        enc_bi_rnn=False,
+        enc_drop_rnn=0.1,
+        enc_gru=False,
+        dec_bi_rnn=False,
+        dec_drop_rnn=0.0,
+        dec_gru=False,
+        d_k=512,
+        pred_dropout=0.1,
+        pred_concat=True,
+        **kwargs,
+    ):
+        super(SARHead, self).__init__()
+
+        # encoder module
+        self.encoder = SAREncoder(
+            enc_bi_rnn=enc_bi_rnn,
+            enc_drop_rnn=enc_drop_rnn,
+            enc_gru=enc_gru,
+            d_model=in_channels,
+            d_enc=enc_dim,
+        )
+
+        # decoder module
+        self.decoder = ParallelSARDecoder(
+            out_channels=out_channels,
+            enc_bi_rnn=enc_bi_rnn,
+            dec_bi_rnn=dec_bi_rnn,
+            dec_drop_rnn=dec_drop_rnn,
+            dec_gru=dec_gru,
+            d_model=in_channels,
+            d_enc=enc_dim,
+            d_k=d_k,
+            pred_dropout=pred_dropout,
+            max_text_length=max_text_length,
+            pred_concat=pred_concat,
+        )
+
+    def forward(self, feat, targets=None):
+        """
+        img_metas: [label, valid_ratio]
+        """
+        holistic_feat = self.encoder(feat, targets)  # bsz c
+
+        if self.training:
+            label = targets[0]  # label
+            final_out = self.decoder(feat, holistic_feat, label, img_metas=targets)
+        else:
+            final_out = self.decoder(
+                feat, holistic_feat, label=None, img_metas=targets, train_mode=False
+            )
+            # (bsz, seq_len, num_classes)
+
+        return final_out

ppocr/modeling/heads/rec_satrn_head.py

@@ -0,0 +1,592 @@
+# 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/open-mmlab/mmocr/blob/1.x/mmocr/models/textrecog/encoders/satrn_encoder.py
+https://github.com/open-mmlab/mmocr/blob/1.x/mmocr/models/textrecog/decoders/nrtr_decoder.py
+"""
+
+import math
+import numpy as np
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle import ParamAttr, reshape, transpose
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+from paddle.nn.initializer import KaimingNormal, Uniform, Constant
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self, num_channels, filter_size, num_filters, stride, padding, num_groups=1
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.conv = nn.Conv2D(
+            in_channels=num_channels,
+            out_channels=num_filters,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=padding,
+            groups=num_groups,
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm2D(
+            num_filters,
+            weight_attr=ParamAttr(initializer=Constant(1)),
+            bias_attr=ParamAttr(initializer=Constant(0)),
+        )
+        self.relu = nn.ReLU()
+
+    def forward(self, inputs):
+        y = self.conv(inputs)
+        y = self.bn(y)
+        y = self.relu(y)
+        return y
+
+
+class SATRNEncoderLayer(nn.Layer):
+    def __init__(
+        self,
+        d_model=512,
+        d_inner=512,
+        n_head=8,
+        d_k=64,
+        d_v=64,
+        dropout=0.1,
+        qkv_bias=False,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.attn = MultiHeadAttention(
+            n_head, d_model, d_k, d_v, qkv_bias=qkv_bias, dropout=dropout
+        )
+        self.norm2 = nn.LayerNorm(d_model)
+        self.feed_forward = LocalityAwareFeedforward(d_model, d_inner, dropout=dropout)
+
+    def forward(self, x, h, w, mask=None):
+        n, hw, c = x.shape
+        residual = x
+        x = self.norm1(x)
+        x = residual + self.attn(x, x, x, mask)
+        residual = x
+        x = self.norm2(x)
+        x = x.transpose([0, 2, 1]).reshape([n, c, h, w])
+        x = self.feed_forward(x)
+        x = x.reshape([n, c, hw]).transpose([0, 2, 1])
+        x = residual + x
+        return x
+
+
+class LocalityAwareFeedforward(nn.Layer):
+    def __init__(
+        self,
+        d_in,
+        d_hid,
+        dropout=0.1,
+    ):
+        super().__init__()
+        self.conv1 = ConvBNLayer(d_in, 1, d_hid, stride=1, padding=0)
+
+        self.depthwise_conv = ConvBNLayer(
+            d_hid, 3, d_hid, stride=1, padding=1, num_groups=d_hid
+        )
+
+        self.conv2 = ConvBNLayer(d_hid, 1, d_in, stride=1, padding=0)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.depthwise_conv(x)
+        x = self.conv2(x)
+
+        return x
+
+
+class Adaptive2DPositionalEncoding(nn.Layer):
+    def __init__(self, d_hid=512, n_height=100, n_width=100, dropout=0.1):
+        super().__init__()
+
+        h_position_encoder = self._get_sinusoid_encoding_table(n_height, d_hid)
+        h_position_encoder = h_position_encoder.transpose([1, 0])
+        h_position_encoder = h_position_encoder.reshape([1, d_hid, n_height, 1])
+
+        w_position_encoder = self._get_sinusoid_encoding_table(n_width, d_hid)
+        w_position_encoder = w_position_encoder.transpose([1, 0])
+        w_position_encoder = w_position_encoder.reshape([1, d_hid, 1, n_width])
+
+        self.register_buffer("h_position_encoder", h_position_encoder)
+        self.register_buffer("w_position_encoder", w_position_encoder)
+
+        self.h_scale = self.scale_factor_generate(d_hid)
+        self.w_scale = self.scale_factor_generate(d_hid)
+        self.pool = nn.AdaptiveAvgPool2D(1)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def _get_sinusoid_encoding_table(self, n_position, d_hid):
+        """Sinusoid position encoding table."""
+        denominator = paddle.to_tensor(
+            [1.0 / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
+        )
+        denominator = denominator.reshape([1, -1])
+        pos_tensor = paddle.cast(paddle.arange(n_position).unsqueeze(-1), "float32")
+        sinusoid_table = pos_tensor * denominator
+        sinusoid_table[:, 0::2] = paddle.sin(sinusoid_table[:, 0::2])
+        sinusoid_table[:, 1::2] = paddle.cos(sinusoid_table[:, 1::2])
+
+        return sinusoid_table
+
+    def scale_factor_generate(self, d_hid):
+        scale_factor = nn.Sequential(
+            nn.Conv2D(d_hid, d_hid, 1),
+            nn.ReLU(),
+            nn.Conv2D(d_hid, d_hid, 1),
+            nn.Sigmoid(),
+        )
+
+        return scale_factor
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+
+        avg_pool = self.pool(x)
+
+        h_pos_encoding = self.h_scale(avg_pool) * self.h_position_encoder[:, :, :h, :]
+        w_pos_encoding = self.w_scale(avg_pool) * self.w_position_encoder[:, :, :, :w]
+
+        out = x + h_pos_encoding + w_pos_encoding
+
+        out = self.dropout(out)
+
+        return out
+
+
+class ScaledDotProductAttention(nn.Layer):
+    def __init__(self, temperature, attn_dropout=0.1):
+        super().__init__()
+        self.temperature = temperature
+        self.dropout = nn.Dropout(attn_dropout)
+
+    def forward(self, q, k, v, mask=None):
+        def masked_fill(x, mask, value):
+            y = paddle.full(x.shape, value, x.dtype)
+            return paddle.where(mask, y, x)
+
+        attn = paddle.matmul(q / self.temperature, k.transpose([0, 1, 3, 2]))
+        if mask is not None:
+            attn = masked_fill(attn, mask == 0, -1e9)
+            # attn = attn.masked_fill(mask == 0, float('-inf'))
+            # attn += mask
+
+        attn = self.dropout(F.softmax(attn, axis=-1))
+        output = paddle.matmul(attn, v)
+
+        return output, attn
+
+
+class MultiHeadAttention(nn.Layer):
+    def __init__(
+        self, n_head=8, d_model=512, d_k=64, d_v=64, dropout=0.1, qkv_bias=False
+    ):
+        super().__init__()
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+
+        self.dim_k = n_head * d_k
+        self.dim_v = n_head * d_v
+
+        self.linear_q = nn.Linear(self.dim_k, self.dim_k, bias_attr=qkv_bias)
+        self.linear_k = nn.Linear(self.dim_k, self.dim_k, bias_attr=qkv_bias)
+        self.linear_v = nn.Linear(self.dim_v, self.dim_v, bias_attr=qkv_bias)
+
+        self.attention = ScaledDotProductAttention(d_k**0.5, dropout)
+
+        self.fc = nn.Linear(self.dim_v, d_model, bias_attr=qkv_bias)
+        self.proj_drop = nn.Dropout(dropout)
+
+    def forward(self, q, k, v, mask=None):
+        batch_size, len_q, _ = q.shape
+        _, len_k, _ = k.shape
+
+        q = self.linear_q(q).reshape([batch_size, len_q, self.n_head, self.d_k])
+        k = self.linear_k(k).reshape([batch_size, len_k, self.n_head, self.d_k])
+        v = self.linear_v(v).reshape([batch_size, len_k, self.n_head, self.d_v])
+
+        q, k, v = (
+            q.transpose([0, 2, 1, 3]),
+            k.transpose([0, 2, 1, 3]),
+            v.transpose([0, 2, 1, 3]),
+        )
+
+        if mask is not None:
+            if mask.dim() == 3:
+                mask = mask.unsqueeze(1)
+            elif mask.dim() == 2:
+                mask = mask.unsqueeze(1).unsqueeze(1)
+
+        attn_out, _ = self.attention(q, k, v, mask=mask)
+
+        attn_out = attn_out.transpose([0, 2, 1, 3]).reshape(
+            [batch_size, len_q, self.dim_v]
+        )
+
+        attn_out = self.fc(attn_out)
+        attn_out = self.proj_drop(attn_out)
+
+        return attn_out
+
+
+class SATRNEncoder(nn.Layer):
+    def __init__(
+        self,
+        n_layers=12,
+        n_head=8,
+        d_k=64,
+        d_v=64,
+        d_model=512,
+        n_position=100,
+        d_inner=256,
+        dropout=0.1,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.position_enc = Adaptive2DPositionalEncoding(
+            d_hid=d_model, n_height=n_position, n_width=n_position, dropout=dropout
+        )
+        self.layer_stack = nn.LayerList(
+            [
+                SATRNEncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+        self.layer_norm = nn.LayerNorm(d_model)
+
+    def forward(self, feat, valid_ratios=None):
+        """
+        Args:
+            feat (Tensor): Feature tensor of shape :math:`(N, D_m, H, W)`.
+            img_metas (dict): A dict that contains meta information of input
+                images. Preferably with the key ``valid_ratio``.
+
+        Returns:
+            Tensor: A tensor of shape :math:`(N, T, D_m)`.
+        """
+        if valid_ratios is None:
+            bs = feat.shape[0]
+            valid_ratios = paddle.full((bs, 1), 1.0, dtype=paddle.float32)
+
+        feat = self.position_enc(feat)
+        n, c, h, w = feat.shape
+
+        mask = paddle.zeros((n, h, w))
+        for i, valid_ratio in enumerate(valid_ratios):
+            valid_width = int(min(w, paddle.ceil(w * valid_ratio)))
+            mask[i, :, :valid_width] = 1
+
+        mask = mask.reshape([n, h * w])
+        feat = feat.reshape([n, c, h * w])
+
+        output = feat.transpose([0, 2, 1])
+        for enc_layer in self.layer_stack:
+            output = enc_layer(output, h, w, mask)
+        output = self.layer_norm(output)
+
+        return output
+
+
+class PositionwiseFeedForward(nn.Layer):
+    def __init__(self, d_in, d_hid, dropout=0.1):
+        super().__init__()
+        self.w_1 = nn.Linear(d_in, d_hid)
+        self.w_2 = nn.Linear(d_hid, d_in)
+        self.act = nn.GELU()
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.w_1(x)
+        x = self.act(x)
+        x = self.w_2(x)
+        x = self.dropout(x)
+
+        return x
+
+
+class PositionalEncoding(nn.Layer):
+    def __init__(self, d_hid=512, n_position=200, dropout=0):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        # Not a parameter
+        # Position table of shape (1, n_position, d_hid)
+        self.register_buffer(
+            "position_table", self._get_sinusoid_encoding_table(n_position, d_hid)
+        )
+
+    def _get_sinusoid_encoding_table(self, n_position, d_hid):
+        """Sinusoid position encoding table."""
+        denominator = paddle.to_tensor(
+            [1.0 / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
+        )
+        denominator = denominator.reshape([1, -1])
+        pos_tensor = paddle.cast(paddle.arange(n_position).unsqueeze(-1), "float32")
+        sinusoid_table = pos_tensor * denominator
+        sinusoid_table[:, 0::2] = paddle.sin(sinusoid_table[:, 0::2])
+        sinusoid_table[:, 1::2] = paddle.cos(sinusoid_table[:, 1::2])
+
+        return sinusoid_table.unsqueeze(0)
+
+    def forward(self, x):
+        x = x + self.position_table[:, : x.shape[1]].clone().detach()
+        return self.dropout(x)
+
+
+class TFDecoderLayer(nn.Layer):
+    def __init__(
+        self,
+        d_model=512,
+        d_inner=256,
+        n_head=8,
+        d_k=64,
+        d_v=64,
+        dropout=0.1,
+        qkv_bias=False,
+        operation_order=None,
+    ):
+        super().__init__()
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+
+        self.self_attn = MultiHeadAttention(
+            n_head, d_model, d_k, d_v, dropout=dropout, qkv_bias=qkv_bias
+        )
+
+        self.enc_attn = MultiHeadAttention(
+            n_head, d_model, d_k, d_v, dropout=dropout, qkv_bias=qkv_bias
+        )
+
+        self.mlp = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+        self.operation_order = operation_order
+        if self.operation_order is None:
+            self.operation_order = (
+                "norm",
+                "self_attn",
+                "norm",
+                "enc_dec_attn",
+                "norm",
+                "ffn",
+            )
+        assert self.operation_order in [
+            ("norm", "self_attn", "norm", "enc_dec_attn", "norm", "ffn"),
+            ("self_attn", "norm", "enc_dec_attn", "norm", "ffn", "norm"),
+        ]
+
+    def forward(
+        self, dec_input, enc_output, self_attn_mask=None, dec_enc_attn_mask=None
+    ):
+        if self.operation_order == (
+            "self_attn",
+            "norm",
+            "enc_dec_attn",
+            "norm",
+            "ffn",
+            "norm",
+        ):
+            dec_attn_out = self.self_attn(
+                dec_input, dec_input, dec_input, self_attn_mask
+            )
+            dec_attn_out += dec_input
+            dec_attn_out = self.norm1(dec_attn_out)
+
+            enc_dec_attn_out = self.enc_attn(
+                dec_attn_out, enc_output, enc_output, dec_enc_attn_mask
+            )
+            enc_dec_attn_out += dec_attn_out
+            enc_dec_attn_out = self.norm2(enc_dec_attn_out)
+
+            mlp_out = self.mlp(enc_dec_attn_out)
+            mlp_out += enc_dec_attn_out
+            mlp_out = self.norm3(mlp_out)
+        elif self.operation_order == (
+            "norm",
+            "self_attn",
+            "norm",
+            "enc_dec_attn",
+            "norm",
+            "ffn",
+        ):
+            dec_input_norm = self.norm1(dec_input)
+            dec_attn_out = self.self_attn(
+                dec_input_norm, dec_input_norm, dec_input_norm, self_attn_mask
+            )
+            dec_attn_out += dec_input
+
+            enc_dec_attn_in = self.norm2(dec_attn_out)
+            enc_dec_attn_out = self.enc_attn(
+                enc_dec_attn_in, enc_output, enc_output, dec_enc_attn_mask
+            )
+            enc_dec_attn_out += dec_attn_out
+
+            mlp_out = self.mlp(self.norm3(enc_dec_attn_out))
+            mlp_out += enc_dec_attn_out
+
+        return mlp_out
+
+
+class SATRNDecoder(nn.Layer):
+    def __init__(
+        self,
+        n_layers=6,
+        d_embedding=512,
+        n_head=8,
+        d_k=64,
+        d_v=64,
+        d_model=512,
+        d_inner=256,
+        n_position=200,
+        dropout=0.1,
+        num_classes=93,
+        max_seq_len=40,
+        start_idx=1,
+        padding_idx=92,
+    ):
+        super().__init__()
+
+        self.padding_idx = padding_idx
+        self.start_idx = start_idx
+        self.max_seq_len = max_seq_len
+
+        self.trg_word_emb = nn.Embedding(
+            num_classes, d_embedding, padding_idx=padding_idx
+        )
+
+        self.position_enc = PositionalEncoding(d_embedding, n_position=n_position)
+        self.dropout = nn.Dropout(p=dropout)
+
+        self.layer_stack = nn.LayerList(
+            [
+                TFDecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+        self.layer_norm = nn.LayerNorm(d_model, epsilon=1e-6)
+
+        pred_num_class = num_classes - 1  # ignore padding_idx
+        self.classifier = nn.Linear(d_model, pred_num_class)
+
+    @staticmethod
+    def get_pad_mask(seq, pad_idx):
+        return (seq != pad_idx).unsqueeze(-2)
+
+    @staticmethod
+    def get_subsequent_mask(seq):
+        """For masking out the subsequent info."""
+        len_s = seq.shape[1]
+        subsequent_mask = 1 - paddle.triu(paddle.ones((len_s, len_s)), diagonal=1)
+        subsequent_mask = paddle.cast(subsequent_mask.unsqueeze(0), "bool")
+
+        return subsequent_mask
+
+    def _attention(self, trg_seq, src, src_mask=None):
+        trg_embedding = self.trg_word_emb(trg_seq)
+        trg_pos_encoded = self.position_enc(trg_embedding)
+        tgt = self.dropout(trg_pos_encoded)
+
+        trg_mask = self.get_pad_mask(
+            trg_seq, pad_idx=self.padding_idx
+        ) & self.get_subsequent_mask(trg_seq)
+        output = tgt
+        for dec_layer in self.layer_stack:
+            output = dec_layer(
+                output, src, self_attn_mask=trg_mask, dec_enc_attn_mask=src_mask
+            )
+        output = self.layer_norm(output)
+
+        return output
+
+    def _get_mask(self, logit, valid_ratios):
+        N, T, _ = logit.shape
+        mask = None
+        if valid_ratios is not None:
+            mask = paddle.zeros((N, T))
+            for i, valid_ratio in enumerate(valid_ratios):
+                valid_width = min(T, math.ceil(T * valid_ratio))
+                mask[i, :valid_width] = 1
+
+        return mask
+
+    def forward_train(self, feat, out_enc, targets, valid_ratio):
+        src_mask = self._get_mask(out_enc, valid_ratio)
+        attn_output = self._attention(targets, out_enc, src_mask=src_mask)
+        outputs = self.classifier(attn_output)
+
+        return outputs
+
+    def forward_test(self, feat, out_enc, valid_ratio):
+        src_mask = self._get_mask(out_enc, valid_ratio)
+        N = out_enc.shape[0]
+        init_target_seq = paddle.full(
+            (N, self.max_seq_len + 1), self.padding_idx, dtype="int64"
+        )
+        # bsz * seq_len
+        init_target_seq[:, 0] = self.start_idx
+
+        outputs = []
+        for step in range(0, paddle.to_tensor(self.max_seq_len)):
+            decoder_output = self._attention(
+                init_target_seq, out_enc, src_mask=src_mask
+            )
+            # bsz * seq_len * C
+            step_result = F.softmax(
+                self.classifier(decoder_output[:, step, :]), axis=-1
+            )
+            # bsz * num_classes
+            outputs.append(step_result)
+            step_max_index = paddle.argmax(step_result, axis=-1)
+            init_target_seq[:, step + 1] = step_max_index
+
+        outputs = paddle.stack(outputs, axis=1)
+
+        return outputs
+
+    def forward(self, feat, out_enc, targets=None, valid_ratio=None):
+        if self.training:
+            return self.forward_train(feat, out_enc, targets, valid_ratio)
+        else:
+            return self.forward_test(feat, out_enc, valid_ratio)
+
+
+class SATRNHead(nn.Layer):
+    def __init__(self, enc_cfg, dec_cfg, **kwargs):
+        super(SATRNHead, self).__init__()
+
+        # encoder module
+        self.encoder = SATRNEncoder(**enc_cfg)
+
+        # decoder module
+        self.decoder = SATRNDecoder(**dec_cfg)
+
+    def forward(self, feat, targets=None):
+        if targets is not None:
+            targets, valid_ratio = targets
+        else:
+            targets, valid_ratio = None, None
+        holistic_feat = self.encoder(feat, valid_ratio)  # bsz c
+        final_out = self.decoder(feat, holistic_feat, targets, valid_ratio)
+
+        return final_out

ppocr/modeling/heads/rec_spin_att_head.py

@@ -0,0 +1,124 @@
+# 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/hikopensource/DAVAR-Lab-OCR/davarocr/davar_rcg/models/sequence_heads/att_head.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class SPINAttentionHead(nn.Layer):
+    def __init__(self, in_channels, out_channels, hidden_size, **kwargs):
+        super(SPINAttentionHead, self).__init__()
+        self.input_size = in_channels
+        self.hidden_size = hidden_size
+        self.num_classes = out_channels
+
+        self.attention_cell = AttentionLSTMCell(
+            in_channels, hidden_size, out_channels, use_gru=False
+        )
+        self.generator = nn.Linear(hidden_size, out_channels)
+
+    def _char_to_onehot(self, input_char, onehot_dim):
+        input_ont_hot = F.one_hot(input_char, onehot_dim)
+        return input_ont_hot
+
+    def forward(self, inputs, targets=None, batch_max_length=25):
+        batch_size = inputs.shape[0]
+        num_steps = batch_max_length + 1  # +1 for [sos] at end of sentence
+
+        hidden = (
+            paddle.zeros((batch_size, self.hidden_size)),
+            paddle.zeros((batch_size, self.hidden_size)),
+        )
+        output_hiddens = []
+        if self.training:  # for train
+            targets = targets[0]
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(
+                    targets[:, i], onehot_dim=self.num_classes
+                )
+                (outputs, hidden), alpha = self.attention_cell(
+                    hidden, inputs, char_onehots
+                )
+                output_hiddens.append(paddle.unsqueeze(outputs, axis=1))
+            output = paddle.concat(output_hiddens, axis=1)
+            probs = self.generator(output)
+        else:
+            targets = paddle.zeros(shape=[batch_size], dtype="int32")
+            probs = None
+            char_onehots = None
+            outputs = None
+            alpha = None
+
+            for i in range(num_steps):
+                char_onehots = self._char_to_onehot(
+                    targets, onehot_dim=self.num_classes
+                )
+                (outputs, hidden), alpha = self.attention_cell(
+                    hidden, inputs, char_onehots
+                )
+                probs_step = self.generator(outputs)
+                if probs is None:
+                    probs = paddle.unsqueeze(probs_step, axis=1)
+                else:
+                    probs = paddle.concat(
+                        [probs, paddle.unsqueeze(probs_step, axis=1)], axis=1
+                    )
+                next_input = probs_step.argmax(axis=1)
+                targets = next_input
+        if not self.training:
+            probs = paddle.nn.functional.softmax(probs, axis=2)
+        return probs
+
+
+class AttentionLSTMCell(nn.Layer):
+    def __init__(self, input_size, hidden_size, num_embeddings, use_gru=False):
+        super(AttentionLSTMCell, self).__init__()
+        self.i2h = nn.Linear(input_size, hidden_size, bias_attr=False)
+        self.h2h = nn.Linear(hidden_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias_attr=False)
+        if not use_gru:
+            self.rnn = nn.LSTMCell(
+                input_size=input_size + num_embeddings, hidden_size=hidden_size
+            )
+        else:
+            self.rnn = nn.GRUCell(
+                input_size=input_size + num_embeddings, hidden_size=hidden_size
+            )
+
+        self.hidden_size = hidden_size
+
+    def forward(self, prev_hidden, batch_H, char_onehots):
+        batch_H_proj = self.i2h(batch_H)
+        prev_hidden_proj = paddle.unsqueeze(self.h2h(prev_hidden[0]), axis=1)
+        res = paddle.add(batch_H_proj, prev_hidden_proj)
+        res = paddle.tanh(res)
+        e = self.score(res)
+
+        alpha = F.softmax(e, axis=1)
+        alpha = paddle.transpose(alpha, [0, 2, 1])
+        context = paddle.squeeze(paddle.mm(alpha, batch_H), axis=1)
+        concat_context = paddle.concat([context, char_onehots], 1)
+        cur_hidden = self.rnn(concat_context, prev_hidden)
+
+        return cur_hidden, alpha

ppocr/modeling/heads/rec_srn_head.py

@@ -0,0 +1,315 @@
+# 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.
+
+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 .self_attention import WrapEncoderForFeature
+from .self_attention import WrapEncoder
+from paddle.static import Program
+from ppocr.modeling.backbones.rec_resnet_fpn import ResNetFPN
+
+from collections import OrderedDict
+
+gradient_clip = 10
+
+
+class PVAM(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        char_num,
+        max_text_length,
+        num_heads,
+        num_encoder_tus,
+        hidden_dims,
+    ):
+        super(PVAM, self).__init__()
+        self.char_num = char_num
+        self.max_length = max_text_length
+        self.num_heads = num_heads
+        self.num_encoder_TUs = num_encoder_tus
+        self.hidden_dims = hidden_dims
+        # Transformer encoder
+        t = 256
+        c = 512
+        self.wrap_encoder_for_feature = WrapEncoderForFeature(
+            src_vocab_size=1,
+            max_length=t,
+            n_layer=self.num_encoder_TUs,
+            n_head=self.num_heads,
+            d_key=int(self.hidden_dims / self.num_heads),
+            d_value=int(self.hidden_dims / self.num_heads),
+            d_model=self.hidden_dims,
+            d_inner_hid=self.hidden_dims,
+            prepostprocess_dropout=0.1,
+            attention_dropout=0.1,
+            relu_dropout=0.1,
+            preprocess_cmd="n",
+            postprocess_cmd="da",
+            weight_sharing=True,
+        )
+
+        # PVAM
+        self.flatten0 = paddle.nn.Flatten(start_axis=0, stop_axis=1)
+        self.fc0 = paddle.nn.Linear(
+            in_features=in_channels,
+            out_features=in_channels,
+        )
+        self.emb = paddle.nn.Embedding(
+            num_embeddings=self.max_length, embedding_dim=in_channels
+        )
+        self.flatten1 = paddle.nn.Flatten(start_axis=0, stop_axis=2)
+        self.fc1 = paddle.nn.Linear(
+            in_features=in_channels, out_features=1, bias_attr=False
+        )
+
+    def forward(self, inputs, encoder_word_pos, gsrm_word_pos):
+        b, c, h, w = inputs.shape
+        conv_features = paddle.reshape(inputs, shape=[-1, c, h * w])
+        conv_features = paddle.transpose(conv_features, perm=[0, 2, 1])
+        # transformer encoder
+        b, t, c = conv_features.shape
+
+        enc_inputs = [conv_features, encoder_word_pos, None]
+        word_features = self.wrap_encoder_for_feature(enc_inputs)
+
+        # pvam
+        b, t, c = word_features.shape
+        word_features = self.fc0(word_features)
+        word_features_ = paddle.reshape(word_features, [-1, 1, t, c])
+        word_features_ = paddle.tile(word_features_, [1, self.max_length, 1, 1])
+        word_pos_feature = self.emb(gsrm_word_pos)
+        word_pos_feature_ = paddle.reshape(
+            word_pos_feature, [-1, self.max_length, 1, c]
+        )
+        word_pos_feature_ = paddle.tile(word_pos_feature_, [1, 1, t, 1])
+        y = word_pos_feature_ + word_features_
+        y = F.tanh(y)
+        attention_weight = self.fc1(y)
+        attention_weight = paddle.reshape(
+            attention_weight, shape=[-1, self.max_length, t]
+        )
+        attention_weight = F.softmax(attention_weight, axis=-1)
+        pvam_features = paddle.matmul(
+            attention_weight, word_features
+        )  # [b, max_length, c]
+        return pvam_features
+
+
+class GSRM(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        char_num,
+        max_text_length,
+        num_heads,
+        num_encoder_tus,
+        num_decoder_tus,
+        hidden_dims,
+    ):
+        super(GSRM, self).__init__()
+        self.char_num = char_num
+        self.max_length = max_text_length
+        self.num_heads = num_heads
+        self.num_encoder_TUs = num_encoder_tus
+        self.num_decoder_TUs = num_decoder_tus
+        self.hidden_dims = hidden_dims
+
+        self.fc0 = paddle.nn.Linear(in_features=in_channels, out_features=self.char_num)
+        self.wrap_encoder0 = WrapEncoder(
+            src_vocab_size=self.char_num + 1,
+            max_length=self.max_length,
+            n_layer=self.num_decoder_TUs,
+            n_head=self.num_heads,
+            d_key=int(self.hidden_dims / self.num_heads),
+            d_value=int(self.hidden_dims / self.num_heads),
+            d_model=self.hidden_dims,
+            d_inner_hid=self.hidden_dims,
+            prepostprocess_dropout=0.1,
+            attention_dropout=0.1,
+            relu_dropout=0.1,
+            preprocess_cmd="n",
+            postprocess_cmd="da",
+            weight_sharing=True,
+        )
+
+        self.wrap_encoder1 = WrapEncoder(
+            src_vocab_size=self.char_num + 1,
+            max_length=self.max_length,
+            n_layer=self.num_decoder_TUs,
+            n_head=self.num_heads,
+            d_key=int(self.hidden_dims / self.num_heads),
+            d_value=int(self.hidden_dims / self.num_heads),
+            d_model=self.hidden_dims,
+            d_inner_hid=self.hidden_dims,
+            prepostprocess_dropout=0.1,
+            attention_dropout=0.1,
+            relu_dropout=0.1,
+            preprocess_cmd="n",
+            postprocess_cmd="da",
+            weight_sharing=True,
+        )
+
+        self.mul = lambda x: paddle.matmul(
+            x=x, y=self.wrap_encoder0.prepare_decoder.emb0.weight, transpose_y=True
+        )
+
+    def forward(self, inputs, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2):
+        # ===== GSRM Visual-to-semantic embedding block =====
+        b, t, c = inputs.shape
+        pvam_features = paddle.reshape(inputs, [-1, c])
+        word_out = self.fc0(pvam_features)
+        word_ids = paddle.argmax(F.softmax(word_out), axis=1)
+        word_ids = paddle.reshape(x=word_ids, shape=[-1, t, 1])
+
+        # ===== GSRM Semantic reasoning block =====
+        """
+        This module is achieved through bi-transformers,
+        ngram_feature1 is the forward one, ngram_fetaure2 is the backward one
+        """
+        pad_idx = self.char_num
+
+        word1 = paddle.cast(word_ids, "float32")
+        word1 = F.pad(word1, [1, 0], value=1.0 * pad_idx, data_format="NLC")
+        word1 = paddle.cast(word1, "int64")
+        word1 = word1[:, :-1, :]
+        word2 = word_ids
+
+        enc_inputs_1 = [word1, gsrm_word_pos, gsrm_slf_attn_bias1]
+        enc_inputs_2 = [word2, gsrm_word_pos, gsrm_slf_attn_bias2]
+
+        gsrm_feature1 = self.wrap_encoder0(enc_inputs_1)
+        gsrm_feature2 = self.wrap_encoder1(enc_inputs_2)
+
+        gsrm_feature2 = F.pad(gsrm_feature2, [0, 1], value=0.0, data_format="NLC")
+        gsrm_feature2 = gsrm_feature2[
+            :,
+            1:,
+        ]
+        gsrm_features = gsrm_feature1 + gsrm_feature2
+
+        gsrm_out = self.mul(gsrm_features)
+
+        b, t, c = gsrm_out.shape
+        gsrm_out = paddle.reshape(gsrm_out, [-1, c])
+
+        return gsrm_features, word_out, gsrm_out
+
+
+class VSFD(nn.Layer):
+    def __init__(self, in_channels=512, pvam_ch=512, char_num=38):
+        super(VSFD, self).__init__()
+        self.char_num = char_num
+        self.fc0 = paddle.nn.Linear(in_features=in_channels * 2, out_features=pvam_ch)
+        self.fc1 = paddle.nn.Linear(in_features=pvam_ch, out_features=self.char_num)
+
+    def forward(self, pvam_feature, gsrm_feature):
+        b, t, c1 = pvam_feature.shape
+        b, t, c2 = gsrm_feature.shape
+        combine_feature_ = paddle.concat([pvam_feature, gsrm_feature], axis=2)
+        img_comb_feature_ = paddle.reshape(combine_feature_, shape=[-1, c1 + c2])
+        img_comb_feature_map = self.fc0(img_comb_feature_)
+        img_comb_feature_map = F.sigmoid(img_comb_feature_map)
+        img_comb_feature_map = paddle.reshape(img_comb_feature_map, shape=[-1, t, c1])
+        combine_feature = (
+            img_comb_feature_map * pvam_feature
+            + (1.0 - img_comb_feature_map) * gsrm_feature
+        )
+        img_comb_feature = paddle.reshape(combine_feature, shape=[-1, c1])
+
+        out = self.fc1(img_comb_feature)
+        return out
+
+
+class SRNHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        max_text_length,
+        num_heads,
+        num_encoder_TUs,
+        num_decoder_TUs,
+        hidden_dims,
+        **kwargs,
+    ):
+        super(SRNHead, self).__init__()
+        self.char_num = out_channels
+        self.max_length = max_text_length
+        self.num_heads = num_heads
+        self.num_encoder_TUs = num_encoder_TUs
+        self.num_decoder_TUs = num_decoder_TUs
+        self.hidden_dims = hidden_dims
+
+        self.pvam = PVAM(
+            in_channels=in_channels,
+            char_num=self.char_num,
+            max_text_length=self.max_length,
+            num_heads=self.num_heads,
+            num_encoder_tus=self.num_encoder_TUs,
+            hidden_dims=self.hidden_dims,
+        )
+
+        self.gsrm = GSRM(
+            in_channels=in_channels,
+            char_num=self.char_num,
+            max_text_length=self.max_length,
+            num_heads=self.num_heads,
+            num_encoder_tus=self.num_encoder_TUs,
+            num_decoder_tus=self.num_decoder_TUs,
+            hidden_dims=self.hidden_dims,
+        )
+        self.vsfd = VSFD(in_channels=in_channels, char_num=self.char_num)
+
+        self.gsrm.wrap_encoder1.prepare_decoder.emb0 = (
+            self.gsrm.wrap_encoder0.prepare_decoder.emb0
+        )
+
+    def forward(self, inputs, targets=None):
+        others = targets[-4:]
+        encoder_word_pos = others[0]
+        gsrm_word_pos = others[1]
+        gsrm_slf_attn_bias1 = others[2]
+        gsrm_slf_attn_bias2 = others[3]
+
+        pvam_feature = self.pvam(inputs, encoder_word_pos, gsrm_word_pos)
+
+        gsrm_feature, word_out, gsrm_out = self.gsrm(
+            pvam_feature, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2
+        )
+
+        final_out = self.vsfd(pvam_feature, gsrm_feature)
+        if not self.training:
+            final_out = F.softmax(final_out, axis=1)
+
+        _, decoded_out = paddle.topk(final_out, k=1)
+
+        predicts = OrderedDict(
+            [
+                ("predict", final_out),
+                ("pvam_feature", pvam_feature),
+                ("decoded_out", decoded_out),
+                ("word_out", word_out),
+                ("gsrm_out", gsrm_out),
+            ]
+        )
+
+        return predicts

ppocr/modeling/heads/rec_visionlan_head.py

@@ -0,0 +1,474 @@
+# 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/wangyuxin87/VisionLAN
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle.nn.initializer import Normal, XavierNormal
+import numpy as np
+
+
+class PositionalEncoding(nn.Layer):
+    def __init__(self, d_hid, n_position=200):
+        super(PositionalEncoding, self).__init__()
+        self.register_buffer(
+            "pos_table", self._get_sinusoid_encoding_table(n_position, d_hid)
+        )
+
+    def _get_sinusoid_encoding_table(self, n_position, d_hid):
+        """Sinusoid position encoding table"""
+
+        def get_position_angle_vec(position):
+            return [
+                position / np.power(10000, 2 * (hid_j // 2) / d_hid)
+                for hid_j in range(d_hid)
+            ]
+
+        sinusoid_table = np.array(
+            [get_position_angle_vec(pos_i) for pos_i in range(n_position)]
+        )
+        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
+        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1
+        sinusoid_table = paddle.to_tensor(sinusoid_table, dtype="float32")
+        sinusoid_table = paddle.unsqueeze(sinusoid_table, axis=0)
+        return sinusoid_table
+
+    def forward(self, x):
+        return x + self.pos_table[:, : x.shape[1]].clone().detach()
+
+
+class ScaledDotProductAttention(nn.Layer):
+    "Scaled Dot-Product Attention"
+
+    def __init__(self, temperature, attn_dropout=0.1):
+        super(ScaledDotProductAttention, self).__init__()
+        self.temperature = temperature
+        self.dropout = nn.Dropout(attn_dropout)
+        self.softmax = nn.Softmax(axis=2)
+
+    def forward(self, q, k, v, mask=None):
+        k = paddle.transpose(k, perm=[0, 2, 1])
+        attn = paddle.bmm(q, k)
+        attn = attn / self.temperature
+        if mask is not None:
+            attn = attn.masked_fill(mask, -1e9)
+            if mask.dim() == 3:
+                mask = paddle.unsqueeze(mask, axis=1)
+            elif mask.dim() == 2:
+                mask = paddle.unsqueeze(mask, axis=1)
+                mask = paddle.unsqueeze(mask, axis=1)
+            repeat_times = [
+                attn.shape[1] // mask.shape[1],
+                attn.shape[2] // mask.shape[2],
+            ]
+            mask = paddle.tile(mask, [1, repeat_times[0], repeat_times[1], 1])
+            attn[mask == 0] = -1e9
+        attn = self.softmax(attn)
+        attn = self.dropout(attn)
+        output = paddle.bmm(attn, v)
+        return output
+
+
+class MultiHeadAttention(nn.Layer):
+    "Multi-Head Attention module"
+
+    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
+        super(MultiHeadAttention, self).__init__()
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+        self.w_qs = nn.Linear(
+            d_model,
+            n_head * d_k,
+            weight_attr=ParamAttr(
+                initializer=Normal(mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
+            ),
+        )
+        self.w_ks = nn.Linear(
+            d_model,
+            n_head * d_k,
+            weight_attr=ParamAttr(
+                initializer=Normal(mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
+            ),
+        )
+        self.w_vs = nn.Linear(
+            d_model,
+            n_head * d_v,
+            weight_attr=ParamAttr(
+                initializer=Normal(mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
+            ),
+        )
+
+        self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5))
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.fc = nn.Linear(
+            n_head * d_v, d_model, weight_attr=ParamAttr(initializer=XavierNormal())
+        )
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, q, k, v, mask=None):
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+        sz_b, len_q, _ = q.shape
+        sz_b, len_k, _ = k.shape
+        sz_b, len_v, _ = v.shape
+        residual = q
+
+        q = self.w_qs(q)
+        q = paddle.reshape(q, shape=[-1, len_q, n_head, d_k])  # 4*21*512 ---- 4*21*8*64
+        k = self.w_ks(k)
+        k = paddle.reshape(k, shape=[-1, len_k, n_head, d_k])
+        v = self.w_vs(v)
+        v = paddle.reshape(v, shape=[-1, len_v, n_head, d_v])
+
+        q = paddle.transpose(q, perm=[2, 0, 1, 3])
+        q = paddle.reshape(q, shape=[-1, len_q, d_k])  # (n*b) x lq x dk
+        k = paddle.transpose(k, perm=[2, 0, 1, 3])
+        k = paddle.reshape(k, shape=[-1, len_k, d_k])  # (n*b) x lk x dk
+        v = paddle.transpose(v, perm=[2, 0, 1, 3])
+        v = paddle.reshape(v, shape=[-1, len_v, d_v])  # (n*b) x lv x dv
+
+        mask = (
+            paddle.tile(mask, [n_head, 1, 1]) if mask is not None else None
+        )  # (n*b) x .. x ..
+        output = self.attention(q, k, v, mask=mask)
+        output = paddle.reshape(output, shape=[n_head, -1, len_q, d_v])
+        output = paddle.transpose(output, perm=[1, 2, 0, 3])
+        output = paddle.reshape(
+            output, shape=[-1, len_q, n_head * d_v]
+        )  # b x lq x (n*dv)
+        output = self.dropout(self.fc(output))
+        output = self.layer_norm(output + residual)
+        return output
+
+
+class PositionwiseFeedForward(nn.Layer):
+    def __init__(self, d_in, d_hid, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Conv1D(d_in, d_hid, 1)  # position-wise
+        self.w_2 = nn.Conv1D(d_hid, d_in, 1)  # position-wise
+        self.layer_norm = nn.LayerNorm(d_in)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        residual = x
+        x = paddle.transpose(x, perm=[0, 2, 1])
+        x = self.w_2(F.relu(self.w_1(x)))
+        x = paddle.transpose(x, perm=[0, 2, 1])
+        x = self.dropout(x)
+        x = self.layer_norm(x + residual)
+        return x
+
+
+class EncoderLayer(nn.Layer):
+    """Compose with two layers"""
+
+    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
+        super(EncoderLayer, self).__init__()
+        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+    def forward(self, enc_input, slf_attn_mask=None):
+        enc_output = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)
+        enc_output = self.pos_ffn(enc_output)
+        return enc_output
+
+
+class Transformer_Encoder(nn.Layer):
+    def __init__(
+        self,
+        n_layers=2,
+        n_head=8,
+        d_word_vec=512,
+        d_k=64,
+        d_v=64,
+        d_model=512,
+        d_inner=2048,
+        dropout=0.1,
+        n_position=256,
+    ):
+        super(Transformer_Encoder, self).__init__()
+        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
+        self.dropout = nn.Dropout(p=dropout)
+        self.layer_stack = nn.LayerList(
+            [
+                EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+                for _ in range(n_layers)
+            ]
+        )
+        self.layer_norm = nn.LayerNorm(d_model, epsilon=1e-6)
+
+    def forward(self, enc_output, src_mask, return_attns=False):
+        enc_output = self.dropout(self.position_enc(enc_output))  # position embedding
+        for enc_layer in self.layer_stack:
+            enc_output = enc_layer(enc_output, slf_attn_mask=src_mask)
+        enc_output = self.layer_norm(enc_output)
+        return enc_output
+
+
+class PP_layer(nn.Layer):
+    def __init__(self, n_dim=512, N_max_character=25, n_position=256):
+        super(PP_layer, self).__init__()
+        self.character_len = N_max_character
+        self.f0_embedding = nn.Embedding(N_max_character, n_dim)
+        self.w0 = nn.Linear(N_max_character, n_position)
+        self.wv = nn.Linear(n_dim, n_dim)
+        self.we = nn.Linear(n_dim, N_max_character)
+        self.active = nn.Tanh()
+        self.softmax = nn.Softmax(axis=2)
+
+    def forward(self, enc_output):
+        # enc_output: b,256,512
+        reading_order = paddle.arange(self.character_len, dtype="int64")
+        reading_order = reading_order.unsqueeze(0).expand(
+            [enc_output.shape[0], self.character_len]
+        )  # (S,) -> (B, S)
+        reading_order = self.f0_embedding(reading_order)  # b,25,512
+
+        # calculate attention
+        reading_order = paddle.transpose(reading_order, perm=[0, 2, 1])
+        t = self.w0(reading_order)  # b,512,256
+        t = self.active(
+            paddle.transpose(t, perm=[0, 2, 1]) + self.wv(enc_output)
+        )  # b,256,512
+        t = self.we(t)  # b,256,25
+        t = self.softmax(paddle.transpose(t, perm=[0, 2, 1]))  # b,25,256
+        g_output = paddle.bmm(t, enc_output)  # b,25,512
+        return g_output
+
+
+class Prediction(nn.Layer):
+    def __init__(self, n_dim=512, n_position=256, N_max_character=25, n_class=37):
+        super(Prediction, self).__init__()
+        self.pp = PP_layer(
+            n_dim=n_dim, N_max_character=N_max_character, n_position=n_position
+        )
+        self.pp_share = PP_layer(
+            n_dim=n_dim, N_max_character=N_max_character, n_position=n_position
+        )
+        self.w_vrm = nn.Linear(n_dim, n_class)  # output layer
+        self.w_share = nn.Linear(n_dim, n_class)  # output layer
+        self.nclass = n_class
+
+    def forward(self, cnn_feature, f_res, f_sub, train_mode=False, use_mlm=True):
+        if train_mode:
+            if not use_mlm:
+                g_output = self.pp(cnn_feature)  # b,25,512
+                g_output = self.w_vrm(g_output)
+                f_res = 0
+                f_sub = 0
+                return g_output, f_res, f_sub
+            g_output = self.pp(cnn_feature)  # b,25,512
+            f_res = self.pp_share(f_res)
+            f_sub = self.pp_share(f_sub)
+            g_output = self.w_vrm(g_output)
+            f_res = self.w_share(f_res)
+            f_sub = self.w_share(f_sub)
+            return g_output, f_res, f_sub
+        else:
+            g_output = self.pp(cnn_feature)  # b,25,512
+            g_output = self.w_vrm(g_output)
+            return g_output
+
+
+class MLM(nn.Layer):
+    "Architecture of MLM"
+
+    def __init__(self, n_dim=512, n_position=256, max_text_length=25):
+        super(MLM, self).__init__()
+        self.MLM_SequenceModeling_mask = Transformer_Encoder(
+            n_layers=2, n_position=n_position
+        )
+        self.MLM_SequenceModeling_WCL = Transformer_Encoder(
+            n_layers=1, n_position=n_position
+        )
+        self.pos_embedding = nn.Embedding(max_text_length, n_dim)
+        self.w0_linear = nn.Linear(1, n_position)
+        self.wv = nn.Linear(n_dim, n_dim)
+        self.active = nn.Tanh()
+        self.we = nn.Linear(n_dim, 1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x, label_pos):
+        # transformer unit for generating mask_c
+        feature_v_seq = self.MLM_SequenceModeling_mask(x, src_mask=None)
+        # position embedding layer
+        label_pos = paddle.to_tensor(label_pos, dtype="int64")
+        pos_emb = self.pos_embedding(label_pos)
+        pos_emb = self.w0_linear(paddle.unsqueeze(pos_emb, axis=2))
+        pos_emb = paddle.transpose(pos_emb, perm=[0, 2, 1])
+        # fusion position embedding with features V & generate mask_c
+        att_map_sub = self.active(pos_emb + self.wv(feature_v_seq))
+        att_map_sub = self.we(att_map_sub)  # b,256,1
+        att_map_sub = paddle.transpose(att_map_sub, perm=[0, 2, 1])
+        att_map_sub = self.sigmoid(att_map_sub)  # b,1,256
+        # WCL
+        ## generate inputs for WCL
+        att_map_sub = paddle.transpose(att_map_sub, perm=[0, 2, 1])
+        f_res = x * (1 - att_map_sub)  # second path with remaining string
+        f_sub = x * att_map_sub  # first path with occluded character
+        ## transformer units in WCL
+        f_res = self.MLM_SequenceModeling_WCL(f_res, src_mask=None)
+        f_sub = self.MLM_SequenceModeling_WCL(f_sub, src_mask=None)
+        return f_res, f_sub, att_map_sub
+
+
+def trans_1d_2d(x):
+    b, w_h, c = x.shape  # b, 256, 512
+    x = paddle.transpose(x, perm=[0, 2, 1])
+    x = paddle.reshape(x, [-1, c, 32, 8])
+    x = paddle.transpose(x, perm=[0, 1, 3, 2])  # [b, c, 8, 32]
+    return x
+
+
+class MLM_VRM(nn.Layer):
+    """
+    MLM+VRM, MLM is only used in training.
+    ratio controls the occluded number in a batch.
+    The pipeline of VisionLAN in testing is very concise with only a backbone + sequence modeling(transformer unit) + prediction layer(pp layer).
+    x: input image
+    label_pos: character index
+    training_step: LF or LA process
+    output
+    text_pre: prediction of VRM
+    test_rem: prediction of remaining string in MLM
+    text_mas: prediction of occluded character in MLM
+    mask_c_show: visualization of Mask_c
+    """
+
+    def __init__(
+        self, n_layers=3, n_position=256, n_dim=512, max_text_length=25, nclass=37
+    ):
+        super(MLM_VRM, self).__init__()
+        self.MLM = MLM(
+            n_dim=n_dim, n_position=n_position, max_text_length=max_text_length
+        )
+        self.SequenceModeling = Transformer_Encoder(
+            n_layers=n_layers, n_position=n_position
+        )
+        self.Prediction = Prediction(
+            n_dim=n_dim,
+            n_position=n_position,
+            N_max_character=max_text_length
+            + 1,  # N_max_character = 1 eos + 25 characters
+            n_class=nclass,
+        )
+        self.nclass = nclass
+        self.max_text_length = max_text_length
+
+    def forward(self, x, label_pos, training_step, train_mode=False):
+        b, c, h, w = x.shape
+        nT = self.max_text_length
+        x = paddle.transpose(x, perm=[0, 1, 3, 2])
+        x = paddle.reshape(x, [-1, c, h * w])
+        x = paddle.transpose(x, perm=[0, 2, 1])
+        if train_mode:
+            if training_step == "LF_1":
+                f_res = 0
+                f_sub = 0
+                x = self.SequenceModeling(x, src_mask=None)
+                text_pre, test_rem, text_mas = self.Prediction(
+                    x, f_res, f_sub, train_mode=True, use_mlm=False
+                )
+                return text_pre, text_pre, text_pre, text_pre
+            elif training_step == "LF_2":
+                # MLM
+                f_res, f_sub, mask_c = self.MLM(x, label_pos)
+                x = self.SequenceModeling(x, src_mask=None)
+                text_pre, test_rem, text_mas = self.Prediction(
+                    x, f_res, f_sub, train_mode=True
+                )
+                mask_c_show = trans_1d_2d(mask_c)
+                return text_pre, test_rem, text_mas, mask_c_show
+            elif training_step == "LA":
+                # MLM
+                f_res, f_sub, mask_c = self.MLM(x, label_pos)
+                ## use the mask_c (1 for occluded character and 0 for remaining characters) to occlude input
+                ## ratio controls the occluded number in a batch
+                character_mask = paddle.zeros_like(mask_c)
+
+                ratio = b // 2
+                if ratio >= 1:
+                    with paddle.no_grad():
+                        character_mask[0:ratio, :, :] = mask_c[0:ratio, :, :]
+                else:
+                    character_mask = mask_c
+                x = x * (1 - character_mask)
+                # VRM
+                ## transformer unit for VRM
+                x = self.SequenceModeling(x, src_mask=None)
+                ## prediction layer for MLM and VSR
+                text_pre, test_rem, text_mas = self.Prediction(
+                    x, f_res, f_sub, train_mode=True
+                )
+                mask_c_show = trans_1d_2d(mask_c)
+                return text_pre, test_rem, text_mas, mask_c_show
+            else:
+                raise NotImplementedError
+        else:  # VRM is only used in the testing stage
+            f_res = 0
+            f_sub = 0
+            contextual_feature = self.SequenceModeling(x, src_mask=None)
+            text_pre = self.Prediction(
+                contextual_feature, f_res, f_sub, train_mode=False, use_mlm=False
+            )
+            text_pre = paddle.transpose(text_pre, perm=[1, 0, 2])  # (26, b, 37))
+            return text_pre, x
+
+
+class VLHead(nn.Layer):
+    """
+    Architecture of VisionLAN
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels=36,
+        n_layers=3,
+        n_position=256,
+        n_dim=512,
+        max_text_length=25,
+        training_step="LA",
+    ):
+        super(VLHead, self).__init__()
+        self.MLM_VRM = MLM_VRM(
+            n_layers=n_layers,
+            n_position=n_position,
+            n_dim=n_dim,
+            max_text_length=max_text_length,
+            nclass=out_channels + 1,
+        )
+        self.training_step = training_step
+
+    def forward(self, feat, targets=None):
+        if self.training:
+            label_pos = targets[-2]
+            text_pre, test_rem, text_mas, mask_map = self.MLM_VRM(
+                feat, label_pos, self.training_step, train_mode=True
+            )
+            return text_pre, test_rem, text_mas, mask_map
+        else:
+            text_pre, x = self.MLM_VRM(
+                feat, targets, self.training_step, train_mode=False
+            )
+            return text_pre, x

ppocr/modeling/heads/self_attention.py

@@ -0,0 +1,460 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+
+import paddle
+from paddle import ParamAttr, nn
+from paddle import nn, ParamAttr
+from paddle.nn import functional as F
+import numpy as np
+
+gradient_clip = 10
+
+
+class WrapEncoderForFeature(nn.Layer):
+    def __init__(
+        self,
+        src_vocab_size,
+        max_length,
+        n_layer,
+        n_head,
+        d_key,
+        d_value,
+        d_model,
+        d_inner_hid,
+        prepostprocess_dropout,
+        attention_dropout,
+        relu_dropout,
+        preprocess_cmd,
+        postprocess_cmd,
+        weight_sharing,
+        bos_idx=0,
+    ):
+        super(WrapEncoderForFeature, self).__init__()
+
+        self.prepare_encoder = PrepareEncoder(
+            src_vocab_size,
+            d_model,
+            max_length,
+            prepostprocess_dropout,
+            bos_idx=bos_idx,
+            word_emb_param_name="src_word_emb_table",
+        )
+        self.encoder = Encoder(
+            n_layer,
+            n_head,
+            d_key,
+            d_value,
+            d_model,
+            d_inner_hid,
+            prepostprocess_dropout,
+            attention_dropout,
+            relu_dropout,
+            preprocess_cmd,
+            postprocess_cmd,
+        )
+
+    def forward(self, enc_inputs):
+        conv_features, src_pos, src_slf_attn_bias = enc_inputs
+        enc_input = self.prepare_encoder(conv_features, src_pos)
+        enc_output = self.encoder(enc_input, src_slf_attn_bias)
+        return enc_output
+
+
+class WrapEncoder(nn.Layer):
+    """
+    embedder + encoder
+    """
+
+    def __init__(
+        self,
+        src_vocab_size,
+        max_length,
+        n_layer,
+        n_head,
+        d_key,
+        d_value,
+        d_model,
+        d_inner_hid,
+        prepostprocess_dropout,
+        attention_dropout,
+        relu_dropout,
+        preprocess_cmd,
+        postprocess_cmd,
+        weight_sharing,
+        bos_idx=0,
+    ):
+        super(WrapEncoder, self).__init__()
+
+        self.prepare_decoder = PrepareDecoder(
+            src_vocab_size, d_model, max_length, prepostprocess_dropout, bos_idx=bos_idx
+        )
+        self.encoder = Encoder(
+            n_layer,
+            n_head,
+            d_key,
+            d_value,
+            d_model,
+            d_inner_hid,
+            prepostprocess_dropout,
+            attention_dropout,
+            relu_dropout,
+            preprocess_cmd,
+            postprocess_cmd,
+        )
+
+    def forward(self, enc_inputs):
+        src_word, src_pos, src_slf_attn_bias = enc_inputs
+        enc_input = self.prepare_decoder(src_word, src_pos)
+        enc_output = self.encoder(enc_input, src_slf_attn_bias)
+        return enc_output
+
+
+class Encoder(nn.Layer):
+    """
+    encoder
+    """
+
+    def __init__(
+        self,
+        n_layer,
+        n_head,
+        d_key,
+        d_value,
+        d_model,
+        d_inner_hid,
+        prepostprocess_dropout,
+        attention_dropout,
+        relu_dropout,
+        preprocess_cmd="n",
+        postprocess_cmd="da",
+    ):
+        super(Encoder, self).__init__()
+
+        self.encoder_layers = list()
+        for i in range(n_layer):
+            self.encoder_layers.append(
+                self.add_sublayer(
+                    "layer_%d" % i,
+                    EncoderLayer(
+                        n_head,
+                        d_key,
+                        d_value,
+                        d_model,
+                        d_inner_hid,
+                        prepostprocess_dropout,
+                        attention_dropout,
+                        relu_dropout,
+                        preprocess_cmd,
+                        postprocess_cmd,
+                    ),
+                )
+            )
+        self.processor = PrePostProcessLayer(
+            preprocess_cmd, d_model, prepostprocess_dropout
+        )
+
+    def forward(self, enc_input, attn_bias):
+        for encoder_layer in self.encoder_layers:
+            enc_output = encoder_layer(enc_input, attn_bias)
+            enc_input = enc_output
+        enc_output = self.processor(enc_output)
+        return enc_output
+
+
+class EncoderLayer(nn.Layer):
+    """
+    EncoderLayer
+    """
+
+    def __init__(
+        self,
+        n_head,
+        d_key,
+        d_value,
+        d_model,
+        d_inner_hid,
+        prepostprocess_dropout,
+        attention_dropout,
+        relu_dropout,
+        preprocess_cmd="n",
+        postprocess_cmd="da",
+    ):
+        super(EncoderLayer, self).__init__()
+        self.preprocesser1 = PrePostProcessLayer(
+            preprocess_cmd, d_model, prepostprocess_dropout
+        )
+        self.self_attn = MultiHeadAttention(
+            d_key, d_value, d_model, n_head, attention_dropout
+        )
+        self.postprocesser1 = PrePostProcessLayer(
+            postprocess_cmd, d_model, prepostprocess_dropout
+        )
+
+        self.preprocesser2 = PrePostProcessLayer(
+            preprocess_cmd, d_model, prepostprocess_dropout
+        )
+        self.ffn = FFN(d_inner_hid, d_model, relu_dropout)
+        self.postprocesser2 = PrePostProcessLayer(
+            postprocess_cmd, d_model, prepostprocess_dropout
+        )
+
+    def forward(self, enc_input, attn_bias):
+        attn_output = self.self_attn(
+            self.preprocesser1(enc_input), None, None, attn_bias
+        )
+        attn_output = self.postprocesser1(attn_output, enc_input)
+        ffn_output = self.ffn(self.preprocesser2(attn_output))
+        ffn_output = self.postprocesser2(ffn_output, attn_output)
+        return ffn_output
+
+
+class MultiHeadAttention(nn.Layer):
+    """
+    Multi-Head Attention
+    """
+
+    def __init__(self, d_key, d_value, d_model, n_head=1, dropout_rate=0.0):
+        super(MultiHeadAttention, self).__init__()
+        self.n_head = n_head
+        self.d_key = d_key
+        self.d_value = d_value
+        self.d_model = d_model
+        self.dropout_rate = dropout_rate
+        self.q_fc = paddle.nn.Linear(
+            in_features=d_model, out_features=d_key * n_head, bias_attr=False
+        )
+        self.k_fc = paddle.nn.Linear(
+            in_features=d_model, out_features=d_key * n_head, bias_attr=False
+        )
+        self.v_fc = paddle.nn.Linear(
+            in_features=d_model, out_features=d_value * n_head, bias_attr=False
+        )
+        self.proj_fc = paddle.nn.Linear(
+            in_features=d_value * n_head, out_features=d_model, bias_attr=False
+        )
+
+    def _prepare_qkv(self, queries, keys, values, cache=None):
+        if keys is None:  # self-attention
+            keys, values = queries, queries
+            static_kv = False
+        else:  # cross-attention
+            static_kv = True
+
+        q = self.q_fc(queries)
+        q = paddle.reshape(x=q, shape=[0, 0, self.n_head, self.d_key])
+        q = paddle.transpose(x=q, perm=[0, 2, 1, 3])
+
+        if cache is not None and static_kv and "static_k" in cache:
+            # for encoder-decoder attention in inference and has cached
+            k = cache["static_k"]
+            v = cache["static_v"]
+        else:
+            k = self.k_fc(keys)
+            v = self.v_fc(values)
+            k = paddle.reshape(x=k, shape=[0, 0, self.n_head, self.d_key])
+            k = paddle.transpose(x=k, perm=[0, 2, 1, 3])
+            v = paddle.reshape(x=v, shape=[0, 0, self.n_head, self.d_value])
+            v = paddle.transpose(x=v, perm=[0, 2, 1, 3])
+
+        if cache is not None:
+            if static_kv and not "static_k" in cache:
+                # for encoder-decoder attention in inference and has not cached
+                cache["static_k"], cache["static_v"] = k, v
+            elif not static_kv:
+                # for decoder self-attention in inference
+                cache_k, cache_v = cache["k"], cache["v"]
+                k = paddle.concat([cache_k, k], axis=2)
+                v = paddle.concat([cache_v, v], axis=2)
+                cache["k"], cache["v"] = k, v
+
+        return q, k, v
+
+    def forward(self, queries, keys, values, attn_bias, cache=None):
+        # compute q ,k ,v
+        keys = queries if keys is None else keys
+        values = keys if values is None else values
+        q, k, v = self._prepare_qkv(queries, keys, values, cache)
+
+        # scale dot product attention
+        product = paddle.matmul(x=q, y=k, transpose_y=True)
+        product = product * self.d_model**-0.5
+        if attn_bias is not None:
+            product += attn_bias.astype(product.dtype)
+        weights = F.softmax(product)
+        if self.dropout_rate:
+            weights = F.dropout(weights, p=self.dropout_rate, mode="downscale_in_infer")
+        out = paddle.matmul(weights, v)
+
+        # combine heads
+        out = paddle.transpose(out, perm=[0, 2, 1, 3])
+        out = paddle.reshape(x=out, shape=[0, 0, out.shape[2] * out.shape[3]])
+
+        # project to output
+        out = self.proj_fc(out)
+
+        return out
+
+
+class PrePostProcessLayer(nn.Layer):
+    """
+    PrePostProcessLayer
+    """
+
+    def __init__(self, process_cmd, d_model, dropout_rate):
+        super(PrePostProcessLayer, self).__init__()
+        self.process_cmd = process_cmd
+        self.functors = []
+        for cmd in self.process_cmd:
+            if cmd == "a":  # add residual connection
+                self.functors.append(lambda x, y: x + y if y is not None else x)
+            elif cmd == "n":  # add layer normalization
+                self.functors.append(
+                    self.add_sublayer(
+                        "layer_norm_%d" % len(self.sublayers()),
+                        paddle.nn.LayerNorm(
+                            normalized_shape=d_model,
+                            weight_attr=paddle.ParamAttr(
+                                initializer=paddle.nn.initializer.Constant(1.0)
+                            ),
+                            bias_attr=paddle.ParamAttr(
+                                initializer=paddle.nn.initializer.Constant(0.0)
+                            ),
+                        ),
+                    )
+                )
+            elif cmd == "d":  # add dropout
+                self.functors.append(
+                    lambda x: (
+                        F.dropout(x, p=dropout_rate, mode="downscale_in_infer")
+                        if dropout_rate
+                        else x
+                    )
+                )
+
+    def forward(self, x, residual=None):
+        for i, cmd in enumerate(self.process_cmd):
+            if cmd == "a":
+                x = self.functors[i](x, residual)
+            else:
+                x = self.functors[i](x)
+        return x
+
+
+class PrepareEncoder(nn.Layer):
+    def __init__(
+        self,
+        src_vocab_size,
+        src_emb_dim,
+        src_max_len,
+        dropout_rate=0,
+        bos_idx=0,
+        word_emb_param_name=None,
+        pos_enc_param_name=None,
+    ):
+        super(PrepareEncoder, self).__init__()
+        self.src_emb_dim = src_emb_dim
+        self.src_max_len = src_max_len
+        self.emb = paddle.nn.Embedding(
+            num_embeddings=self.src_max_len, embedding_dim=self.src_emb_dim
+        )
+        self.dropout_rate = dropout_rate
+
+    def forward(self, src_word, src_pos):
+        src_word_emb = src_word
+        src_word_emb = paddle.cast(src_word_emb, "float32")
+        src_word_emb = paddle.scale(x=src_word_emb, scale=self.src_emb_dim**0.5)
+        src_pos = paddle.squeeze(src_pos, axis=-1)
+        src_pos_enc = self.emb(src_pos)
+        src_pos_enc.stop_gradient = True
+        enc_input = src_word_emb + src_pos_enc
+        if self.dropout_rate:
+            out = F.dropout(x=enc_input, p=self.dropout_rate, mode="downscale_in_infer")
+        else:
+            out = enc_input
+        return out
+
+
+class PrepareDecoder(nn.Layer):
+    def __init__(
+        self,
+        src_vocab_size,
+        src_emb_dim,
+        src_max_len,
+        dropout_rate=0,
+        bos_idx=0,
+        word_emb_param_name=None,
+        pos_enc_param_name=None,
+    ):
+        super(PrepareDecoder, self).__init__()
+        self.src_emb_dim = src_emb_dim
+        """
+        self.emb0 = Embedding(num_embeddings=src_vocab_size,
+                              embedding_dim=src_emb_dim)
+        """
+        self.emb0 = paddle.nn.Embedding(
+            num_embeddings=src_vocab_size,
+            embedding_dim=self.src_emb_dim,
+            padding_idx=bos_idx,
+            weight_attr=paddle.ParamAttr(
+                name=word_emb_param_name,
+                initializer=nn.initializer.Normal(0.0, src_emb_dim**-0.5),
+            ),
+        )
+        self.emb1 = paddle.nn.Embedding(
+            num_embeddings=src_max_len,
+            embedding_dim=self.src_emb_dim,
+            weight_attr=paddle.ParamAttr(name=pos_enc_param_name),
+        )
+        self.dropout_rate = dropout_rate
+
+    def forward(self, src_word, src_pos):
+        src_word = paddle.cast(src_word, "int64")
+        src_word = paddle.squeeze(src_word, axis=-1)
+        src_word_emb = self.emb0(src_word)
+        src_word_emb = paddle.scale(x=src_word_emb, scale=self.src_emb_dim**0.5)
+        src_pos = paddle.squeeze(src_pos, axis=-1)
+        src_pos_enc = self.emb1(src_pos)
+        src_pos_enc.stop_gradient = True
+        enc_input = src_word_emb + src_pos_enc
+        if self.dropout_rate:
+            out = F.dropout(x=enc_input, p=self.dropout_rate, mode="downscale_in_infer")
+        else:
+            out = enc_input
+        return out
+
+
+class FFN(nn.Layer):
+    """
+    Feed-Forward Network
+    """
+
+    def __init__(self, d_inner_hid, d_model, dropout_rate):
+        super(FFN, self).__init__()
+        self.dropout_rate = dropout_rate
+        self.fc1 = paddle.nn.Linear(in_features=d_model, out_features=d_inner_hid)
+        self.fc2 = paddle.nn.Linear(in_features=d_inner_hid, out_features=d_model)
+
+    def forward(self, x):
+        hidden = self.fc1(x)
+        hidden = F.relu(hidden)
+        if self.dropout_rate:
+            hidden = F.dropout(hidden, p=self.dropout_rate, mode="downscale_in_infer")
+        out = self.fc2(hidden)
+        return out

ppocr/modeling/heads/sr_rensnet_transformer.py

@@ -0,0 +1,427 @@
+# 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/text-gestalt/loss/transformer_english_decomposition.py
+"""
+import copy
+import math
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+def subsequent_mask(size):
+    """Generate a square mask for the sequence. The masked positions are filled with float('-inf').
+    Unmasked positions are filled with float(0.0).
+    """
+    mask = paddle.ones([1, size, size], dtype="float32")
+    mask_inf = paddle.triu(
+        paddle.full(shape=[1, size, size], dtype="float32", fill_value="-inf"),
+        diagonal=1,
+    )
+    mask = mask + mask_inf
+    padding_mask = paddle.equal(mask, paddle.to_tensor(1, dtype=mask.dtype))
+    return padding_mask
+
+
+def clones(module, N):
+    return nn.LayerList([copy.deepcopy(module) for _ in range(N)])
+
+
+def masked_fill(x, mask, value):
+    y = paddle.full(x.shape, value, x.dtype)
+    return paddle.where(mask, y, x)
+
+
+def attention(query, key, value, mask=None, dropout=None, attention_map=None):
+    d_k = query.shape[-1]
+    scores = paddle.matmul(query, paddle.transpose(key, [0, 1, 3, 2])) / math.sqrt(d_k)
+
+    if mask is not None:
+        scores = masked_fill(scores, mask == 0, float("-inf"))
+    else:
+        pass
+
+    p_attn = F.softmax(scores, axis=-1)
+
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return paddle.matmul(p_attn, value), p_attn
+
+
+class MultiHeadedAttention(nn.Layer):
+    def __init__(self, h, d_model, dropout=0.1, compress_attention=False):
+        super(MultiHeadedAttention, self).__init__()
+        assert d_model % h == 0
+        self.d_k = d_model // h
+        self.h = h
+        self.linears = clones(nn.Linear(d_model, d_model), 4)
+        self.attn = None
+        self.dropout = nn.Dropout(p=dropout, mode="downscale_in_infer")
+        self.compress_attention = compress_attention
+        self.compress_attention_linear = nn.Linear(h, 1)
+
+    def forward(self, query, key, value, mask=None, attention_map=None):
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+        nbatches = query.shape[0]
+
+        query, key, value = [
+            paddle.transpose(
+                l(x).reshape([nbatches, -1, self.h, self.d_k]), [0, 2, 1, 3]
+            )
+            for l, x in zip(self.linears, (query, key, value))
+        ]
+
+        x, attention_map = attention(
+            query,
+            key,
+            value,
+            mask=mask,
+            dropout=self.dropout,
+            attention_map=attention_map,
+        )
+
+        x = paddle.reshape(
+            paddle.transpose(x, [0, 2, 1, 3]), [nbatches, -1, self.h * self.d_k]
+        )
+
+        return self.linears[-1](x), attention_map
+
+
+class ResNet(nn.Layer):
+    def __init__(self, num_in, block, layers):
+        super(ResNet, self).__init__()
+
+        self.conv1 = nn.Conv2D(num_in, 64, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2D(64, use_global_stats=True)
+        self.relu1 = nn.ReLU()
+        self.pool = nn.MaxPool2D((2, 2), (2, 2))
+
+        self.conv2 = nn.Conv2D(64, 128, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2D(128, use_global_stats=True)
+        self.relu2 = nn.ReLU()
+
+        self.layer1_pool = nn.MaxPool2D((2, 2), (2, 2))
+        self.layer1 = self._make_layer(block, 128, 256, layers[0])
+        self.layer1_conv = nn.Conv2D(256, 256, 3, 1, 1)
+        self.layer1_bn = nn.BatchNorm2D(256, use_global_stats=True)
+        self.layer1_relu = nn.ReLU()
+
+        self.layer2_pool = nn.MaxPool2D((2, 2), (2, 2))
+        self.layer2 = self._make_layer(block, 256, 256, layers[1])
+        self.layer2_conv = nn.Conv2D(256, 256, 3, 1, 1)
+        self.layer2_bn = nn.BatchNorm2D(256, use_global_stats=True)
+        self.layer2_relu = nn.ReLU()
+
+        self.layer3_pool = nn.MaxPool2D((2, 2), (2, 2))
+        self.layer3 = self._make_layer(block, 256, 512, layers[2])
+        self.layer3_conv = nn.Conv2D(512, 512, 3, 1, 1)
+        self.layer3_bn = nn.BatchNorm2D(512, use_global_stats=True)
+        self.layer3_relu = nn.ReLU()
+
+        self.layer4_pool = nn.MaxPool2D((2, 2), (2, 2))
+        self.layer4 = self._make_layer(block, 512, 512, layers[3])
+        self.layer4_conv2 = nn.Conv2D(512, 1024, 3, 1, 1)
+        self.layer4_conv2_bn = nn.BatchNorm2D(1024, use_global_stats=True)
+        self.layer4_conv2_relu = nn.ReLU()
+
+    def _make_layer(self, block, inplanes, planes, blocks):
+        if inplanes != planes:
+            downsample = nn.Sequential(
+                nn.Conv2D(inplanes, planes, 3, 1, 1),
+                nn.BatchNorm2D(planes, use_global_stats=True),
+            )
+        else:
+            downsample = None
+        layers = []
+        layers.append(block(inplanes, planes, downsample))
+        for i in range(1, blocks):
+            layers.append(block(planes, planes, downsample=None))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu1(x)
+        x = self.pool(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu2(x)
+
+        x = self.layer1_pool(x)
+        x = self.layer1(x)
+        x = self.layer1_conv(x)
+        x = self.layer1_bn(x)
+        x = self.layer1_relu(x)
+
+        x = self.layer2(x)
+        x = self.layer2_conv(x)
+        x = self.layer2_bn(x)
+        x = self.layer2_relu(x)
+
+        x = self.layer3(x)
+        x = self.layer3_conv(x)
+        x = self.layer3_bn(x)
+        x = self.layer3_relu(x)
+
+        x = self.layer4(x)
+        x = self.layer4_conv2(x)
+        x = self.layer4_conv2_bn(x)
+        x = self.layer4_conv2_relu(x)
+
+        return x
+
+
+class Bottleneck(nn.Layer):
+    def __init__(self, input_dim):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2D(input_dim, input_dim, 1)
+        self.bn1 = nn.BatchNorm2D(input_dim, use_global_stats=True)
+        self.relu = nn.ReLU()
+
+        self.conv2 = nn.Conv2D(input_dim, input_dim, 3, 1, 1)
+        self.bn2 = nn.BatchNorm2D(input_dim, use_global_stats=True)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class PositionalEncoding(nn.Layer):
+    "Implement the PE function."
+
+    def __init__(self, dropout, dim, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout, mode="downscale_in_infer")
+
+        pe = paddle.zeros([max_len, dim])
+        position = paddle.arange(0, max_len, dtype=paddle.float32).unsqueeze(1)
+        div_term = paddle.exp(
+            paddle.arange(0, dim, 2).astype("float32") * (-math.log(10000.0) / dim)
+        )
+        pe[:, 0::2] = paddle.sin(position * div_term)
+        pe[:, 1::2] = paddle.cos(position * div_term)
+        pe = paddle.unsqueeze(pe, 0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + self.pe[:, : x.shape[1]]
+        return self.dropout(x)
+
+
+class PositionwiseFeedForward(nn.Layer):
+    "Implements FFN equation."
+
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout, mode="downscale_in_infer")
+
+    def forward(self, x):
+        return self.w_2(self.dropout(F.relu(self.w_1(x))))
+
+
+class Generator(nn.Layer):
+    "Define standard linear + softmax generation step."
+
+    def __init__(self, d_model, vocab):
+        super(Generator, self).__init__()
+        self.proj = nn.Linear(d_model, vocab)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        out = self.proj(x)
+        return out
+
+
+class Embeddings(nn.Layer):
+    def __init__(self, d_model, vocab):
+        super(Embeddings, self).__init__()
+        self.lut = nn.Embedding(vocab, d_model)
+        self.d_model = d_model
+
+    def forward(self, x):
+        embed = self.lut(x) * math.sqrt(self.d_model)
+        return embed
+
+
+class LayerNorm(nn.Layer):
+    "Construct a layernorm module (See citation for details)."
+
+    def __init__(self, features, eps=1e-6):
+        super(LayerNorm, self).__init__()
+        self.a_2 = self.create_parameter(
+            shape=[features], default_initializer=paddle.nn.initializer.Constant(1.0)
+        )
+        self.b_2 = self.create_parameter(
+            shape=[features], default_initializer=paddle.nn.initializer.Constant(0.0)
+        )
+        self.eps = eps
+
+    def forward(self, x):
+        mean = x.mean(-1, keepdim=True)
+        std = x.std(-1, keepdim=True)
+        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
+
+
+class Decoder(nn.Layer):
+    def __init__(self):
+        super(Decoder, self).__init__()
+
+        self.mask_multihead = MultiHeadedAttention(h=16, d_model=1024, dropout=0.1)
+        self.mul_layernorm1 = LayerNorm(1024)
+
+        self.multihead = MultiHeadedAttention(h=16, d_model=1024, dropout=0.1)
+        self.mul_layernorm2 = LayerNorm(1024)
+
+        self.pff = PositionwiseFeedForward(1024, 2048)
+        self.mul_layernorm3 = LayerNorm(1024)
+
+    def forward(self, text, conv_feature, attention_map=None):
+        text_max_length = text.shape[1]
+        mask = subsequent_mask(text_max_length)
+        result = text
+        result = self.mul_layernorm1(
+            result + self.mask_multihead(text, text, text, mask=mask)[0]
+        )
+        b, c, h, w = conv_feature.shape
+        conv_feature = paddle.transpose(conv_feature.reshape([b, c, h * w]), [0, 2, 1])
+        word_image_align, attention_map = self.multihead(
+            result, conv_feature, conv_feature, mask=None, attention_map=attention_map
+        )
+        result = self.mul_layernorm2(result + word_image_align)
+        result = self.mul_layernorm3(result + self.pff(result))
+
+        return result, attention_map
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self, inplanes, planes, downsample):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2D(inplanes, planes, kernel_size=3, stride=1, padding=1)
+        self.bn1 = nn.BatchNorm2D(planes, use_global_stats=True)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2D(planes, planes, kernel_size=3, stride=1, padding=1)
+        self.bn2 = nn.BatchNorm2D(planes, use_global_stats=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample != None:
+            residual = self.downsample(residual)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Encoder(nn.Layer):
+    def __init__(self):
+        super(Encoder, self).__init__()
+        self.cnn = ResNet(num_in=1, block=BasicBlock, layers=[1, 2, 5, 3])
+
+    def forward(self, input):
+        conv_result = self.cnn(input)
+        return conv_result
+
+
+class Transformer(nn.Layer):
+    def __init__(self, in_channels=1, alphabet="0123456789"):
+        super(Transformer, self).__init__()
+        self.alphabet = alphabet
+        word_n_class = self.get_alphabet_len()
+        self.embedding_word_with_upperword = Embeddings(512, word_n_class)
+        self.pe = PositionalEncoding(dim=512, dropout=0.1, max_len=5000)
+
+        self.encoder = Encoder()
+        self.decoder = Decoder()
+        self.generator_word_with_upperword = Generator(1024, word_n_class)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.initializer.XavierNormal(p)
+
+    def get_alphabet_len(self):
+        return len(self.alphabet)
+
+    def forward(self, image, text_length, text_input, attention_map=None):
+        if image.shape[1] == 3:
+            R = image[:, 0:1, :, :]
+            G = image[:, 1:2, :, :]
+            B = image[:, 2:3, :, :]
+            image = 0.299 * R + 0.587 * G + 0.114 * B
+
+        conv_feature = self.encoder(image)  # batch, 1024, 8, 32
+        max_length = max(text_length)
+        text_input = text_input[:, :max_length]
+
+        text_embedding = self.embedding_word_with_upperword(
+            text_input
+        )  # batch, text_max_length, 512
+        postion_embedding = self.pe(
+            paddle.zeros(text_embedding.shape)
+        )  # batch, text_max_length, 512
+        text_input_with_pe = paddle.concat(
+            [text_embedding, postion_embedding], 2
+        )  # batch, text_max_length, 1024
+        batch, seq_len, _ = text_input_with_pe.shape
+
+        text_input_with_pe, word_attention_map = self.decoder(
+            text_input_with_pe, conv_feature
+        )
+
+        word_decoder_result = self.generator_word_with_upperword(text_input_with_pe)
+
+        if self.training:
+            total_length = paddle.sum(text_length)
+            probs_res = paddle.zeros([total_length, self.get_alphabet_len()])
+            start = 0
+
+            for index, length in enumerate(text_length):
+                length = int(length.numpy())
+                probs_res[start : start + length, :] = word_decoder_result[
+                    index, 0 : 0 + length, :
+                ]
+
+                start = start + length
+
+            return probs_res, word_attention_map, None
+        else:
+            return word_decoder_result

ppocr/modeling/heads/table_att_head.py

@@ -0,0 +1,421 @@
+# copyright (c) 2021 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
+import paddle.nn as nn
+from paddle import ParamAttr
+import paddle.nn.functional as F
+import numpy as np
+
+from .rec_att_head import AttentionGRUCell
+from ppocr.modeling.backbones.rec_svtrnet import DropPath, Identity, Mlp
+
+
+def get_para_bias_attr(l2_decay, k):
+    if l2_decay > 0:
+        regularizer = paddle.regularizer.L2Decay(l2_decay)
+        stdv = 1.0 / math.sqrt(k * 1.0)
+        initializer = nn.initializer.Uniform(-stdv, stdv)
+    else:
+        regularizer = None
+        initializer = None
+    weight_attr = ParamAttr(regularizer=regularizer, initializer=initializer)
+    bias_attr = ParamAttr(regularizer=regularizer, initializer=initializer)
+    return [weight_attr, bias_attr]
+
+
+class TableAttentionHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        in_max_len=488,
+        max_text_length=800,
+        out_channels=30,
+        loc_reg_num=4,
+        **kwargs,
+    ):
+        super(TableAttentionHead, self).__init__()
+        self.input_size = in_channels[-1]
+        self.hidden_size = hidden_size
+        self.out_channels = out_channels
+        self.max_text_length = max_text_length
+
+        self.structure_attention_cell = AttentionGRUCell(
+            self.input_size, hidden_size, self.out_channels, use_gru=False
+        )
+        self.structure_generator = nn.Linear(hidden_size, self.out_channels)
+        self.in_max_len = in_max_len
+
+        if self.in_max_len == 640:
+            self.loc_fea_trans = nn.Linear(400, self.max_text_length + 1)
+        elif self.in_max_len == 800:
+            self.loc_fea_trans = nn.Linear(625, self.max_text_length + 1)
+        else:
+            self.loc_fea_trans = nn.Linear(256, self.max_text_length + 1)
+        self.loc_generator = nn.Linear(self.input_size + hidden_size, loc_reg_num)
+
+    def _char_to_onehot(self, input_char, onehot_dim):
+        input_ont_hot = F.one_hot(input_char, onehot_dim)
+        return input_ont_hot
+
+    def forward(self, inputs, targets=None):
+        # if and else branch are both needed when you want to assign a variable
+        # if you modify the var in just one branch, then the modification will not work.
+        fea = inputs[-1]
+        last_shape = int(np.prod(fea.shape[2:]))  # gry added
+        fea = paddle.reshape(fea, [fea.shape[0], fea.shape[1], last_shape])
+        fea = fea.transpose([0, 2, 1])  # (NTC)(batch, width, channels)
+        batch_size = fea.shape[0]
+
+        hidden = paddle.zeros((batch_size, self.hidden_size))
+        output_hiddens = paddle.zeros(
+            (batch_size, self.max_text_length + 1, self.hidden_size)
+        )
+        if self.training and targets is not None:
+            structure = targets[0]
+            for i in range(self.max_text_length + 1):
+                elem_onehots = self._char_to_onehot(
+                    structure[:, i], onehot_dim=self.out_channels
+                )
+                (outputs, hidden), alpha = self.structure_attention_cell(
+                    hidden, fea, elem_onehots
+                )
+                output_hiddens[:, i, :] = outputs
+            structure_probs = self.structure_generator(output_hiddens)
+            loc_fea = fea.transpose([0, 2, 1])
+            loc_fea = self.loc_fea_trans(loc_fea)
+            loc_fea = loc_fea.transpose([0, 2, 1])
+            loc_concat = paddle.concat([output_hiddens, loc_fea], axis=2)
+            loc_preds = self.loc_generator(loc_concat)
+            loc_preds = F.sigmoid(loc_preds)
+        else:
+            temp_elem = paddle.zeros(shape=[batch_size], dtype="int32")
+            structure_probs = None
+            loc_preds = None
+            elem_onehots = None
+            outputs = None
+            alpha = None
+            max_text_length = paddle.to_tensor(self.max_text_length)
+            for i in range(max_text_length + 1):
+                elem_onehots = self._char_to_onehot(
+                    temp_elem, onehot_dim=self.out_channels
+                )
+                (outputs, hidden), alpha = self.structure_attention_cell(
+                    hidden, fea, elem_onehots
+                )
+                output_hiddens[:, i, :] = outputs
+                structure_probs_step = self.structure_generator(outputs)
+                temp_elem = structure_probs_step.argmax(axis=1, dtype="int32")
+
+            structure_probs = self.structure_generator(output_hiddens)
+            structure_probs = F.softmax(structure_probs)
+            loc_fea = fea.transpose([0, 2, 1])
+            loc_fea = self.loc_fea_trans(loc_fea)
+            loc_fea = loc_fea.transpose([0, 2, 1])
+            loc_concat = paddle.concat([output_hiddens, loc_fea], axis=2)
+            loc_preds = self.loc_generator(loc_concat)
+            loc_preds = F.sigmoid(loc_preds)
+        return {"structure_probs": structure_probs, "loc_preds": loc_preds}
+
+
+class HWAttention(nn.Layer):
+    def __init__(
+        self,
+        head_dim=32,
+        qk_scale=None,
+        attn_drop=0.0,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.scale = qk_scale or self.head_dim**-0.5
+        self.attn_drop = nn.Dropout(attn_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        C = C // 3
+        qkv = x.reshape([B, N, 3, C // self.head_dim, self.head_dim]).transpose(
+            [2, 0, 3, 1, 4]
+        )
+        q, k, v = qkv.unbind(0)
+        attn = q @ k.transpose([0, 1, 3, 2]) * self.scale
+        attn = F.softmax(attn, -1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+        x = x.transpose([0, 2, 1]).reshape([B, N, C])
+        return x
+
+
+def img2windows(img, H_sp, W_sp):
+    """
+    img: B C H W
+    """
+    B, H, W, C = img.shape
+    img_reshape = img.reshape([B, H // H_sp, H_sp, W // W_sp, W_sp, C])
+    img_perm = img_reshape.transpose([0, 1, 3, 2, 4, 5]).reshape([-1, H_sp * W_sp, C])
+    return img_perm
+
+
+def windows2img(img_splits_hw, H_sp, W_sp, H, W):
+    """
+    img_splits_hw: B' H W C
+    """
+    B = int(img_splits_hw.shape[0] / (H * W / H_sp / W_sp))
+
+    img = img_splits_hw.reshape([B, H // H_sp, W // W_sp, H_sp, W_sp, -1])
+    img = img.transpose([0, 1, 3, 2, 4, 5]).flatten(1, 4)
+    return img
+
+
+class Block(nn.Layer):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        split_h=4,
+        split_w=4,
+        h_num_heads=None,
+        w_num_heads=None,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        eps=1e-6,
+    ):
+        super().__init__()
+        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+        self.split_h = split_h
+        self.split_w = split_w
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.norm1 = norm_layer(dim, epsilon=eps)
+        self.h_num_heads = h_num_heads if h_num_heads is not None else num_heads // 2
+        self.w_num_heads = w_num_heads if w_num_heads is not None else num_heads // 2
+        self.head_dim = dim // num_heads
+        self.mixer = HWAttention(
+            head_dim=dim // num_heads,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
+        self.norm2 = norm_layer(dim, epsilon=eps)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose([0, 2, 1])
+
+        qkv = self.qkv(x).reshape([B, H, W, 3 * C])
+
+        x1 = qkv[:, :, :, : 3 * self.h_num_heads * self.head_dim]  # b, h, w, 3ch
+        x2 = qkv[:, :, :, 3 * self.h_num_heads * self.head_dim :]  # b, h, w, 3cw
+
+        x1 = self.mixer(img2windows(x1, self.split_h, W))  # b*splith, W, 3ch
+        x2 = self.mixer(img2windows(x2, H, self.split_w))  # b*splitw, h, 3ch
+        x1 = windows2img(x1, self.split_h, W, H, W)
+        x2 = windows2img(x2, H, self.split_w, H, W)
+
+        attened_x = paddle.concat([x1, x2], 2)
+        attened_x = self.proj(attened_x)
+
+        x = self.norm1(x + self.drop_path(attened_x))
+        x = self.norm2(x + self.drop_path(self.mlp(x)))
+        x = x.transpose([0, 2, 1]).reshape([-1, C, H, W])
+        return x
+
+
+class SLAHead(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        hidden_size,
+        out_channels=30,
+        max_text_length=500,
+        loc_reg_num=4,
+        fc_decay=0.0,
+        use_attn=False,
+        **kwargs,
+    ):
+        """
+        @param in_channels: input shape
+        @param hidden_size: hidden_size for RNN and Embedding
+        @param out_channels: num_classes to rec
+        @param max_text_length: max text pred
+        """
+        super().__init__()
+
+        if isinstance(in_channels, int):
+            self.is_next = True
+            in_channels = 512
+        else:
+            self.is_next = False
+            in_channels = in_channels[-1]
+        self.hidden_size = hidden_size
+        self.max_text_length = max_text_length
+        self.emb = self._char_to_onehot
+        self.num_embeddings = out_channels
+        self.loc_reg_num = loc_reg_num
+        self.eos = self.num_embeddings - 1
+
+        # structure
+        self.structure_attention_cell = AttentionGRUCell(
+            in_channels, hidden_size, self.num_embeddings
+        )
+        weight_attr, bias_attr = get_para_bias_attr(l2_decay=fc_decay, k=hidden_size)
+        weight_attr1_1, bias_attr1_1 = get_para_bias_attr(
+            l2_decay=fc_decay, k=hidden_size
+        )
+        weight_attr1_2, bias_attr1_2 = get_para_bias_attr(
+            l2_decay=fc_decay, k=hidden_size
+        )
+        self.structure_generator = nn.Sequential(
+            nn.Linear(
+                self.hidden_size,
+                self.hidden_size,
+                weight_attr=weight_attr1_2,
+                bias_attr=bias_attr1_2,
+            ),
+            nn.Linear(
+                hidden_size, out_channels, weight_attr=weight_attr, bias_attr=bias_attr
+            ),
+        )
+        dpr = np.linspace(0, 0.1, 2)
+
+        self.use_attn = use_attn
+        if use_attn:
+            layer_list = [
+                Block(
+                    in_channels,
+                    num_heads=2,
+                    mlp_ratio=4.0,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                )
+                for i in range(2)
+            ]
+            self.cross_atten = nn.Sequential(*layer_list)
+        # loc
+        weight_attr1, bias_attr1 = get_para_bias_attr(
+            l2_decay=fc_decay, k=self.hidden_size
+        )
+        weight_attr2, bias_attr2 = get_para_bias_attr(
+            l2_decay=fc_decay, k=self.hidden_size
+        )
+        self.loc_generator = nn.Sequential(
+            nn.Linear(
+                self.hidden_size,
+                self.hidden_size,
+                weight_attr=weight_attr1,
+                bias_attr=bias_attr1,
+            ),
+            nn.Linear(
+                self.hidden_size,
+                loc_reg_num,
+                weight_attr=weight_attr2,
+                bias_attr=bias_attr2,
+            ),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, inputs, targets=None):
+        if self.is_next == True:
+            fea = inputs
+            batch_size = fea.shape[0]
+        else:
+            fea = inputs[-1]
+            batch_size = fea.shape[0]
+            if self.use_attn:
+                fea = fea + self.cross_atten(fea)
+            # reshape
+            fea = paddle.reshape(fea, [fea.shape[0], fea.shape[1], -1])
+            fea = fea.transpose([0, 2, 1])  # (NTC)(batch, width, channels)
+
+        hidden = paddle.zeros((batch_size, self.hidden_size))
+        structure_preds = paddle.zeros(
+            (batch_size, self.max_text_length + 1, self.num_embeddings)
+        )
+        loc_preds = paddle.zeros(
+            (batch_size, self.max_text_length + 1, self.loc_reg_num)
+        )
+        structure_preds.stop_gradient = True
+        loc_preds.stop_gradient = True
+
+        if self.training and targets is not None:
+            structure = targets[0]
+            max_len = targets[-2].max().astype("int32")
+            for i in range(max_len + 1):
+                hidden, structure_step, loc_step = self._decode(
+                    structure[:, i], fea, hidden
+                )
+                structure_preds[:, i, :] = structure_step
+                loc_preds[:, i, :] = loc_step
+            structure_preds = structure_preds[:, : max_len + 1]
+            loc_preds = loc_preds[:, : max_len + 1]
+        else:
+            structure_ids = paddle.zeros(
+                (batch_size, self.max_text_length + 1), dtype="int32"
+            )
+            pre_chars = paddle.zeros(shape=[batch_size], dtype="int32")
+            max_text_length = paddle.to_tensor(self.max_text_length)
+            for i in range(max_text_length + 1):
+                hidden, structure_step, loc_step = self._decode(pre_chars, fea, hidden)
+                pre_chars = structure_step.argmax(axis=1, dtype="int32")
+                structure_preds[:, i, :] = structure_step
+                loc_preds[:, i, :] = loc_step
+
+                structure_ids[:, i] = pre_chars
+                if (structure_ids == self.eos).any(-1).all():
+                    break
+        if not self.training:
+            structure_preds = F.softmax(structure_preds[:, : i + 1])
+            loc_preds = loc_preds[:, : i + 1]
+        return {"structure_probs": structure_preds, "loc_preds": loc_preds}
+
+    def _decode(self, pre_chars, features, hidden):
+        """
+        Predict table label and coordinates for each step
+        @param pre_chars: Table label in previous step
+        @param features:
+        @param hidden: hidden status in previous step
+        @return:
+        """
+        emb_feature = self.emb(pre_chars)
+        # output shape is b * self.hidden_size
+        (output, hidden), alpha = self.structure_attention_cell(
+            hidden, features, emb_feature
+        )
+
+        # structure
+        structure_step = self.structure_generator(output)
+        # loc
+        loc_step = self.loc_generator(output)
+        return hidden, structure_step, loc_step
+
+    def _char_to_onehot(self, input_char):
+        input_ont_hot = F.one_hot(input_char, self.num_embeddings)
+        return input_ont_hot

ppocr/modeling/heads/table_master_head.py

@@ -0,0 +1,285 @@
+# Copyright (c) 2021 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.
+"""
+This code is refer from:
+https://github.com/JiaquanYe/TableMASTER-mmocr/blob/master/mmocr/models/textrecog/decoders/master_decoder.py
+"""
+
+import copy
+import math
+import paddle
+from paddle import nn
+from paddle.nn import functional as F
+
+
+class TableMasterHead(nn.Layer):
+    """
+    Split to two transformer header at the last layer.
+    Cls_layer is used to structure token classification.
+    Bbox_layer is used to regress bbox coord.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels=30,
+        headers=8,
+        d_ff=2048,
+        dropout=0,
+        max_text_length=500,
+        loc_reg_num=4,
+        **kwargs,
+    ):
+        super(TableMasterHead, self).__init__()
+        hidden_size = in_channels[-1]
+        self.layers = clones(DecoderLayer(headers, hidden_size, dropout, d_ff), 2)
+        self.cls_layer = clones(DecoderLayer(headers, hidden_size, dropout, d_ff), 1)
+        self.bbox_layer = clones(DecoderLayer(headers, hidden_size, dropout, d_ff), 1)
+        self.cls_fc = nn.Linear(hidden_size, out_channels)
+        self.bbox_fc = nn.Sequential(
+            # nn.Linear(hidden_size, hidden_size),
+            nn.Linear(hidden_size, loc_reg_num),
+            nn.Sigmoid(),
+        )
+        self.norm = nn.LayerNorm(hidden_size)
+        self.embedding = Embeddings(d_model=hidden_size, vocab=out_channels)
+        self.positional_encoding = PositionalEncoding(d_model=hidden_size)
+
+        self.SOS = out_channels - 3
+        self.PAD = out_channels - 1
+        self.out_channels = out_channels
+        self.loc_reg_num = loc_reg_num
+        self.max_text_length = max_text_length
+
+    def make_mask(self, tgt):
+        """
+        Make mask for self attention.
+        :param src: [b, c, h, l_src]
+        :param tgt: [b, l_tgt]
+        :return:
+        """
+        trg_pad_mask = (tgt != self.PAD).unsqueeze(1).unsqueeze(3)
+
+        tgt_len = tgt.shape[1]
+        trg_sub_mask = paddle.tril(
+            paddle.ones(([tgt_len, tgt_len]), dtype=paddle.float32)
+        )
+
+        tgt_mask = paddle.logical_and(trg_pad_mask.astype(paddle.float32), trg_sub_mask)
+        return tgt_mask.astype(paddle.float32)
+
+    def decode(self, input, feature, src_mask, tgt_mask):
+        # main process of transformer decoder.
+        x = self.embedding(input)  # x: 1*x*512, feature: 1*3600,512
+        x = self.positional_encoding(x)
+
+        # origin transformer layers
+        for i, layer in enumerate(self.layers):
+            x = layer(x, feature, src_mask, tgt_mask)
+
+        # cls head
+        cls_x = x
+        for layer in self.cls_layer:
+            cls_x = layer(x, feature, src_mask, tgt_mask)
+        cls_x = self.norm(cls_x)
+
+        # bbox head
+        bbox_x = x
+        for layer in self.bbox_layer:
+            bbox_x = layer(x, feature, src_mask, tgt_mask)
+        bbox_x = self.norm(bbox_x)
+        return self.cls_fc(cls_x), self.bbox_fc(bbox_x)
+
+    def greedy_forward(self, SOS, feature):
+        input = SOS
+        output = paddle.zeros(
+            [input.shape[0], self.max_text_length + 1, self.out_channels]
+        )
+        bbox_output = paddle.zeros(
+            [input.shape[0], self.max_text_length + 1, self.loc_reg_num]
+        )
+        max_text_length = paddle.to_tensor(self.max_text_length)
+        for i in range(max_text_length + 1):
+            target_mask = self.make_mask(input)
+            out_step, bbox_output_step = self.decode(input, feature, None, target_mask)
+            prob = F.softmax(out_step, axis=-1)
+            next_word = prob.argmax(axis=2, dtype="int64")
+            input = paddle.concat([input, next_word[:, -1].unsqueeze(-1)], axis=1)
+            if i == self.max_text_length:
+                output = out_step
+                bbox_output = bbox_output_step
+        return output, bbox_output
+
+    def forward_train(self, out_enc, targets):
+        # x is token of label
+        # feat is feature after backbone before pe.
+        # out_enc is feature after pe.
+        padded_targets = targets[0]
+        src_mask = None
+        tgt_mask = self.make_mask(padded_targets[:, :-1])
+        output, bbox_output = self.decode(
+            padded_targets[:, :-1], out_enc, src_mask, tgt_mask
+        )
+        return {"structure_probs": output, "loc_preds": bbox_output}
+
+    def forward_test(self, out_enc):
+        batch_size = out_enc.shape[0]
+        SOS = paddle.zeros([batch_size, 1], dtype="int64") + self.SOS
+        output, bbox_output = self.greedy_forward(SOS, out_enc)
+        output = F.softmax(output)
+        return {"structure_probs": output, "loc_preds": bbox_output}
+
+    def forward(self, feat, targets=None):
+        feat = feat[-1]
+        b, c, h, w = feat.shape
+        feat = feat.reshape([b, c, h * w])  # flatten 2D feature map
+        feat = feat.transpose((0, 2, 1))
+        out_enc = self.positional_encoding(feat)
+        if self.training:
+            return self.forward_train(out_enc, targets)
+
+        return self.forward_test(out_enc)
+
+
+class DecoderLayer(nn.Layer):
+    """
+    Decoder is made of self attention, source attention and feed forward.
+    """
+
+    def __init__(self, headers, d_model, dropout, d_ff):
+        super(DecoderLayer, self).__init__()
+        self.self_attn = MultiHeadAttention(headers, d_model, dropout)
+        self.src_attn = MultiHeadAttention(headers, d_model, dropout)
+        self.feed_forward = FeedForward(d_model, d_ff, dropout)
+        self.sublayer = clones(SubLayerConnection(d_model, dropout), 3)
+
+    def forward(self, x, feature, src_mask, tgt_mask):
+        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
+        x = self.sublayer[1](x, lambda x: self.src_attn(x, feature, feature, src_mask))
+        return self.sublayer[2](x, self.feed_forward)
+
+
+class MultiHeadAttention(nn.Layer):
+    def __init__(self, headers, d_model, dropout):
+        super(MultiHeadAttention, self).__init__()
+
+        assert d_model % headers == 0
+        self.d_k = int(d_model / headers)
+        self.headers = headers
+        self.linears = clones(nn.Linear(d_model, d_model), 4)
+        self.attn = None
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, query, key, value, mask=None):
+        B = query.shape[0]
+
+        # 1) Do all the linear projections in batch from d_model => h x d_k
+        query, key, value = [
+            l(x).reshape([B, 0, self.headers, self.d_k]).transpose([0, 2, 1, 3])
+            for l, x in zip(self.linears, (query, key, value))
+        ]
+        # 2) Apply attention on all the projected vectors in batch
+        x, self.attn = self_attention(
+            query, key, value, mask=mask, dropout=self.dropout
+        )
+        x = x.transpose([0, 2, 1, 3]).reshape([B, 0, self.headers * self.d_k])
+        return self.linears[-1](x)
+
+
+class FeedForward(nn.Layer):
+    def __init__(self, d_model, d_ff, dropout):
+        super(FeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        return self.w_2(self.dropout(F.relu(self.w_1(x))))
+
+
+class SubLayerConnection(nn.Layer):
+    """
+    A residual connection followed by a layer norm.
+    Note for code simplicity the norm is first as opposed to last.
+    """
+
+    def __init__(self, size, dropout):
+        super(SubLayerConnection, self).__init__()
+        self.norm = nn.LayerNorm(size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, sublayer):
+        return x + self.dropout(sublayer(self.norm(x)))
+
+
+def masked_fill(x, mask, value):
+    mask = mask.astype(x.dtype)
+    return x * paddle.logical_not(mask).astype(x.dtype) + mask * value
+
+
+def self_attention(query, key, value, mask=None, dropout=None):
+    """
+    Compute 'Scale Dot Product Attention'
+    """
+    d_k = value.shape[-1]
+
+    score = paddle.matmul(query, key.transpose([0, 1, 3, 2]) / math.sqrt(d_k))
+    if mask is not None:
+        # score = score.masked_fill(mask == 0, -1e9) # b, h, L, L
+        score = masked_fill(score, mask == 0, -6.55e4)  # for fp16
+
+    p_attn = F.softmax(score, axis=-1)
+
+    if dropout is not None:
+        p_attn = dropout(p_attn)
+    return paddle.matmul(p_attn, value), p_attn
+
+
+def clones(module, N):
+    """Produce N identical layers"""
+    return nn.LayerList([copy.deepcopy(module) for _ in range(N)])
+
+
+class Embeddings(nn.Layer):
+    def __init__(self, d_model, vocab):
+        super(Embeddings, self).__init__()
+        self.lut = nn.Embedding(vocab, d_model)
+        self.d_model = d_model
+
+    def forward(self, *input):
+        x = input[0]
+        return self.lut(x) * math.sqrt(self.d_model)
+
+
+class PositionalEncoding(nn.Layer):
+    """Implement the PE function."""
+
+    def __init__(self, d_model, dropout=0.0, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        # Compute the positional encodings once in log space.
+        pe = paddle.zeros([max_len, d_model])
+        position = paddle.arange(0, max_len).unsqueeze(1).astype("float32")
+        div_term = paddle.exp(
+            paddle.arange(0, d_model, 2) * -math.log(10000.0) / d_model
+        )
+        pe[:, 0::2] = paddle.sin(position * div_term)
+        pe[:, 1::2] = paddle.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, feat, **kwargs):
+        feat = feat + self.pe[:, : feat.shape[1]]  # pe 1*5000*512
+        return self.dropout(feat)

ppocr/modeling/necks/__init__.py

@@ -0,0 +1,57 @@
+# 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_neck"]
+
+
+def build_neck(config):
+    from .db_fpn import DBFPN, RSEFPN, LKPAN
+    from .east_fpn import EASTFPN
+    from .sast_fpn import SASTFPN
+    from .rnn import SequenceEncoder
+    from .pg_fpn import PGFPN
+    from .table_fpn import TableFPN
+    from .fpn import FPN
+    from .fce_fpn import FCEFPN
+    from .pren_fpn import PRENFPN
+    from .csp_pan import CSPPAN
+    from .ct_fpn import CTFPN
+    from .fpn_unet import FPN_UNet
+    from .rf_adaptor import RFAdaptor
+
+    support_dict = [
+        "FPN",
+        "FCEFPN",
+        "LKPAN",
+        "DBFPN",
+        "RSEFPN",
+        "EASTFPN",
+        "SASTFPN",
+        "SequenceEncoder",
+        "PGFPN",
+        "TableFPN",
+        "PRENFPN",
+        "CSPPAN",
+        "CTFPN",
+        "RFAdaptor",
+        "FPN_UNet",
+    ]
+
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "neck only support {}".format(support_dict)
+    )
+
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/modeling/necks/csp_pan.py

@@ -0,0 +1,337 @@
+# Copyright (c) 2021 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.
+
+# The code is based on:
+# https://github.com/PaddlePaddle/PaddleDetection/blob/release%2F2.3/ppdet/modeling/necks/csp_pan.py
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+__all__ = ["CSPPAN"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channel=96,
+        out_channel=96,
+        kernel_size=3,
+        stride=1,
+        groups=1,
+        act="leaky_relu",
+    ):
+        super(ConvBNLayer, self).__init__()
+        initializer = nn.initializer.KaimingUniform()
+        self.act = act
+        assert self.act in ["leaky_relu", "hard_swish"]
+        self.conv = nn.Conv2D(
+            in_channels=in_channel,
+            out_channels=out_channel,
+            kernel_size=kernel_size,
+            groups=groups,
+            padding=(kernel_size - 1) // 2,
+            stride=stride,
+            weight_attr=ParamAttr(initializer=initializer),
+            bias_attr=False,
+        )
+        self.bn = nn.BatchNorm2D(out_channel)
+
+    def forward(self, x):
+        x = self.bn(self.conv(x))
+        if self.act == "leaky_relu":
+            x = F.leaky_relu(x)
+        elif self.act == "hard_swish":
+            x = F.hardswish(x)
+        return x
+
+
+class DPModule(nn.Layer):
+    """
+    Depth-wise and point-wise module.
+     Args:
+        in_channel (int): The input channels of this Module.
+        out_channel (int): The output channels of this Module.
+        kernel_size (int): The conv2d kernel size of this Module.
+        stride (int): The conv2d's stride of this Module.
+        act (str): The activation function of this Module,
+                   Now support `leaky_relu` and `hard_swish`.
+    """
+
+    def __init__(
+        self, in_channel=96, out_channel=96, kernel_size=3, stride=1, act="leaky_relu"
+    ):
+        super(DPModule, self).__init__()
+        initializer = nn.initializer.KaimingUniform()
+        self.act = act
+        self.dwconv = nn.Conv2D(
+            in_channels=in_channel,
+            out_channels=out_channel,
+            kernel_size=kernel_size,
+            groups=out_channel,
+            padding=(kernel_size - 1) // 2,
+            stride=stride,
+            weight_attr=ParamAttr(initializer=initializer),
+            bias_attr=False,
+        )
+        self.bn1 = nn.BatchNorm2D(out_channel)
+        self.pwconv = nn.Conv2D(
+            in_channels=out_channel,
+            out_channels=out_channel,
+            kernel_size=1,
+            groups=1,
+            padding=0,
+            weight_attr=ParamAttr(initializer=initializer),
+            bias_attr=False,
+        )
+        self.bn2 = nn.BatchNorm2D(out_channel)
+
+    def act_func(self, x):
+        if self.act == "leaky_relu":
+            x = F.leaky_relu(x)
+        elif self.act == "hard_swish":
+            x = F.hardswish(x)
+        return x
+
+    def forward(self, x):
+        x = self.act_func(self.bn1(self.dwconv(x)))
+        x = self.act_func(self.bn2(self.pwconv(x)))
+        return x
+
+
+class DarknetBottleneck(nn.Layer):
+    """The basic bottleneck block used in Darknet.
+    Each Block consists of two ConvModules and the input is added to the
+    final output. Each ConvModule is composed of Conv, BN, and act.
+    The first convLayer has filter size of 1x1 and the second one has the
+    filter size of 3x3.
+    Args:
+        in_channels (int): The input channels of this Module.
+        out_channels (int): The output channels of this Module.
+        expansion (int): The kernel size of the convolution. Default: 0.5
+        add_identity (bool): Whether to add identity to the out.
+            Default: True
+        use_depthwise (bool): Whether to use depthwise separable convolution.
+            Default: False
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        expansion=0.5,
+        add_identity=True,
+        use_depthwise=False,
+        act="leaky_relu",
+    ):
+        super(DarknetBottleneck, self).__init__()
+        hidden_channels = int(out_channels * expansion)
+        conv_func = DPModule if use_depthwise else ConvBNLayer
+        self.conv1 = ConvBNLayer(
+            in_channel=in_channels, out_channel=hidden_channels, kernel_size=1, act=act
+        )
+        self.conv2 = conv_func(
+            in_channel=hidden_channels,
+            out_channel=out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            act=act,
+        )
+        self.add_identity = add_identity and in_channels == out_channels
+
+    def forward(self, x):
+        identity = x
+        out = self.conv1(x)
+        out = self.conv2(out)
+
+        if self.add_identity:
+            return out + identity
+        else:
+            return out
+
+
+class CSPLayer(nn.Layer):
+    """Cross Stage Partial Layer.
+    Args:
+        in_channels (int): The input channels of the CSP layer.
+        out_channels (int): The output channels of the CSP layer.
+        expand_ratio (float): Ratio to adjust the number of channels of the
+            hidden layer. Default: 0.5
+        num_blocks (int): Number of blocks. Default: 1
+        add_identity (bool): Whether to add identity in blocks.
+            Default: True
+        use_depthwise (bool): Whether to depthwise separable convolution in
+            blocks. Default: False
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        expand_ratio=0.5,
+        num_blocks=1,
+        add_identity=True,
+        use_depthwise=False,
+        act="leaky_relu",
+    ):
+        super().__init__()
+        mid_channels = int(out_channels * expand_ratio)
+        self.main_conv = ConvBNLayer(in_channels, mid_channels, 1, act=act)
+        self.short_conv = ConvBNLayer(in_channels, mid_channels, 1, act=act)
+        self.final_conv = ConvBNLayer(2 * mid_channels, out_channels, 1, act=act)
+
+        self.blocks = nn.Sequential(
+            *[
+                DarknetBottleneck(
+                    mid_channels,
+                    mid_channels,
+                    kernel_size,
+                    1.0,
+                    add_identity,
+                    use_depthwise,
+                    act=act,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+
+    def forward(self, x):
+        x_short = self.short_conv(x)
+
+        x_main = self.main_conv(x)
+        x_main = self.blocks(x_main)
+
+        x_final = paddle.concat((x_main, x_short), axis=1)
+        return self.final_conv(x_final)
+
+
+class Channel_T(nn.Layer):
+    def __init__(self, in_channels=[116, 232, 464], out_channels=96, act="leaky_relu"):
+        super(Channel_T, self).__init__()
+        self.convs = nn.LayerList()
+        for i in range(len(in_channels)):
+            self.convs.append(ConvBNLayer(in_channels[i], out_channels, 1, act=act))
+
+    def forward(self, x):
+        outs = [self.convs[i](x[i]) for i in range(len(x))]
+        return outs
+
+
+class CSPPAN(nn.Layer):
+    """Path Aggregation Network with CSP module.
+    Args:
+        in_channels (List[int]): Number of input channels per scale.
+        out_channels (int): Number of output channels (used at each scale)
+        kernel_size (int): The conv2d kernel size of this Module.
+        num_csp_blocks (int): Number of bottlenecks in CSPLayer. Default: 1
+        use_depthwise (bool): Whether to depthwise separable convolution in
+            blocks. Default: True
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=5,
+        num_csp_blocks=1,
+        use_depthwise=True,
+        act="hard_swish",
+    ):
+        super(CSPPAN, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = [out_channels] * len(in_channels)
+        conv_func = DPModule if use_depthwise else ConvBNLayer
+
+        self.conv_t = Channel_T(in_channels, out_channels, act=act)
+
+        # build top-down blocks
+        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
+        self.top_down_blocks = nn.LayerList()
+        for idx in range(len(in_channels) - 1, 0, -1):
+            self.top_down_blocks.append(
+                CSPLayer(
+                    out_channels * 2,
+                    out_channels,
+                    kernel_size=kernel_size,
+                    num_blocks=num_csp_blocks,
+                    add_identity=False,
+                    use_depthwise=use_depthwise,
+                    act=act,
+                )
+            )
+
+        # build bottom-up blocks
+        self.downsamples = nn.LayerList()
+        self.bottom_up_blocks = nn.LayerList()
+        for idx in range(len(in_channels) - 1):
+            self.downsamples.append(
+                conv_func(
+                    out_channels,
+                    out_channels,
+                    kernel_size=kernel_size,
+                    stride=2,
+                    act=act,
+                )
+            )
+            self.bottom_up_blocks.append(
+                CSPLayer(
+                    out_channels * 2,
+                    out_channels,
+                    kernel_size=kernel_size,
+                    num_blocks=num_csp_blocks,
+                    add_identity=False,
+                    use_depthwise=use_depthwise,
+                    act=act,
+                )
+            )
+
+    def forward(self, inputs):
+        """
+        Args:
+            inputs (tuple[Tensor]): input features.
+        Returns:
+            tuple[Tensor]: CSPPAN features.
+        """
+        assert len(inputs) == len(self.in_channels)
+        inputs = self.conv_t(inputs)
+
+        # top-down path
+        inner_outs = [inputs[-1]]
+        for idx in range(len(self.in_channels) - 1, 0, -1):
+            feat_heigh = inner_outs[0]
+            feat_low = inputs[idx - 1]
+            upsample_feat = F.upsample(
+                feat_heigh, size=feat_low.shape[2:4], mode="nearest"
+            )
+
+            inner_out = self.top_down_blocks[len(self.in_channels) - 1 - idx](
+                paddle.concat([upsample_feat, feat_low], 1)
+            )
+            inner_outs.insert(0, inner_out)
+
+        # bottom-up path
+        outs = [inner_outs[0]]
+        for idx in range(len(self.in_channels) - 1):
+            feat_low = outs[-1]
+            feat_height = inner_outs[idx + 1]
+            downsample_feat = self.downsamples[idx](feat_low)
+            out = self.bottom_up_blocks[idx](
+                paddle.concat([downsample_feat, feat_height], 1)
+            )
+            outs.append(out)
+
+        return tuple(outs)

ppocr/modeling/necks/ct_fpn.py

@@ -0,0 +1,188 @@
+# copyright (c) 2019 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+import os
+import sys
+
+import math
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal
+
+ones_ = Constant(value=1.0)
+zeros_ = Constant(value=0.0)
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../..")))
+
+
+class Conv_BN_ReLU(nn.Layer):
+    def __init__(self, in_planes, out_planes, kernel_size=1, stride=1, padding=0):
+        super(Conv_BN_ReLU, self).__init__()
+        self.conv = nn.Conv2D(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias_attr=False,
+        )
+        self.bn = nn.BatchNorm2D(out_planes)
+        self.relu = nn.ReLU()
+
+        for m in self.sublayers():
+            if isinstance(m, nn.Conv2D):
+                n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels
+                normal_ = Normal(mean=0.0, std=math.sqrt(2.0 / n))
+                normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2D):
+                zeros_(m.bias)
+                ones_(m.weight)
+
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+
+class FPEM(nn.Layer):
+    def __init__(self, in_channels, out_channels):
+        super(FPEM, self).__init__()
+        planes = out_channels
+        self.dwconv3_1 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer3_1 = Conv_BN_ReLU(planes, planes)
+
+        self.dwconv2_1 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer2_1 = Conv_BN_ReLU(planes, planes)
+
+        self.dwconv1_1 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer1_1 = Conv_BN_ReLU(planes, planes)
+
+        self.dwconv2_2 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer2_2 = Conv_BN_ReLU(planes, planes)
+
+        self.dwconv3_2 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer3_2 = Conv_BN_ReLU(planes, planes)
+
+        self.dwconv4_2 = nn.Conv2D(
+            planes,
+            planes,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            groups=planes,
+            bias_attr=False,
+        )
+        self.smooth_layer4_2 = Conv_BN_ReLU(planes, planes)
+
+    def _upsample_add(self, x, y):
+        return F.upsample(x, scale_factor=2, mode="bilinear") + y
+
+    def forward(self, f1, f2, f3, f4):
+        # up-down
+        f3 = self.smooth_layer3_1(self.dwconv3_1(self._upsample_add(f4, f3)))
+        f2 = self.smooth_layer2_1(self.dwconv2_1(self._upsample_add(f3, f2)))
+        f1 = self.smooth_layer1_1(self.dwconv1_1(self._upsample_add(f2, f1)))
+
+        # down-up
+        f2 = self.smooth_layer2_2(self.dwconv2_2(self._upsample_add(f2, f1)))
+        f3 = self.smooth_layer3_2(self.dwconv3_2(self._upsample_add(f3, f2)))
+        f4 = self.smooth_layer4_2(self.dwconv4_2(self._upsample_add(f4, f3)))
+
+        return f1, f2, f3, f4
+
+
+class CTFPN(nn.Layer):
+    def __init__(self, in_channels, out_channel=128):
+        super(CTFPN, self).__init__()
+        self.out_channels = out_channel * 4
+
+        self.reduce_layer1 = Conv_BN_ReLU(in_channels[0], 128)
+        self.reduce_layer2 = Conv_BN_ReLU(in_channels[1], 128)
+        self.reduce_layer3 = Conv_BN_ReLU(in_channels[2], 128)
+        self.reduce_layer4 = Conv_BN_ReLU(in_channels[3], 128)
+
+        self.fpem1 = FPEM(in_channels=(64, 128, 256, 512), out_channels=128)
+        self.fpem2 = FPEM(in_channels=(64, 128, 256, 512), out_channels=128)
+
+    def _upsample(self, x, scale=1):
+        return F.upsample(x, scale_factor=scale, mode="bilinear")
+
+    def forward(self, f):
+        # # reduce channel
+        f1 = self.reduce_layer1(f[0])  # N,64,160,160    --> N, 128, 160, 160
+        f2 = self.reduce_layer2(f[1])  # N, 128, 80, 80  --> N, 128, 80, 80
+        f3 = self.reduce_layer3(f[2])  # N, 256, 40, 40  --> N, 128, 40, 40
+        f4 = self.reduce_layer4(f[3])  # N, 512, 20, 20  --> N, 128, 20, 20
+
+        # FPEM
+        f1_1, f2_1, f3_1, f4_1 = self.fpem1(f1, f2, f3, f4)
+        f1_2, f2_2, f3_2, f4_2 = self.fpem2(f1_1, f2_1, f3_1, f4_1)
+
+        # FFM
+        f1 = f1_1 + f1_2
+        f2 = f2_1 + f2_2
+        f3 = f3_1 + f3_2
+        f4 = f4_1 + f4_2
+
+        f2 = self._upsample(f2, scale=2)
+        f3 = self._upsample(f3, scale=4)
+        f4 = self._upsample(f4, scale=8)
+        ff = paddle.concat((f1, f2, f3, f4), 1)  # N,512, 160,160
+        return ff

ppocr/modeling/necks/db_fpn.py

@@ -0,0 +1,492 @@
+# copyright (c) 2019 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+import os
+import sys
+from ppocr.modeling.necks.intracl import IntraCLBlock
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(__dir__)
+sys.path.insert(0, os.path.abspath(os.path.join(__dir__, "../../..")))
+
+from ppocr.modeling.backbones.det_mobilenet_v3 import SEModule
+
+
+class DSConv(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        padding,
+        stride=1,
+        groups=None,
+        if_act=True,
+        act="relu",
+        **kwargs,
+    ):
+        super(DSConv, self).__init__()
+        if groups == None:
+            groups = in_channels
+        self.if_act = if_act
+        self.act = act
+        self.conv1 = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            bias_attr=False,
+        )
+
+        self.bn1 = nn.BatchNorm(num_channels=in_channels, act=None)
+
+        self.conv2 = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=int(in_channels * 4),
+            kernel_size=1,
+            stride=1,
+            bias_attr=False,
+        )
+
+        self.bn2 = nn.BatchNorm(num_channels=int(in_channels * 4), act=None)
+
+        self.conv3 = nn.Conv2D(
+            in_channels=int(in_channels * 4),
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            bias_attr=False,
+        )
+        self._c = [in_channels, out_channels]
+        if in_channels != out_channels:
+            self.conv_end = nn.Conv2D(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                bias_attr=False,
+            )
+
+    def forward(self, inputs):
+        x = self.conv1(inputs)
+        x = self.bn1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        if self.if_act:
+            if self.act == "relu":
+                x = F.relu(x)
+            elif self.act == "hardswish":
+                x = F.hardswish(x)
+            else:
+                print(
+                    "The activation function({}) is selected incorrectly.".format(
+                        self.act
+                    )
+                )
+                exit()
+
+        x = self.conv3(x)
+        if self._c[0] != self._c[1]:
+            x = x + self.conv_end(inputs)
+        return x
+
+
+class DBFPN(nn.Layer):
+    def __init__(self, in_channels, out_channels, use_asf=False, **kwargs):
+        super(DBFPN, self).__init__()
+        self.out_channels = out_channels
+        self.use_asf = use_asf
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+
+        self.in2_conv = nn.Conv2D(
+            in_channels=in_channels[0],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in3_conv = nn.Conv2D(
+            in_channels=in_channels[1],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in4_conv = nn.Conv2D(
+            in_channels=in_channels[2],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in5_conv = nn.Conv2D(
+            in_channels=in_channels[3],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p5_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p4_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p3_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p2_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+
+        if self.use_asf is True:
+            self.asf = ASFBlock(self.out_channels, self.out_channels // 4)
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.in5_conv(c5)
+        in4 = self.in4_conv(c4)
+        in3 = self.in3_conv(c3)
+        in2 = self.in2_conv(c2)
+
+        out4 = in4 + F.upsample(
+            in5, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/16
+        out3 = in3 + F.upsample(
+            out4, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/8
+        out2 = in2 + F.upsample(
+            out3, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/4
+
+        p5 = self.p5_conv(in5)
+        p4 = self.p4_conv(out4)
+        p3 = self.p3_conv(out3)
+        p2 = self.p2_conv(out2)
+        p5 = F.upsample(p5, scale_factor=8, mode="nearest", align_mode=1)
+        p4 = F.upsample(p4, scale_factor=4, mode="nearest", align_mode=1)
+        p3 = F.upsample(p3, scale_factor=2, mode="nearest", align_mode=1)
+
+        fuse = paddle.concat([p5, p4, p3, p2], axis=1)
+
+        if self.use_asf is True:
+            fuse = self.asf(fuse, [p5, p4, p3, p2])
+
+        return fuse
+
+
+class RSELayer(nn.Layer):
+    def __init__(self, in_channels, out_channels, kernel_size, shortcut=True):
+        super(RSELayer, self).__init__()
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+        self.out_channels = out_channels
+        self.in_conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=self.out_channels,
+            kernel_size=kernel_size,
+            padding=int(kernel_size // 2),
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.se_block = SEModule(self.out_channels)
+        self.shortcut = shortcut
+
+    def forward(self, ins):
+        x = self.in_conv(ins)
+        if self.shortcut:
+            out = x + self.se_block(x)
+        else:
+            out = self.se_block(x)
+        return out
+
+
+class RSEFPN(nn.Layer):
+    def __init__(self, in_channels, out_channels, shortcut=True, **kwargs):
+        super(RSEFPN, self).__init__()
+        self.out_channels = out_channels
+        self.ins_conv = nn.LayerList()
+        self.inp_conv = nn.LayerList()
+        self.intracl = False
+        if "intracl" in kwargs.keys() and kwargs["intracl"] is True:
+            self.intracl = kwargs["intracl"]
+            self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+
+        for i in range(len(in_channels)):
+            self.ins_conv.append(
+                RSELayer(in_channels[i], out_channels, kernel_size=1, shortcut=shortcut)
+            )
+            self.inp_conv.append(
+                RSELayer(
+                    out_channels, out_channels // 4, kernel_size=3, shortcut=shortcut
+                )
+            )
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.ins_conv[3](c5)
+        in4 = self.ins_conv[2](c4)
+        in3 = self.ins_conv[1](c3)
+        in2 = self.ins_conv[0](c2)
+
+        out4 = in4 + F.upsample(
+            in5, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/16
+        out3 = in3 + F.upsample(
+            out4, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/8
+        out2 = in2 + F.upsample(
+            out3, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/4
+
+        p5 = self.inp_conv[3](in5)
+        p4 = self.inp_conv[2](out4)
+        p3 = self.inp_conv[1](out3)
+        p2 = self.inp_conv[0](out2)
+
+        if self.intracl is True:
+            p5 = self.incl4(p5)
+            p4 = self.incl3(p4)
+            p3 = self.incl2(p3)
+            p2 = self.incl1(p2)
+
+        p5 = F.upsample(p5, scale_factor=8, mode="nearest", align_mode=1)
+        p4 = F.upsample(p4, scale_factor=4, mode="nearest", align_mode=1)
+        p3 = F.upsample(p3, scale_factor=2, mode="nearest", align_mode=1)
+
+        fuse = paddle.concat([p5, p4, p3, p2], axis=1)
+        return fuse
+
+
+class LKPAN(nn.Layer):
+    def __init__(self, in_channels, out_channels, mode="large", **kwargs):
+        super(LKPAN, self).__init__()
+        self.out_channels = out_channels
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+
+        self.ins_conv = nn.LayerList()
+        self.inp_conv = nn.LayerList()
+        # pan head
+        self.pan_head_conv = nn.LayerList()
+        self.pan_lat_conv = nn.LayerList()
+
+        if mode.lower() == "lite":
+            p_layer = DSConv
+        elif mode.lower() == "large":
+            p_layer = nn.Conv2D
+        else:
+            raise ValueError(
+                "mode can only be one of ['lite', 'large'], but received {}".format(
+                    mode
+                )
+            )
+
+        for i in range(len(in_channels)):
+            self.ins_conv.append(
+                nn.Conv2D(
+                    in_channels=in_channels[i],
+                    out_channels=self.out_channels,
+                    kernel_size=1,
+                    weight_attr=ParamAttr(initializer=weight_attr),
+                    bias_attr=False,
+                )
+            )
+
+            self.inp_conv.append(
+                p_layer(
+                    in_channels=self.out_channels,
+                    out_channels=self.out_channels // 4,
+                    kernel_size=9,
+                    padding=4,
+                    weight_attr=ParamAttr(initializer=weight_attr),
+                    bias_attr=False,
+                )
+            )
+
+            if i > 0:
+                self.pan_head_conv.append(
+                    nn.Conv2D(
+                        in_channels=self.out_channels // 4,
+                        out_channels=self.out_channels // 4,
+                        kernel_size=3,
+                        padding=1,
+                        stride=2,
+                        weight_attr=ParamAttr(initializer=weight_attr),
+                        bias_attr=False,
+                    )
+                )
+            self.pan_lat_conv.append(
+                p_layer(
+                    in_channels=self.out_channels // 4,
+                    out_channels=self.out_channels // 4,
+                    kernel_size=9,
+                    padding=4,
+                    weight_attr=ParamAttr(initializer=weight_attr),
+                    bias_attr=False,
+                )
+            )
+
+        self.intracl = False
+        if "intracl" in kwargs.keys() and kwargs["intracl"] is True:
+            self.intracl = kwargs["intracl"]
+            self.incl1 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl2 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl3 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+            self.incl4 = IntraCLBlock(self.out_channels // 4, reduce_factor=2)
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.ins_conv[3](c5)
+        in4 = self.ins_conv[2](c4)
+        in3 = self.ins_conv[1](c3)
+        in2 = self.ins_conv[0](c2)
+
+        out4 = in4 + F.upsample(
+            in5, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/16
+        out3 = in3 + F.upsample(
+            out4, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/8
+        out2 = in2 + F.upsample(
+            out3, scale_factor=2, mode="nearest", align_mode=1
+        )  # 1/4
+
+        f5 = self.inp_conv[3](in5)
+        f4 = self.inp_conv[2](out4)
+        f3 = self.inp_conv[1](out3)
+        f2 = self.inp_conv[0](out2)
+
+        pan3 = f3 + self.pan_head_conv[0](f2)
+        pan4 = f4 + self.pan_head_conv[1](pan3)
+        pan5 = f5 + self.pan_head_conv[2](pan4)
+
+        p2 = self.pan_lat_conv[0](f2)
+        p3 = self.pan_lat_conv[1](pan3)
+        p4 = self.pan_lat_conv[2](pan4)
+        p5 = self.pan_lat_conv[3](pan5)
+
+        if self.intracl is True:
+            p5 = self.incl4(p5)
+            p4 = self.incl3(p4)
+            p3 = self.incl2(p3)
+            p2 = self.incl1(p2)
+
+        p5 = F.upsample(p5, scale_factor=8, mode="nearest", align_mode=1)
+        p4 = F.upsample(p4, scale_factor=4, mode="nearest", align_mode=1)
+        p3 = F.upsample(p3, scale_factor=2, mode="nearest", align_mode=1)
+
+        fuse = paddle.concat([p5, p4, p3, p2], axis=1)
+        return fuse
+
+
+class ASFBlock(nn.Layer):
+    """
+    This code is referred from:
+        https://github.com/MhLiao/DB/blob/master/decoders/feature_attention.py
+    """
+
+    def __init__(self, in_channels, inter_channels, out_features_num=4):
+        """
+        Adaptive Scale Fusion (ASF) block of DBNet++
+        Args:
+            in_channels: the number of channels in the input data
+            inter_channels: the number of middle channels
+            out_features_num: the number of fused stages
+        """
+        super(ASFBlock, self).__init__()
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+        self.in_channels = in_channels
+        self.inter_channels = inter_channels
+        self.out_features_num = out_features_num
+        self.conv = nn.Conv2D(in_channels, inter_channels, 3, padding=1)
+
+        self.spatial_scale = nn.Sequential(
+            # Nx1xHxW
+            nn.Conv2D(
+                in_channels=1,
+                out_channels=1,
+                kernel_size=3,
+                bias_attr=False,
+                padding=1,
+                weight_attr=ParamAttr(initializer=weight_attr),
+            ),
+            nn.ReLU(),
+            nn.Conv2D(
+                in_channels=1,
+                out_channels=1,
+                kernel_size=1,
+                bias_attr=False,
+                weight_attr=ParamAttr(initializer=weight_attr),
+            ),
+            nn.Sigmoid(),
+        )
+
+        self.channel_scale = nn.Sequential(
+            nn.Conv2D(
+                in_channels=inter_channels,
+                out_channels=out_features_num,
+                kernel_size=1,
+                bias_attr=False,
+                weight_attr=ParamAttr(initializer=weight_attr),
+            ),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, fuse_features, features_list):
+        fuse_features = self.conv(fuse_features)
+        spatial_x = paddle.mean(fuse_features, axis=1, keepdim=True)
+        attention_scores = self.spatial_scale(spatial_x) + fuse_features
+        attention_scores = self.channel_scale(attention_scores)
+        assert len(features_list) == self.out_features_num
+
+        out_list = []
+        for i in range(self.out_features_num):
+            out_list.append(attention_scores[:, i : i + 1] * features_list[i])
+        return paddle.concat(out_list, axis=1)

ppocr/modeling/necks/east_fpn.py

@@ -0,0 +1,203 @@
+# copyright (c) 2019 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="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 DeConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(DeConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.deconv = nn.Conv2DTranspose(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+        )
+
+    def forward(self, x):
+        x = self.deconv(x)
+        x = self.bn(x)
+        return x
+
+
+class EASTFPN(nn.Layer):
+    def __init__(self, in_channels, model_name, **kwargs):
+        super(EASTFPN, self).__init__()
+        self.model_name = model_name
+        if self.model_name == "large":
+            self.out_channels = 128
+        else:
+            self.out_channels = 64
+        self.in_channels = in_channels[::-1]
+        self.h1_conv = ConvBNLayer(
+            in_channels=self.out_channels + self.in_channels[1],
+            out_channels=self.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_h_1",
+        )
+        self.h2_conv = ConvBNLayer(
+            in_channels=self.out_channels + self.in_channels[2],
+            out_channels=self.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_h_2",
+        )
+        self.h3_conv = ConvBNLayer(
+            in_channels=self.out_channels + self.in_channels[3],
+            out_channels=self.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_h_3",
+        )
+        self.g0_deconv = DeConvBNLayer(
+            in_channels=self.in_channels[0],
+            out_channels=self.out_channels,
+            kernel_size=4,
+            stride=2,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_g_0",
+        )
+        self.g1_deconv = DeConvBNLayer(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels,
+            kernel_size=4,
+            stride=2,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_g_1",
+        )
+        self.g2_deconv = DeConvBNLayer(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels,
+            kernel_size=4,
+            stride=2,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_g_2",
+        )
+        self.g3_conv = ConvBNLayer(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            if_act=True,
+            act="relu",
+            name="unet_g_3",
+        )
+
+    def forward(self, x):
+        f = x[::-1]
+
+        h = f[0]
+        g = self.g0_deconv(h)
+        h = paddle.concat([g, f[1]], axis=1)
+        h = self.h1_conv(h)
+        g = self.g1_deconv(h)
+        h = paddle.concat([g, f[2]], axis=1)
+        h = self.h2_conv(h)
+        g = self.g2_deconv(h)
+        h = paddle.concat([g, f[3]], axis=1)
+        h = self.h3_conv(h)
+        g = self.g3_conv(h)
+
+        return g

ppocr/modeling/necks/fce_fpn.py

@@ -0,0 +1,304 @@
+# 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/PaddlePaddle/PaddleDetection/blob/release/2.3/ppdet/modeling/necks/fpn.py
+"""
+
+import paddle.nn as nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+from paddle.nn.initializer import XavierUniform
+from paddle.nn.initializer import Normal
+from paddle.regularizer import L2Decay
+
+__all__ = ["FCEFPN"]
+
+
+class ConvNormLayer(nn.Layer):
+    def __init__(
+        self,
+        ch_in,
+        ch_out,
+        filter_size,
+        stride,
+        groups=1,
+        norm_type="bn",
+        norm_decay=0.0,
+        norm_groups=32,
+        lr_scale=1.0,
+        freeze_norm=False,
+        initializer=Normal(mean=0.0, std=0.01),
+    ):
+        super(ConvNormLayer, self).__init__()
+        assert norm_type in ["bn", "sync_bn", "gn"]
+
+        bias_attr = False
+
+        self.conv = nn.Conv2D(
+            in_channels=ch_in,
+            out_channels=ch_out,
+            kernel_size=filter_size,
+            stride=stride,
+            padding=(filter_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(initializer=initializer, learning_rate=1.0),
+            bias_attr=bias_attr,
+        )
+
+        norm_lr = 0.0 if freeze_norm else 1.0
+        param_attr = ParamAttr(
+            learning_rate=norm_lr,
+            regularizer=L2Decay(norm_decay) if norm_decay is not None else None,
+        )
+        bias_attr = ParamAttr(
+            learning_rate=norm_lr,
+            regularizer=L2Decay(norm_decay) if norm_decay is not None else None,
+        )
+        if norm_type == "bn":
+            self.norm = nn.BatchNorm2D(
+                ch_out, weight_attr=param_attr, bias_attr=bias_attr
+            )
+        elif norm_type == "sync_bn":
+            self.norm = nn.SyncBatchNorm(
+                ch_out, weight_attr=param_attr, bias_attr=bias_attr
+            )
+        elif norm_type == "gn":
+            self.norm = nn.GroupNorm(
+                num_groups=norm_groups,
+                num_channels=ch_out,
+                weight_attr=param_attr,
+                bias_attr=bias_attr,
+            )
+
+    def forward(self, inputs):
+        out = self.conv(inputs)
+        out = self.norm(out)
+        return out
+
+
+class FCEFPN(nn.Layer):
+    """
+    Feature Pyramid Network, see https://arxiv.org/abs/1612.03144
+    Args:
+        in_channels (list[int]): input channels of each level which can be
+            derived from the output shape of backbone by from_config
+        out_channels (list[int]): output channel of each level
+        spatial_scales (list[float]): the spatial scales between input feature
+            maps and original input image which can be derived from the output
+            shape of backbone by from_config
+        has_extra_convs (bool): whether to add extra conv to the last level.
+            default False
+        extra_stage (int): the number of extra stages added to the last level.
+            default 1
+        use_c5 (bool): Whether to use c5 as the input of extra stage,
+            otherwise p5 is used. default True
+        norm_type (string|None): The normalization type in FPN module. If
+            norm_type is None, norm will not be used after conv and if
+            norm_type is string, bn, gn, sync_bn are available. default None
+        norm_decay (float): weight decay for normalization layer weights.
+            default 0.
+        freeze_norm (bool): whether to freeze normalization layer.
+            default False
+        relu_before_extra_convs (bool): whether to add relu before extra convs.
+            default False
+
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        spatial_scales=[0.25, 0.125, 0.0625, 0.03125],
+        has_extra_convs=False,
+        extra_stage=1,
+        use_c5=True,
+        norm_type=None,
+        norm_decay=0.0,
+        freeze_norm=False,
+        relu_before_extra_convs=True,
+    ):
+        super(FCEFPN, self).__init__()
+        self.out_channels = out_channels
+        for s in range(extra_stage):
+            spatial_scales = spatial_scales + [spatial_scales[-1] / 2.0]
+        self.spatial_scales = spatial_scales
+        self.has_extra_convs = has_extra_convs
+        self.extra_stage = extra_stage
+        self.use_c5 = use_c5
+        self.relu_before_extra_convs = relu_before_extra_convs
+        self.norm_type = norm_type
+        self.norm_decay = norm_decay
+        self.freeze_norm = freeze_norm
+
+        self.lateral_convs = []
+        self.fpn_convs = []
+        fan = out_channels * 3 * 3
+
+        # stage index 0,1,2,3 stands for res2,res3,res4,res5 on ResNet Backbone
+        # 0 <= st_stage < ed_stage <= 3
+        st_stage = 4 - len(in_channels)
+        ed_stage = st_stage + len(in_channels) - 1
+        for i in range(st_stage, ed_stage + 1):
+            if i == 3:
+                lateral_name = "fpn_inner_res5_sum"
+            else:
+                lateral_name = "fpn_inner_res{}_sum_lateral".format(i + 2)
+            in_c = in_channels[i - st_stage]
+            if self.norm_type is not None:
+                lateral = self.add_sublayer(
+                    lateral_name,
+                    ConvNormLayer(
+                        ch_in=in_c,
+                        ch_out=out_channels,
+                        filter_size=1,
+                        stride=1,
+                        norm_type=self.norm_type,
+                        norm_decay=self.norm_decay,
+                        freeze_norm=self.freeze_norm,
+                        initializer=XavierUniform(fan_out=in_c),
+                    ),
+                )
+            else:
+                lateral = self.add_sublayer(
+                    lateral_name,
+                    nn.Conv2D(
+                        in_channels=in_c,
+                        out_channels=out_channels,
+                        kernel_size=1,
+                        weight_attr=ParamAttr(initializer=XavierUniform(fan_out=in_c)),
+                    ),
+                )
+            self.lateral_convs.append(lateral)
+
+        for i in range(st_stage, ed_stage + 1):
+            fpn_name = "fpn_res{}_sum".format(i + 2)
+            if self.norm_type is not None:
+                fpn_conv = self.add_sublayer(
+                    fpn_name,
+                    ConvNormLayer(
+                        ch_in=out_channels,
+                        ch_out=out_channels,
+                        filter_size=3,
+                        stride=1,
+                        norm_type=self.norm_type,
+                        norm_decay=self.norm_decay,
+                        freeze_norm=self.freeze_norm,
+                        initializer=XavierUniform(fan_out=fan),
+                    ),
+                )
+            else:
+                fpn_conv = self.add_sublayer(
+                    fpn_name,
+                    nn.Conv2D(
+                        in_channels=out_channels,
+                        out_channels=out_channels,
+                        kernel_size=3,
+                        padding=1,
+                        weight_attr=ParamAttr(initializer=XavierUniform(fan_out=fan)),
+                    ),
+                )
+            self.fpn_convs.append(fpn_conv)
+
+        # add extra conv levels for RetinaNet(use_c5)/FCOS(use_p5)
+        if self.has_extra_convs:
+            for i in range(self.extra_stage):
+                lvl = ed_stage + 1 + i
+                if i == 0 and self.use_c5:
+                    in_c = in_channels[-1]
+                else:
+                    in_c = out_channels
+                extra_fpn_name = "fpn_{}".format(lvl + 2)
+                if self.norm_type is not None:
+                    extra_fpn_conv = self.add_sublayer(
+                        extra_fpn_name,
+                        ConvNormLayer(
+                            ch_in=in_c,
+                            ch_out=out_channels,
+                            filter_size=3,
+                            stride=2,
+                            norm_type=self.norm_type,
+                            norm_decay=self.norm_decay,
+                            freeze_norm=self.freeze_norm,
+                            initializer=XavierUniform(fan_out=fan),
+                        ),
+                    )
+                else:
+                    extra_fpn_conv = self.add_sublayer(
+                        extra_fpn_name,
+                        nn.Conv2D(
+                            in_channels=in_c,
+                            out_channels=out_channels,
+                            kernel_size=3,
+                            stride=2,
+                            padding=1,
+                            weight_attr=ParamAttr(
+                                initializer=XavierUniform(fan_out=fan)
+                            ),
+                        ),
+                    )
+                self.fpn_convs.append(extra_fpn_conv)
+
+    @classmethod
+    def from_config(cls, cfg, input_shape):
+        return {
+            "in_channels": [i.channels for i in input_shape],
+            "spatial_scales": [1.0 / i.stride for i in input_shape],
+        }
+
+    def forward(self, body_feats):
+        laterals = []
+        num_levels = len(body_feats)
+
+        for i in range(num_levels):
+            laterals.append(self.lateral_convs[i](body_feats[i]))
+
+        for i in range(1, num_levels):
+            lvl = num_levels - i
+            upsample = F.interpolate(
+                laterals[lvl],
+                scale_factor=2.0,
+                mode="nearest",
+            )
+            laterals[lvl - 1] += upsample
+
+        fpn_output = []
+        for lvl in range(num_levels):
+            fpn_output.append(self.fpn_convs[lvl](laterals[lvl]))
+
+        if self.extra_stage > 0:
+            # use max pool to get more levels on top of outputs (Faster R-CNN, Mask R-CNN)
+            if not self.has_extra_convs:
+                assert (
+                    self.extra_stage == 1
+                ), "extra_stage should be 1 if FPN has not extra convs"
+                fpn_output.append(F.max_pool2d(fpn_output[-1], 1, stride=2))
+            # add extra conv levels for RetinaNet(use_c5)/FCOS(use_p5)
+            else:
+                if self.use_c5:
+                    extra_source = body_feats[-1]
+                else:
+                    extra_source = fpn_output[-1]
+                fpn_output.append(self.fpn_convs[num_levels](extra_source))
+
+                for i in range(1, self.extra_stage):
+                    if self.relu_before_extra_convs:
+                        fpn_output.append(
+                            self.fpn_convs[num_levels + i](F.relu(fpn_output[-1]))
+                        )
+                    else:
+                        fpn_output.append(
+                            self.fpn_convs[num_levels + i](fpn_output[-1])
+                        )
+        return fpn_output

ppocr/modeling/necks/fpn.py

@@ -0,0 +1,149 @@
+# copyright (c) 2021 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/whai362/PSENet/blob/python3/models/neck/fpn.py
+"""
+
+import paddle.nn as nn
+import paddle
+import math
+import paddle.nn.functional as F
+
+
+class Conv_BN_ReLU(nn.Layer):
+    def __init__(self, in_planes, out_planes, kernel_size=1, stride=1, padding=0):
+        super(Conv_BN_ReLU, self).__init__()
+        self.conv = nn.Conv2D(
+            in_planes,
+            out_planes,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            bias_attr=False,
+        )
+        self.bn = nn.BatchNorm2D(out_planes, momentum=0.1)
+        self.relu = nn.ReLU()
+
+        for m in self.sublayers():
+            if isinstance(m, nn.Conv2D):
+                n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels
+                m.weight = paddle.create_parameter(
+                    shape=m.weight.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Normal(
+                        0, math.sqrt(2.0 / n)
+                    ),
+                )
+            elif isinstance(m, nn.BatchNorm2D):
+                m.weight = paddle.create_parameter(
+                    shape=m.weight.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Constant(1.0),
+                )
+                m.bias = paddle.create_parameter(
+                    shape=m.bias.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Constant(0.0),
+                )
+
+    def forward(self, x):
+        return self.relu(self.bn(self.conv(x)))
+
+
+class FPN(nn.Layer):
+    def __init__(self, in_channels, out_channels):
+        super(FPN, self).__init__()
+
+        # Top layer
+        self.toplayer_ = Conv_BN_ReLU(
+            in_channels[3], out_channels, kernel_size=1, stride=1, padding=0
+        )
+        # Lateral layers
+        self.latlayer1_ = Conv_BN_ReLU(
+            in_channels[2], out_channels, kernel_size=1, stride=1, padding=0
+        )
+
+        self.latlayer2_ = Conv_BN_ReLU(
+            in_channels[1], out_channels, kernel_size=1, stride=1, padding=0
+        )
+
+        self.latlayer3_ = Conv_BN_ReLU(
+            in_channels[0], out_channels, kernel_size=1, stride=1, padding=0
+        )
+
+        # Smooth layers
+        self.smooth1_ = Conv_BN_ReLU(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        self.smooth2_ = Conv_BN_ReLU(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        self.smooth3_ = Conv_BN_ReLU(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        self.out_channels = out_channels * 4
+        for m in self.sublayers():
+            if isinstance(m, nn.Conv2D):
+                n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels
+                m.weight = paddle.create_parameter(
+                    shape=m.weight.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Normal(
+                        0, math.sqrt(2.0 / n)
+                    ),
+                )
+            elif isinstance(m, nn.BatchNorm2D):
+                m.weight = paddle.create_parameter(
+                    shape=m.weight.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Constant(1.0),
+                )
+                m.bias = paddle.create_parameter(
+                    shape=m.bias.shape,
+                    dtype="float32",
+                    default_initializer=paddle.nn.initializer.Constant(0.0),
+                )
+
+    def _upsample(self, x, scale=1):
+        return F.upsample(x, scale_factor=scale, mode="bilinear")
+
+    def _upsample_add(self, x, y, scale=1):
+        return F.upsample(x, scale_factor=scale, mode="bilinear") + y
+
+    def forward(self, x):
+        f2, f3, f4, f5 = x
+        p5 = self.toplayer_(f5)
+
+        f4 = self.latlayer1_(f4)
+        p4 = self._upsample_add(p5, f4, 2)
+        p4 = self.smooth1_(p4)
+
+        f3 = self.latlayer2_(f3)
+        p3 = self._upsample_add(p4, f3, 2)
+        p3 = self.smooth2_(p3)
+
+        f2 = self.latlayer3_(f2)
+        p2 = self._upsample_add(p3, f2, 2)
+        p2 = self.smooth3_(p2)
+
+        p3 = self._upsample(p3, 2)
+        p4 = self._upsample(p4, 4)
+        p5 = self._upsample(p5, 8)
+
+        fuse = paddle.concat([p2, p3, p4, p5], axis=1)
+        return fuse

ppocr/modeling/necks/fpn_unet.py

@@ -0,0 +1,103 @@
+# 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/open-mmlab/mmocr/blob/main/mmocr/models/textdet/necks/fpn_unet.py
+"""
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class UpBlock(nn.Layer):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        assert isinstance(in_channels, int)
+        assert isinstance(out_channels, int)
+
+        self.conv1x1 = nn.Conv2D(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.conv3x3 = nn.Conv2D(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.deconv = nn.Conv2DTranspose(
+            out_channels, out_channels, kernel_size=4, stride=2, padding=1
+        )
+
+    def forward(self, x):
+        x = F.relu(self.conv1x1(x))
+        x = F.relu(self.conv3x3(x))
+        x = self.deconv(x)
+        return x
+
+
+class FPN_UNet(nn.Layer):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        assert len(in_channels) == 4
+        assert isinstance(out_channels, int)
+        self.out_channels = out_channels
+
+        blocks_out_channels = [out_channels] + [
+            min(out_channels * 2**i, 256) for i in range(4)
+        ]
+        blocks_in_channels = (
+            [blocks_out_channels[1]]
+            + [in_channels[i] + blocks_out_channels[i + 2] for i in range(3)]
+            + [in_channels[3]]
+        )
+
+        self.up4 = nn.Conv2DTranspose(
+            blocks_in_channels[4],
+            blocks_out_channels[4],
+            kernel_size=4,
+            stride=2,
+            padding=1,
+        )
+        self.up_block3 = UpBlock(blocks_in_channels[3], blocks_out_channels[3])
+        self.up_block2 = UpBlock(blocks_in_channels[2], blocks_out_channels[2])
+        self.up_block1 = UpBlock(blocks_in_channels[1], blocks_out_channels[1])
+        self.up_block0 = UpBlock(blocks_in_channels[0], blocks_out_channels[0])
+
+    def forward(self, x):
+        """
+        Args:
+            x (list[Tensor] | tuple[Tensor]): A list of four tensors of shape
+                :math:`(N, C_i, H_i, W_i)`, representing C2, C3, C4, C5
+                features respectively. :math:`C_i` should matches the number in
+                ``in_channels``.
+
+        Returns:
+            Tensor: Shape :math:`(N, C, H, W)` where :math:`H=4H_0` and
+            :math:`W=4W_0`.
+        """
+        c2, c3, c4, c5 = x
+
+        x = F.relu(self.up4(c5))
+
+        x = paddle.concat([x, c4], axis=1)
+        x = F.relu(self.up_block3(x))
+
+        x = paddle.concat([x, c3], axis=1)
+        x = F.relu(self.up_block2(x))
+
+        x = paddle.concat([x, c2], axis=1)
+        x = F.relu(self.up_block1(x))
+
+        x = self.up_block0(x)
+        return x

ppocr/modeling/necks/intracl.py

@@ -0,0 +1,121 @@
+import paddle
+from paddle import nn
+
+# refer from: https://github.com/ViTAE-Transformer/I3CL/blob/736c80237f66d352d488e83b05f3e33c55201317/mmdet/models/detectors/intra_cl_module.py
+
+
+class IntraCLBlock(nn.Layer):
+    def __init__(self, in_channels=96, reduce_factor=4):
+        super(IntraCLBlock, self).__init__()
+        self.channels = in_channels
+        self.rf = reduce_factor
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+        self.conv1x1_reduce_channel = nn.Conv2D(
+            self.channels, self.channels // self.rf, kernel_size=1, stride=1, padding=0
+        )
+        self.conv1x1_return_channel = nn.Conv2D(
+            self.channels // self.rf, self.channels, kernel_size=1, stride=1, padding=0
+        )
+
+        self.v_layer_7x1 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(7, 1),
+            stride=(1, 1),
+            padding=(3, 0),
+        )
+        self.v_layer_5x1 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(5, 1),
+            stride=(1, 1),
+            padding=(2, 0),
+        )
+        self.v_layer_3x1 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(3, 1),
+            stride=(1, 1),
+            padding=(1, 0),
+        )
+
+        self.q_layer_1x7 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 7),
+            stride=(1, 1),
+            padding=(0, 3),
+        )
+        self.q_layer_1x5 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 5),
+            stride=(1, 1),
+            padding=(0, 2),
+        )
+        self.q_layer_1x3 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(1, 3),
+            stride=(1, 1),
+            padding=(0, 1),
+        )
+
+        # base
+        self.c_layer_7x7 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(7, 7),
+            stride=(1, 1),
+            padding=(3, 3),
+        )
+        self.c_layer_5x5 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(5, 5),
+            stride=(1, 1),
+            padding=(2, 2),
+        )
+        self.c_layer_3x3 = nn.Conv2D(
+            self.channels // self.rf,
+            self.channels // self.rf,
+            kernel_size=(3, 3),
+            stride=(1, 1),
+            padding=(1, 1),
+        )
+
+        self.bn = nn.BatchNorm2D(self.channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x_new = self.conv1x1_reduce_channel(x)
+
+        x_7_c = self.c_layer_7x7(x_new)
+        x_7_v = self.v_layer_7x1(x_new)
+        x_7_q = self.q_layer_1x7(x_new)
+        x_7 = x_7_c + x_7_v + x_7_q
+
+        x_5_c = self.c_layer_5x5(x_7)
+        x_5_v = self.v_layer_5x1(x_7)
+        x_5_q = self.q_layer_1x5(x_7)
+        x_5 = x_5_c + x_5_v + x_5_q
+
+        x_3_c = self.c_layer_3x3(x_5)
+        x_3_v = self.v_layer_3x1(x_5)
+        x_3_q = self.q_layer_1x3(x_5)
+        x_3 = x_3_c + x_3_v + x_3_q
+
+        x_relation = self.conv1x1_return_channel(x_3)
+
+        x_relation = self.bn(x_relation)
+        x_relation = self.relu(x_relation)
+
+        return x + x_relation
+
+
+def build_intraclblock_list(num_block):
+    IntraCLBlock_list = nn.LayerList()
+    for i in range(num_block):
+        IntraCLBlock_list.append(IntraCLBlock())
+
+    return IntraCLBlock_list

ppocr/modeling/necks/pg_fpn.py

@@ -0,0 +1,345 @@
+# copyright (c) 2021 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        is_vd_mode=False,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True
+        )
+        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,
+        )
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = 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",
+            use_global_stats=False,
+        )
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class DeConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=4,
+        stride=2,
+        padding=1,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(DeConvBNLayer, self).__init__()
+
+        self.if_act = if_act
+        self.act = act
+        self.deconv = nn.Conv2DTranspose(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False,
+        )
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+            use_global_stats=False,
+        )
+
+    def forward(self, x):
+        x = self.deconv(x)
+        x = self.bn(x)
+        return x
+
+
+class PGFPN(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super(PGFPN, self).__init__()
+        num_inputs = [2048, 2048, 1024, 512, 256]
+        num_outputs = [256, 256, 192, 192, 128]
+        self.out_channels = 128
+        self.conv_bn_layer_1 = ConvBNLayer(
+            in_channels=3,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name="FPN_d1",
+        )
+        self.conv_bn_layer_2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name="FPN_d2",
+        )
+        self.conv_bn_layer_3 = ConvBNLayer(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name="FPN_d3",
+        )
+        self.conv_bn_layer_4 = ConvBNLayer(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=3,
+            stride=2,
+            act=None,
+            name="FPN_d4",
+        )
+        self.conv_bn_layer_5 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="FPN_d5",
+        )
+        self.conv_bn_layer_6 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=3,
+            stride=2,
+            act=None,
+            name="FPN_d6",
+        )
+        self.conv_bn_layer_7 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="FPN_d7",
+        )
+        self.conv_bn_layer_8 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="FPN_d8",
+        )
+
+        self.conv_h0 = ConvBNLayer(
+            in_channels=num_inputs[0],
+            out_channels=num_outputs[0],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(0),
+        )
+        self.conv_h1 = ConvBNLayer(
+            in_channels=num_inputs[1],
+            out_channels=num_outputs[1],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(1),
+        )
+        self.conv_h2 = ConvBNLayer(
+            in_channels=num_inputs[2],
+            out_channels=num_outputs[2],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(2),
+        )
+        self.conv_h3 = ConvBNLayer(
+            in_channels=num_inputs[3],
+            out_channels=num_outputs[3],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(3),
+        )
+        self.conv_h4 = ConvBNLayer(
+            in_channels=num_inputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(4),
+        )
+
+        self.dconv0 = DeConvBNLayer(
+            in_channels=num_outputs[0],
+            out_channels=num_outputs[0 + 1],
+            name="dconv_{}".format(0),
+        )
+        self.dconv1 = DeConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[1 + 1],
+            act=None,
+            name="dconv_{}".format(1),
+        )
+        self.dconv2 = DeConvBNLayer(
+            in_channels=num_outputs[2],
+            out_channels=num_outputs[2 + 1],
+            act=None,
+            name="dconv_{}".format(2),
+        )
+        self.dconv3 = DeConvBNLayer(
+            in_channels=num_outputs[3],
+            out_channels=num_outputs[3 + 1],
+            act=None,
+            name="dconv_{}".format(3),
+        )
+        self.conv_g1 = ConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[1],
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv_g{}".format(1),
+        )
+        self.conv_g2 = ConvBNLayer(
+            in_channels=num_outputs[2],
+            out_channels=num_outputs[2],
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv_g{}".format(2),
+        )
+        self.conv_g3 = ConvBNLayer(
+            in_channels=num_outputs[3],
+            out_channels=num_outputs[3],
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv_g{}".format(3),
+        )
+        self.conv_g4 = ConvBNLayer(
+            in_channels=num_outputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=3,
+            stride=1,
+            act="relu",
+            name="conv_g{}".format(4),
+        )
+        self.convf = ConvBNLayer(
+            in_channels=num_outputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_f{}".format(4),
+        )
+
+    def forward(self, x):
+        c0, c1, c2, c3, c4, c5, c6 = x
+        # FPN_Down_Fusion
+        f = [c0, c1, c2]
+        g = [None, None, None]
+        h = [None, None, None]
+        h[0] = self.conv_bn_layer_1(f[0])
+        h[1] = self.conv_bn_layer_2(f[1])
+        h[2] = self.conv_bn_layer_3(f[2])
+
+        g[0] = self.conv_bn_layer_4(h[0])
+        g[1] = paddle.add(g[0], h[1])
+        g[1] = F.relu(g[1])
+        g[1] = self.conv_bn_layer_5(g[1])
+        g[1] = self.conv_bn_layer_6(g[1])
+
+        g[2] = paddle.add(g[1], h[2])
+        g[2] = F.relu(g[2])
+        g[2] = self.conv_bn_layer_7(g[2])
+        f_down = self.conv_bn_layer_8(g[2])
+
+        # FPN UP Fusion
+        f1 = [c6, c5, c4, c3, c2]
+        g = [None, None, None, None, None]
+        h = [None, None, None, None, None]
+        h[0] = self.conv_h0(f1[0])
+        h[1] = self.conv_h1(f1[1])
+        h[2] = self.conv_h2(f1[2])
+        h[3] = self.conv_h3(f1[3])
+        h[4] = self.conv_h4(f1[4])
+
+        g[0] = self.dconv0(h[0])
+        g[1] = paddle.add(g[0], h[1])
+        g[1] = F.relu(g[1])
+        g[1] = self.conv_g1(g[1])
+        g[1] = self.dconv1(g[1])
+
+        g[2] = paddle.add(g[1], h[2])
+        g[2] = F.relu(g[2])
+        g[2] = self.conv_g2(g[2])
+        g[2] = self.dconv2(g[2])
+
+        g[3] = paddle.add(g[2], h[3])
+        g[3] = F.relu(g[3])
+        g[3] = self.conv_g3(g[3])
+        g[3] = self.dconv3(g[3])
+
+        g[4] = paddle.add(x=g[3], y=h[4])
+        g[4] = F.relu(g[4])
+        g[4] = self.conv_g4(g[4])
+        f_up = self.convf(g[4])
+        f_common = paddle.add(f_down, f_up)
+        f_common = F.relu(f_common)
+        return f_common

ppocr/modeling/necks/pren_fpn.py

@@ -0,0 +1,177 @@
+# 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.
+"""
+Code is refer from:
+https://github.com/RuijieJ/pren/blob/main/Nets/Aggregation.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+
+
+class PoolAggregate(nn.Layer):
+    def __init__(self, n_r, d_in, d_middle=None, d_out=None):
+        super(PoolAggregate, self).__init__()
+        if not d_middle:
+            d_middle = d_in
+        if not d_out:
+            d_out = d_in
+
+        self.d_in = d_in
+        self.d_middle = d_middle
+        self.d_out = d_out
+        self.act = nn.Swish()
+
+        self.n_r = n_r
+        self.aggs = self._build_aggs()
+
+    def _build_aggs(self):
+        aggs = []
+        for i in range(self.n_r):
+            aggs.append(
+                self.add_sublayer(
+                    "{}".format(i),
+                    nn.Sequential(
+                        (
+                            "conv1",
+                            nn.Conv2D(
+                                self.d_in, self.d_middle, 3, 2, 1, bias_attr=False
+                            ),
+                        ),
+                        ("bn1", nn.BatchNorm(self.d_middle)),
+                        ("act", self.act),
+                        (
+                            "conv2",
+                            nn.Conv2D(
+                                self.d_middle, self.d_out, 3, 2, 1, bias_attr=False
+                            ),
+                        ),
+                        ("bn2", nn.BatchNorm(self.d_out)),
+                    ),
+                )
+            )
+        return aggs
+
+    def forward(self, x):
+        b = x.shape[0]
+        outs = []
+        for agg in self.aggs:
+            y = agg(x)
+            p = F.adaptive_avg_pool2d(y, 1)
+            outs.append(p.reshape((b, 1, self.d_out)))
+        out = paddle.concat(outs, 1)
+        return out
+
+
+class WeightAggregate(nn.Layer):
+    def __init__(self, n_r, d_in, d_middle=None, d_out=None):
+        super(WeightAggregate, self).__init__()
+        if not d_middle:
+            d_middle = d_in
+        if not d_out:
+            d_out = d_in
+
+        self.n_r = n_r
+        self.d_out = d_out
+        self.act = nn.Swish()
+
+        self.conv_n = nn.Sequential(
+            ("conv1", nn.Conv2D(d_in, d_in, 3, 1, 1, bias_attr=False)),
+            ("bn1", nn.BatchNorm(d_in)),
+            ("act1", self.act),
+            ("conv2", nn.Conv2D(d_in, n_r, 1, bias_attr=False)),
+            ("bn2", nn.BatchNorm(n_r)),
+            ("act2", nn.Sigmoid()),
+        )
+        self.conv_d = nn.Sequential(
+            ("conv1", nn.Conv2D(d_in, d_middle, 3, 1, 1, bias_attr=False)),
+            ("bn1", nn.BatchNorm(d_middle)),
+            ("act1", self.act),
+            ("conv2", nn.Conv2D(d_middle, d_out, 1, bias_attr=False)),
+            ("bn2", nn.BatchNorm(d_out)),
+        )
+
+    def forward(self, x):
+        b, _, h, w = x.shape
+
+        hmaps = self.conv_n(x)
+        fmaps = self.conv_d(x)
+        r = paddle.bmm(
+            hmaps.reshape((b, self.n_r, h * w)),
+            fmaps.reshape((b, self.d_out, h * w)).transpose((0, 2, 1)),
+        )
+        return r
+
+
+class GCN(nn.Layer):
+    def __init__(self, d_in, n_in, d_out=None, n_out=None, dropout=0.1):
+        super(GCN, self).__init__()
+        if not d_out:
+            d_out = d_in
+        if not n_out:
+            n_out = d_in
+
+        self.conv_n = nn.Conv1D(n_in, n_out, 1)
+        self.linear = nn.Linear(d_in, d_out)
+        self.dropout = nn.Dropout(dropout)
+        self.act = nn.Swish()
+
+    def forward(self, x):
+        x = self.conv_n(x)
+        x = self.dropout(self.linear(x))
+        return self.act(x)
+
+
+class PRENFPN(nn.Layer):
+    def __init__(self, in_channels, n_r, d_model, max_len, dropout):
+        super(PRENFPN, self).__init__()
+        assert len(in_channels) == 3, "in_channels' length must be 3."
+        c1, c2, c3 = in_channels  # the depths are from big to small
+        # build fpn
+        assert d_model % 3 == 0, "{} can't be divided by 3.".format(d_model)
+        self.agg_p1 = PoolAggregate(n_r, c1, d_out=d_model // 3)
+        self.agg_p2 = PoolAggregate(n_r, c2, d_out=d_model // 3)
+        self.agg_p3 = PoolAggregate(n_r, c3, d_out=d_model // 3)
+
+        self.agg_w1 = WeightAggregate(n_r, c1, 4 * c1, d_model // 3)
+        self.agg_w2 = WeightAggregate(n_r, c2, 4 * c2, d_model // 3)
+        self.agg_w3 = WeightAggregate(n_r, c3, 4 * c3, d_model // 3)
+
+        self.gcn_pool = GCN(d_model, n_r, d_model, max_len, dropout)
+        self.gcn_weight = GCN(d_model, n_r, d_model, max_len, dropout)
+
+        self.out_channels = d_model
+
+    def forward(self, inputs):
+        f3, f5, f7 = inputs
+
+        rp1 = self.agg_p1(f3)
+        rp2 = self.agg_p2(f5)
+        rp3 = self.agg_p3(f7)
+        rp = paddle.concat([rp1, rp2, rp3], 2)  # [b,nr,d]
+
+        rw1 = self.agg_w1(f3)
+        rw2 = self.agg_w2(f5)
+        rw3 = self.agg_w3(f7)
+        rw = paddle.concat([rw1, rw2, rw3], 2)  # [b,nr,d]
+
+        y1 = self.gcn_pool(rp)
+        y2 = self.gcn_weight(rw)
+        y = 0.5 * (y1 + y2)
+        return y  # [b,max_len,d]

ppocr/modeling/necks/rf_adaptor.py

@@ -0,0 +1,146 @@
+# 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/hikopensource/DAVAR-Lab-OCR/blob/main/davarocr/davar_rcg/models/connects/single_block/RFAdaptor.py
+"""
+
+import paddle
+import paddle.nn as nn
+from paddle.nn.initializer import TruncatedNormal, Constant, Normal, KaimingNormal
+
+kaiming_init_ = KaimingNormal()
+zeros_ = Constant(value=0.0)
+ones_ = Constant(value=1.0)
+
+
+class S2VAdaptor(nn.Layer):
+    """Semantic to Visual adaptation module"""
+
+    def __init__(self, in_channels=512):
+        super(S2VAdaptor, self).__init__()
+
+        self.in_channels = in_channels  # 512
+
+        # feature strengthen module, channel attention
+        self.channel_inter = nn.Linear(
+            self.in_channels, self.in_channels, bias_attr=False
+        )
+        self.channel_bn = nn.BatchNorm1D(self.in_channels)
+        self.channel_act = nn.ReLU()
+        self.apply(self.init_weights)
+
+    def init_weights(self, m):
+        if isinstance(m, nn.Conv2D):
+            kaiming_init_(m.weight)
+            if isinstance(m, nn.Conv2D) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, (nn.BatchNorm, nn.BatchNorm2D, nn.BatchNorm1D)):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward(self, semantic):
+        semantic_source = semantic  # batch, channel, height, width
+
+        # feature transformation
+        semantic = semantic.squeeze(2).transpose([0, 2, 1])  # batch, width, channel
+        channel_att = self.channel_inter(semantic)  # batch, width, channel
+        channel_att = channel_att.transpose([0, 2, 1])  # batch, channel, width
+        channel_bn = self.channel_bn(channel_att)  # batch, channel, width
+        channel_att = self.channel_act(channel_bn)  # batch, channel, width
+
+        # Feature enhancement
+        channel_output = semantic_source * channel_att.unsqueeze(
+            -2
+        )  # batch, channel, 1, width
+
+        return channel_output
+
+
+class V2SAdaptor(nn.Layer):
+    """Visual to Semantic adaptation module"""
+
+    def __init__(self, in_channels=512, return_mask=False):
+        super(V2SAdaptor, self).__init__()
+
+        # parameter initialization
+        self.in_channels = in_channels
+        self.return_mask = return_mask
+
+        # output transformation
+        self.channel_inter = nn.Linear(
+            self.in_channels, self.in_channels, bias_attr=False
+        )
+        self.channel_bn = nn.BatchNorm1D(self.in_channels)
+        self.channel_act = nn.ReLU()
+
+    def forward(self, visual):
+        # Feature enhancement
+        visual = visual.squeeze(2).transpose([0, 2, 1])  # batch, width, channel
+        channel_att = self.channel_inter(visual)  # batch, width, channel
+        channel_att = channel_att.transpose([0, 2, 1])  # batch, channel, width
+        channel_bn = self.channel_bn(channel_att)  # batch, channel, width
+        channel_att = self.channel_act(channel_bn)  # batch, channel, width
+
+        # size alignment
+        channel_output = channel_att.unsqueeze(-2)  # batch, width, channel
+
+        if self.return_mask:
+            return channel_output, channel_att
+        return channel_output
+
+
+class RFAdaptor(nn.Layer):
+    def __init__(self, in_channels=512, use_v2s=True, use_s2v=True, **kwargs):
+        super(RFAdaptor, self).__init__()
+        if use_v2s is True:
+            self.neck_v2s = V2SAdaptor(in_channels=in_channels, **kwargs)
+        else:
+            self.neck_v2s = None
+        if use_s2v is True:
+            self.neck_s2v = S2VAdaptor(in_channels=in_channels, **kwargs)
+        else:
+            self.neck_s2v = None
+        self.out_channels = in_channels
+
+    def forward(self, x):
+        visual_feature, rcg_feature = x
+        if visual_feature is not None:
+            (
+                batch,
+                source_channels,
+                v_source_height,
+                v_source_width,
+            ) = visual_feature.shape
+            visual_feature = visual_feature.reshape(
+                [batch, source_channels, 1, v_source_height * v_source_width]
+            )
+
+        if self.neck_v2s is not None:
+            v_rcg_feature = rcg_feature * self.neck_v2s(visual_feature)
+        else:
+            v_rcg_feature = rcg_feature
+
+        if self.neck_s2v is not None:
+            v_visual_feature = visual_feature + self.neck_s2v(rcg_feature)
+        else:
+            v_visual_feature = visual_feature
+        if v_rcg_feature is not None:
+            batch, source_channels, source_height, source_width = v_rcg_feature.shape
+            v_rcg_feature = v_rcg_feature.reshape(
+                [batch, source_channels, 1, source_height * source_width]
+            )
+
+            v_rcg_feature = v_rcg_feature.squeeze(2).transpose([0, 2, 1])
+        return v_visual_feature, v_rcg_feature

ppocr/modeling/necks/rnn.py

@@ -0,0 +1,284 @@
+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+from ppocr.modeling.heads.rec_ctc_head import get_para_bias_attr
+from ppocr.modeling.backbones.rec_svtrnet import (
+    Block,
+    ConvBNLayer,
+    trunc_normal_,
+    zeros_,
+    ones_,
+)
+
+
+class Im2Seq(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super().__init__()
+        self.out_channels = in_channels
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == 1
+        x = x.squeeze(axis=2)
+        x = x.transpose([0, 2, 1])  # (NTC)(batch, width, channels)
+        return x
+
+
+class EncoderWithRNN(nn.Layer):
+    def __init__(self, in_channels, hidden_size):
+        super(EncoderWithRNN, self).__init__()
+        self.out_channels = hidden_size * 2
+        self.lstm = nn.LSTM(
+            in_channels, hidden_size, direction="bidirectional", num_layers=2
+        )
+
+    def forward(self, x):
+        x, _ = self.lstm(x)
+        return x
+
+
+class BidirectionalLSTM(nn.Layer):
+    def __init__(
+        self,
+        input_size,
+        hidden_size,
+        output_size=None,
+        num_layers=1,
+        dropout=0,
+        direction=False,
+        time_major=False,
+        with_linear=False,
+    ):
+        super(BidirectionalLSTM, self).__init__()
+        self.with_linear = with_linear
+        self.rnn = nn.LSTM(
+            input_size,
+            hidden_size,
+            num_layers=num_layers,
+            dropout=dropout,
+            direction=direction,
+            time_major=time_major,
+        )
+
+        # text recognition the specified structure LSTM with linear
+        if self.with_linear:
+            self.linear = nn.Linear(hidden_size * 2, output_size)
+
+    def forward(self, input_feature):
+        recurrent, _ = self.rnn(
+            input_feature
+        )  # batch_size x T x input_size -> batch_size x T x (2*hidden_size)
+        if self.with_linear:
+            output = self.linear(recurrent)  # batch_size x T x output_size
+            return output
+        return recurrent
+
+
+class EncoderWithCascadeRNN(nn.Layer):
+    def __init__(
+        self, in_channels, hidden_size, out_channels, num_layers=2, with_linear=False
+    ):
+        super(EncoderWithCascadeRNN, self).__init__()
+        self.out_channels = out_channels[-1]
+        self.encoder = nn.LayerList(
+            [
+                BidirectionalLSTM(
+                    in_channels if i == 0 else out_channels[i - 1],
+                    hidden_size,
+                    output_size=out_channels[i],
+                    num_layers=1,
+                    direction="bidirectional",
+                    with_linear=with_linear,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def forward(self, x):
+        for i, l in enumerate(self.encoder):
+            x = l(x)
+        return x
+
+
+class EncoderWithFC(nn.Layer):
+    def __init__(self, in_channels, hidden_size):
+        super(EncoderWithFC, self).__init__()
+        self.out_channels = hidden_size
+        weight_attr, bias_attr = get_para_bias_attr(l2_decay=0.00001, k=in_channels)
+        self.fc = nn.Linear(
+            in_channels,
+            hidden_size,
+            weight_attr=weight_attr,
+            bias_attr=bias_attr,
+            name="reduce_encoder_fea",
+        )
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x
+
+
+class EncoderWithSVTR(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        dims=64,  # XS
+        depth=2,
+        hidden_dims=120,
+        use_guide=False,
+        num_heads=8,
+        qkv_bias=True,
+        mlp_ratio=2.0,
+        drop_rate=0.1,
+        attn_drop_rate=0.1,
+        drop_path=0.0,
+        kernel_size=[3, 3],
+        qk_scale=None,
+    ):
+        super(EncoderWithSVTR, self).__init__()
+        self.depth = depth
+        self.use_guide = use_guide
+        self.conv1 = ConvBNLayer(
+            in_channels,
+            in_channels // 8,
+            kernel_size=kernel_size,
+            padding=[kernel_size[0] // 2, kernel_size[1] // 2],
+            act=nn.Swish,
+        )
+        self.conv2 = ConvBNLayer(
+            in_channels // 8, hidden_dims, kernel_size=1, act=nn.Swish
+        )
+
+        self.svtr_block = nn.LayerList(
+            [
+                Block(
+                    dim=hidden_dims,
+                    num_heads=num_heads,
+                    mixer="Global",
+                    HW=None,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    act_layer=nn.Swish,
+                    attn_drop=attn_drop_rate,
+                    drop_path=drop_path,
+                    norm_layer="nn.LayerNorm",
+                    epsilon=1e-05,
+                    prenorm=False,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(hidden_dims, epsilon=1e-6)
+        self.conv3 = ConvBNLayer(hidden_dims, in_channels, kernel_size=1, act=nn.Swish)
+        # last conv-nxn, the input is concat of input tensor and conv3 output tensor
+        self.conv4 = ConvBNLayer(
+            2 * in_channels,
+            in_channels // 8,
+            kernel_size=kernel_size,
+            padding=[kernel_size[0] // 2, kernel_size[1] // 2],
+            act=nn.Swish,
+        )
+
+        self.conv1x1 = ConvBNLayer(in_channels // 8, dims, kernel_size=1, act=nn.Swish)
+        self.out_channels = dims
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                zeros_(m.bias)
+        elif isinstance(m, nn.LayerNorm):
+            zeros_(m.bias)
+            ones_(m.weight)
+
+    def forward(self, x):
+        # for use guide
+        if self.use_guide:
+            z = x.clone()
+            z.stop_gradient = True
+        else:
+            z = x
+        # for short cut
+        h = z
+        # reduce dim
+        z = self.conv1(z)
+        z = self.conv2(z)
+        # SVTR global block
+        B, C, H, W = z.shape
+        z = z.flatten(2).transpose([0, 2, 1])
+        for blk in self.svtr_block:
+            z = blk(z)
+        z = self.norm(z)
+        # last stage
+        z = z.reshape([0, H, W, C]).transpose([0, 3, 1, 2])
+        z = self.conv3(z)
+        z = paddle.concat((h, z), axis=1)
+        z = self.conv1x1(self.conv4(z))
+        return z
+
+
+class SequenceEncoder(nn.Layer):
+    def __init__(self, in_channels, encoder_type, hidden_size=48, **kwargs):
+        super(SequenceEncoder, self).__init__()
+        self.encoder_reshape = Im2Seq(in_channels)
+        self.out_channels = self.encoder_reshape.out_channels
+        self.encoder_type = encoder_type
+        if encoder_type == "reshape":
+            self.only_reshape = True
+        else:
+            support_encoder_dict = {
+                "reshape": Im2Seq,
+                "fc": EncoderWithFC,
+                "rnn": EncoderWithRNN,
+                "svtr": EncoderWithSVTR,
+                "cascadernn": EncoderWithCascadeRNN,
+            }
+            assert encoder_type in support_encoder_dict, "{} must in {}".format(
+                encoder_type, support_encoder_dict.keys()
+            )
+            if encoder_type == "svtr":
+                self.encoder = support_encoder_dict[encoder_type](
+                    self.encoder_reshape.out_channels, **kwargs
+                )
+            elif encoder_type == "cascadernn":
+                self.encoder = support_encoder_dict[encoder_type](
+                    self.encoder_reshape.out_channels, hidden_size, **kwargs
+                )
+            else:
+                self.encoder = support_encoder_dict[encoder_type](
+                    self.encoder_reshape.out_channels, hidden_size
+                )
+            self.out_channels = self.encoder.out_channels
+            self.only_reshape = False
+
+    def forward(self, x):
+        if self.encoder_type != "svtr":
+            x = self.encoder_reshape(x)
+            if not self.only_reshape:
+                x = self.encoder(x)
+            return x
+        else:
+            x = self.encoder(x)
+            x = self.encoder_reshape(x)
+            return x

ppocr/modeling/necks/sast_fpn.py

@@ -0,0 +1,368 @@
+# copyright (c) 2019 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        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,
+        )
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="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 DeConvBNLayer(nn.Layer):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        groups=1,
+        if_act=True,
+        act=None,
+        name=None,
+    ):
+        super(DeConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.deconv = nn.Conv2DTranspose(
+            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,
+        )
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+        )
+
+    def forward(self, x):
+        x = self.deconv(x)
+        x = self.bn(x)
+        return x
+
+
+class FPN_Up_Fusion(nn.Layer):
+    def __init__(self, in_channels):
+        super(FPN_Up_Fusion, self).__init__()
+        in_channels = in_channels[::-1]
+        out_channels = [256, 256, 192, 192, 128]
+
+        self.h0_conv = ConvBNLayer(
+            in_channels[0], out_channels[0], 1, 1, act=None, name="fpn_up_h0"
+        )
+        self.h1_conv = ConvBNLayer(
+            in_channels[1], out_channels[1], 1, 1, act=None, name="fpn_up_h1"
+        )
+        self.h2_conv = ConvBNLayer(
+            in_channels[2], out_channels[2], 1, 1, act=None, name="fpn_up_h2"
+        )
+        self.h3_conv = ConvBNLayer(
+            in_channels[3], out_channels[3], 1, 1, act=None, name="fpn_up_h3"
+        )
+        self.h4_conv = ConvBNLayer(
+            in_channels[4], out_channels[4], 1, 1, act=None, name="fpn_up_h4"
+        )
+
+        self.g0_conv = DeConvBNLayer(
+            out_channels[0], out_channels[1], 4, 2, act=None, name="fpn_up_g0"
+        )
+
+        self.g1_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[1], out_channels[1], 3, 1, act="relu", name="fpn_up_g1_1"
+            ),
+            DeConvBNLayer(
+                out_channels[1], out_channels[2], 4, 2, act=None, name="fpn_up_g1_2"
+            ),
+        )
+        self.g2_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[2], out_channels[2], 3, 1, act="relu", name="fpn_up_g2_1"
+            ),
+            DeConvBNLayer(
+                out_channels[2], out_channels[3], 4, 2, act=None, name="fpn_up_g2_2"
+            ),
+        )
+        self.g3_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[3], out_channels[3], 3, 1, act="relu", name="fpn_up_g3_1"
+            ),
+            DeConvBNLayer(
+                out_channels[3], out_channels[4], 4, 2, act=None, name="fpn_up_g3_2"
+            ),
+        )
+
+        self.g4_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[4],
+                out_channels[4],
+                3,
+                1,
+                act="relu",
+                name="fpn_up_fusion_1",
+            ),
+            ConvBNLayer(
+                out_channels[4], out_channels[4], 1, 1, act=None, name="fpn_up_fusion_2"
+            ),
+        )
+
+    def _add_relu(self, x1, x2):
+        x = paddle.add(x=x1, y=x2)
+        x = F.relu(x)
+        return x
+
+    def forward(self, x):
+        f = x[2:][::-1]
+        h0 = self.h0_conv(f[0])
+        h1 = self.h1_conv(f[1])
+        h2 = self.h2_conv(f[2])
+        h3 = self.h3_conv(f[3])
+        h4 = self.h4_conv(f[4])
+
+        g0 = self.g0_conv(h0)
+        g1 = self._add_relu(g0, h1)
+        g1 = self.g1_conv(g1)
+        g2 = self.g2_conv(self._add_relu(g1, h2))
+        g3 = self.g3_conv(self._add_relu(g2, h3))
+        g4 = self.g4_conv(self._add_relu(g3, h4))
+
+        return g4
+
+
+class FPN_Down_Fusion(nn.Layer):
+    def __init__(self, in_channels):
+        super(FPN_Down_Fusion, self).__init__()
+        out_channels = [32, 64, 128]
+
+        self.h0_conv = ConvBNLayer(
+            in_channels[0], out_channels[0], 3, 1, act=None, name="fpn_down_h0"
+        )
+        self.h1_conv = ConvBNLayer(
+            in_channels[1], out_channels[1], 3, 1, act=None, name="fpn_down_h1"
+        )
+        self.h2_conv = ConvBNLayer(
+            in_channels[2], out_channels[2], 3, 1, act=None, name="fpn_down_h2"
+        )
+
+        self.g0_conv = ConvBNLayer(
+            out_channels[0], out_channels[1], 3, 2, act=None, name="fpn_down_g0"
+        )
+
+        self.g1_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[1], out_channels[1], 3, 1, act="relu", name="fpn_down_g1_1"
+            ),
+            ConvBNLayer(
+                out_channels[1], out_channels[2], 3, 2, act=None, name="fpn_down_g1_2"
+            ),
+        )
+
+        self.g2_conv = nn.Sequential(
+            ConvBNLayer(
+                out_channels[2],
+                out_channels[2],
+                3,
+                1,
+                act="relu",
+                name="fpn_down_fusion_1",
+            ),
+            ConvBNLayer(
+                out_channels[2],
+                out_channels[2],
+                1,
+                1,
+                act=None,
+                name="fpn_down_fusion_2",
+            ),
+        )
+
+    def forward(self, x):
+        f = x[:3]
+        h0 = self.h0_conv(f[0])
+        h1 = self.h1_conv(f[1])
+        h2 = self.h2_conv(f[2])
+        g0 = self.g0_conv(h0)
+        g1 = paddle.add(x=g0, y=h1)
+        g1 = F.relu(g1)
+        g1 = self.g1_conv(g1)
+        g2 = paddle.add(x=g1, y=h2)
+        g2 = F.relu(g2)
+        g2 = self.g2_conv(g2)
+        return g2
+
+
+class Cross_Attention(nn.Layer):
+    def __init__(self, in_channels):
+        super(Cross_Attention, self).__init__()
+        self.theta_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act="relu", name="f_theta"
+        )
+        self.phi_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act="relu", name="f_phi"
+        )
+        self.g_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act="relu", name="f_g"
+        )
+
+        self.fh_weight_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act=None, name="fh_weight"
+        )
+        self.fh_sc_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act=None, name="fh_sc"
+        )
+
+        self.fv_weight_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act=None, name="fv_weight"
+        )
+        self.fv_sc_conv = ConvBNLayer(
+            in_channels, in_channels, 1, 1, act=None, name="fv_sc"
+        )
+
+        self.f_attn_conv = ConvBNLayer(
+            in_channels * 2, in_channels, 1, 1, act="relu", name="f_attn"
+        )
+
+    def _cal_fweight(self, f, shape):
+        f_theta, f_phi, f_g = f
+        # flatten
+        f_theta = paddle.transpose(f_theta, [0, 2, 3, 1])
+        f_theta = paddle.reshape(f_theta, [shape[0] * shape[1], shape[2], 128])
+        f_phi = paddle.transpose(f_phi, [0, 2, 3, 1])
+        f_phi = paddle.reshape(f_phi, [shape[0] * shape[1], shape[2], 128])
+        f_g = paddle.transpose(f_g, [0, 2, 3, 1])
+        f_g = paddle.reshape(f_g, [shape[0] * shape[1], shape[2], 128])
+        # correlation
+        f_attn = paddle.matmul(f_theta, paddle.transpose(f_phi, [0, 2, 1]))
+        # scale
+        f_attn = f_attn / (128**0.5)
+        f_attn = F.softmax(f_attn)
+        # weighted sum
+        f_weight = paddle.matmul(f_attn, f_g)
+        f_weight = paddle.reshape(f_weight, [shape[0], shape[1], shape[2], 128])
+        return f_weight
+
+    def forward(self, f_common):
+        f_shape = f_common.shape
+        # print('f_shape: ', f_shape)
+
+        f_theta = self.theta_conv(f_common)
+        f_phi = self.phi_conv(f_common)
+        f_g = self.g_conv(f_common)
+
+        ######## horizon ########
+        fh_weight = self._cal_fweight(
+            [f_theta, f_phi, f_g], [f_shape[0], f_shape[2], f_shape[3]]
+        )
+        fh_weight = paddle.transpose(fh_weight, [0, 3, 1, 2])
+        fh_weight = self.fh_weight_conv(fh_weight)
+        # short cut
+        fh_sc = self.fh_sc_conv(f_common)
+        f_h = F.relu(fh_weight + fh_sc)
+
+        ######## vertical ########
+        fv_theta = paddle.transpose(f_theta, [0, 1, 3, 2])
+        fv_phi = paddle.transpose(f_phi, [0, 1, 3, 2])
+        fv_g = paddle.transpose(f_g, [0, 1, 3, 2])
+        fv_weight = self._cal_fweight(
+            [fv_theta, fv_phi, fv_g], [f_shape[0], f_shape[3], f_shape[2]]
+        )
+        fv_weight = paddle.transpose(fv_weight, [0, 3, 2, 1])
+        fv_weight = self.fv_weight_conv(fv_weight)
+        # short cut
+        fv_sc = self.fv_sc_conv(f_common)
+        f_v = F.relu(fv_weight + fv_sc)
+
+        ######## merge ########
+        f_attn = paddle.concat([f_h, f_v], axis=1)
+        f_attn = self.f_attn_conv(f_attn)
+        return f_attn
+
+
+class SASTFPN(nn.Layer):
+    def __init__(self, in_channels, with_cab=False, **kwargs):
+        super(SASTFPN, self).__init__()
+        self.in_channels = in_channels
+        self.with_cab = with_cab
+        self.FPN_Down_Fusion = FPN_Down_Fusion(self.in_channels)
+        self.FPN_Up_Fusion = FPN_Up_Fusion(self.in_channels)
+        self.out_channels = 128
+        self.cross_attention = Cross_Attention(self.out_channels)
+
+    def forward(self, x):
+        # down fpn
+        f_down = self.FPN_Down_Fusion(x)
+
+        # up fpn
+        f_up = self.FPN_Up_Fusion(x)
+
+        # fusion
+        f_common = paddle.add(x=f_down, y=f_up)
+        f_common = F.relu(f_common)
+
+        if self.with_cab:
+            # print('enhence f_common with CAB.')
+            f_common = self.cross_attention(f_common)
+
+        return f_common

ppocr/modeling/necks/table_fpn.py

@@ -0,0 +1,123 @@
+# copyright (c) 2021 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 paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class TableFPN(nn.Layer):
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(TableFPN, self).__init__()
+        self.out_channels = 512
+        weight_attr = paddle.nn.initializer.KaimingUniform()
+        self.in2_conv = nn.Conv2D(
+            in_channels=in_channels[0],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in3_conv = nn.Conv2D(
+            in_channels=in_channels[1],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            stride=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in4_conv = nn.Conv2D(
+            in_channels=in_channels[2],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.in5_conv = nn.Conv2D(
+            in_channels=in_channels[3],
+            out_channels=self.out_channels,
+            kernel_size=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p5_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p4_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p3_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.p2_conv = nn.Conv2D(
+            in_channels=self.out_channels,
+            out_channels=self.out_channels // 4,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+        self.fuse_conv = nn.Conv2D(
+            in_channels=self.out_channels * 4,
+            out_channels=512,
+            kernel_size=3,
+            padding=1,
+            weight_attr=ParamAttr(initializer=weight_attr),
+            bias_attr=False,
+        )
+
+    def forward(self, x):
+        c2, c3, c4, c5 = x
+
+        in5 = self.in5_conv(c5)
+        in4 = self.in4_conv(c4)
+        in3 = self.in3_conv(c3)
+        in2 = self.in2_conv(c2)
+
+        out4 = in4 + F.upsample(
+            in5, size=in4.shape[2:4], mode="nearest", align_mode=1
+        )  # 1/16
+        out3 = in3 + F.upsample(
+            out4, size=in3.shape[2:4], mode="nearest", align_mode=1
+        )  # 1/8
+        out2 = in2 + F.upsample(
+            out3, size=in2.shape[2:4], mode="nearest", align_mode=1
+        )  # 1/4
+
+        p4 = F.upsample(out4, size=in5.shape[2:4], mode="nearest", align_mode=1)
+        p3 = F.upsample(out3, size=in5.shape[2:4], mode="nearest", align_mode=1)
+        p2 = F.upsample(out2, size=in5.shape[2:4], mode="nearest", align_mode=1)
+        fuse = paddle.concat([in5, p4, p3, p2], axis=1)
+        fuse_conv = self.fuse_conv(fuse) * 0.005
+        return [c5 + fuse_conv]

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