first commit

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)