chore: add missing scripts of PyTorch module

python/rapidocr/inference_engine/pytorch/networks/backbones/rec_lcnetv4.py

@@ -0,0 +1,689 @@
+# 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.
+#
+# PyTorch implementation of PPLCNetV4 backbone for PP-OCRv6.
+# Converted from PaddlePaddle: PaddleOCR/ppocr/modeling/backbones/rec_lcnetv4.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Network configuration tables (identical to PaddleOCR)
+# ---------------------------------------------------------------------------
+
+NET_CONFIG_DET = {
+    "tiny": {
+        # stem(mid=16, out=32)  channels: 32 -> 48 -> 64 -> 160
+        "stem": (16, 32),
+        "blocks_s1": [[3, 32, 32, 1, True], [3, 32, 32, 1, False]],
+        "blocks_s2": [
+            [3, 32, 48, 2, False],
+            [3, 48, 48, 1, True],
+            [3, 48, 48, 1, False],
+        ],
+        "blocks_s3": [
+            [3, 48, 64, 2, False],
+            [3, 64, 64, 1, True],
+            [3, 64, 64, 1, False],
+            [3, 64, 64, 1, True],
+            [3, 64, 64, 1, False],
+        ],
+        "blocks_s4": [
+            [3, 64, 160, 2, False],
+            [3, 160, 160, 1, True],
+            [3, 160, 160, 1, False],
+        ],
+    },
+    "small": {
+        # stem(mid=24, out=48)  channels: 48 -> 96 -> 192 -> 384
+        "stem": (24, 48),
+        "blocks_s1": [[3, 48, 48, 1, True], [3, 48, 48, 1, False]],
+        "blocks_s2": [
+            [3, 48, 96, 2, False],
+            [3, 96, 96, 1, True],
+            [3, 96, 96, 1, False],
+        ],
+        "blocks_s3": [
+            [3, 96, 192, 2, False],
+            [3, 192, 192, 1, True],
+            [3, 192, 192, 1, False],
+            [3, 192, 192, 1, True],
+            [3, 192, 192, 1, False],
+        ],
+        "blocks_s4": [
+            [3, 192, 384, 2, False],
+            [3, 384, 384, 1, True],
+            [3, 384, 384, 1, False],
+        ],
+    },
+    "medium": {
+        # stem(mid=64, out=128)  channels: 128 -> 256 -> 512 -> 896
+        "stem": (64, 128),
+        "blocks_s1": [[3, 128, 128, 1, True], [3, 128, 128, 1, False]],
+        "blocks_s2": [
+            [3, 128, 256, 2, False],
+            [3, 256, 256, 1, True],
+            [3, 256, 256, 1, False],
+        ],
+        "blocks_s3": [
+            [3, 256, 512, 2, False],
+            [3, 512, 512, 1, True],
+            [3, 512, 512, 1, False],
+            [3, 512, 512, 1, True],
+            [3, 512, 512, 1, False],
+        ],
+        "blocks_s4": [
+            [3, 512, 896, 2, False],
+            [3, 896, 896, 1, True],
+            [3, 896, 896, 1, False],
+        ],
+    },
+}
+
+
+NET_CONFIG_REC = {
+    "tiny": {
+        # stem: simple (2xConv2D_BN+GELU, mid=24, out=48)  channels: 48 -> 96 -> 160
+        "stem": (24, 48),
+        "stem_type": "simple",
+        "blocks2": [[3, 48, 48, 1, True]],
+        "blocks3": [[3, 48, 48, 1, False]],
+        "blocks4": [
+            [3, 48, 96, (2, 1), False],
+            [3, 96, 96, 1, True],
+            [3, 96, 96, 1, False],
+        ],
+        "blocks5": [
+            [3, 96, 160, (2, 1), False],
+            [3, 160, 160, 1, True],
+            [3, 160, 160, 1, False],
+            [3, 160, 160, 1, False],
+        ],
+        "blocks6": [],
+    },
+    "small": {
+        # stem: branch StemBlock (mid=48, out=96)  channels: 96 -> 192 -> 384
+        "stem": (48, 96),
+        "stem_type": "branch",
+        "blocks2": [[3, 96, 96, 1, True]],
+        "blocks3": [[3, 96, 96, 1, False], [3, 96, 96, 1, False]],
+        "blocks4": [
+            [3, 96, 192, (2, 1), False],
+            [3, 192, 192, 1, True],
+            [3, 192, 192, 1, False],
+            [3, 192, 192, 1, True],
+            [3, 192, 192, 1, False],
+            [3, 192, 192, 1, True],
+            [3, 192, 192, 1, False],
+        ],
+        "blocks5": [
+            [3, 192, 384, (2, 1), False],
+            [3, 384, 384, 1, True],
+            [3, 384, 384, 1, False],
+        ],
+        "blocks6": [],
+    },
+    "medium": {
+        # stem: branch StemBlock (mid=64, out=128)  channels: 128 -> 256 -> 512 -> 768
+        "stem": (64, 128),
+        "stem_type": "branch",
+        "blocks2": [[3, 128, 128, 1, True]],
+        "blocks3": [
+            [3, 128, 256, 1, False],
+            [3, 256, 256, 1, False],
+            [3, 256, 256, 1, True],
+        ],
+        "blocks4": [
+            [3, 256, 512, (2, 1), False],
+            [3, 512, 512, 1, True],
+            [3, 512, 512, 1, False],
+            [3, 512, 512, 1, True],
+            [3, 512, 512, 1, False],
+            [3, 512, 512, 1, True],
+            [3, 512, 512, 1, False],
+        ],
+        "blocks5": [
+            [3, 512, 768, (2, 1), False],
+            [3, 768, 768, 1, True],
+            [3, 768, 768, 1, False],
+        ],
+        "blocks6": [],
+    },
+}
+
+# ---------------------------------------------------------------------------
+# PyTorch building blocks
+# ---------------------------------------------------------------------------
+
+
+class Conv2D_BN(nn.Sequential):
+    """Conv2D + BatchNorm2d (bias=False).  Supports fuse() for deployment."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=1,
+        stride=1,
+        padding=0,
+        groups=1,
+        bn_weight_init=1.0,
+    ):
+        super().__init__()
+        self.add_module(
+            "conv",
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                groups=groups,
+                bias=False,
+            ),
+        )
+        bn = nn.BatchNorm2d(out_channels)
+        nn.init.constant_(bn.weight, 1.0 if bn_weight_init == 1.0 else 0.0)
+        nn.init.constant_(bn.bias, 0.0)
+        self.add_module("bn", bn)
+
+    @torch.no_grad()
+    def fuse(self):
+        """Fuse Conv2d + BatchNorm2d into a single Conv2d (bias=True)."""
+        c, bn = self.conv, self.bn
+        w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+        w = c.weight * w[:, None, None, None]
+        b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
+        # Handle 'same' padding
+        if isinstance(c.padding, str):
+            padding = (c.kernel_size[0] - 1) // 2
+        else:
+            padding = c.padding
+        m = nn.Conv2d(
+            w.shape[1] * c.groups,
+            w.shape[0],
+            w.shape[2:],
+            stride=c.stride,
+            padding=padding,
+            groups=c.groups,
+        )
+        m.weight.data.copy_(w)
+        m.bias.data.copy_(b)
+        return m
+
+
+class ConvBNAct(nn.Module):
+    """Conv2d + BN + ReLU.  Supports rep() that fuses Conv+BN into one Conv."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=3,
+        stride=1,
+        padding=1,
+        groups=1,
+        use_act=True,
+        lr_mult=1.0,
+    ):
+        super().__init__()
+        self.use_act = use_act
+        self.is_repped = False
+        if isinstance(padding, str) and padding.upper() == "SAME":
+            # Use PyTorch's native 'same' padding (handles asymmetric case correctly)
+            padding = "same"
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding=padding,
+            groups=groups,
+            bias=False,
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        if self.use_act:
+            self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if not self.is_repped:
+            x = self.bn(x)
+        if self.use_act:
+            x = self.act(x)
+        return x
+
+    @torch.no_grad()
+    def rep(self):
+        """Fuse Conv2d + BatchNorm2d into a single Conv2d (bias=True)."""
+        if self.is_repped:
+            return
+        c, bn = self.conv, self.bn
+        w = bn.weight / (bn.running_var + bn.eps) ** 0.5
+        fused_w = c.weight * w[:, None, None, None]
+        fused_b = (
+            bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
+        )
+        # Handle 'same' padding
+        if isinstance(c.padding, str):
+            padding = (
+                c.padding[0]
+                if isinstance(c.padding, tuple)
+                else (c.kernel_size[0] - 1) // 2
+            )
+        else:
+            padding = c.padding
+        m = nn.Conv2d(
+            c.in_channels,
+            c.out_channels,
+            c.kernel_size,
+            stride=c.stride,
+            padding=padding,
+            dilation=c.dilation,
+            groups=c.groups,
+        )
+        m.weight.data.copy_(fused_w)
+        m.bias.data.copy_(fused_b)
+        self.conv = m
+        del self.bn
+        self.is_repped = True
+
+
+class StemBlock(nn.Module):
+    """Multi-branch stem with total stride 4 (stem1 stride=2 + stem3 stride=2)."""
+
+    def __init__(self, in_channels=3, mid_channels=48, out_channels=96, lr_mult=1.0):
+        super().__init__()
+        self.is_repped = False
+        self.stem1 = ConvBNAct(
+            in_channels, mid_channels, 3, 2, use_act=True, lr_mult=lr_mult
+        )
+        self.stem2a = ConvBNAct(
+            mid_channels,
+            mid_channels // 2,
+            2,
+            1,
+            padding="SAME",
+            use_act=True,
+            lr_mult=lr_mult,
+        )
+        self.stem2b = ConvBNAct(
+            mid_channels // 2,
+            mid_channels,
+            2,
+            1,
+            padding="SAME",
+            use_act=True,
+            lr_mult=lr_mult,
+        )
+        self.stem3 = ConvBNAct(
+            mid_channels * 2, mid_channels, 3, 2, use_act=True, lr_mult=lr_mult
+        )
+        self.stem4 = ConvBNAct(
+            mid_channels, out_channels, 1, 1, padding=0, use_act=True, lr_mult=lr_mult
+        )
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0)
+        # Note: Paddle uses padding="SAME" (asymmetric), we use explicit F.pad
+
+    def forward(self, x):
+        x = self.stem1(x)
+        x2 = self.stem2b(self.stem2a(x))
+        # Explicit padding to match Paddle's MaxPool2D with padding="SAME"
+        # kernel=2, stride=1, "SAME" -> pad (right=1, bottom=1) only
+        x_pool = F.pad(x, [0, 1, 0, 1])
+        x1 = self.pool(x_pool)
+        x = self.stem4(self.stem3(torch.cat([x1, x2], dim=1)))
+        return x
+
+    def rep(self, fuse_lab=None):
+        if self.is_repped:
+            return
+        for attr in ("stem1", "stem2a", "stem2b", "stem3", "stem4"):
+            getattr(self, attr).rep()
+        self.is_repped = True
+
+
+class SELayer(nn.Module):
+    """Squeeze-and-Excitation layer (reduction=4, Hardsigmoid gating)."""
+
+    def __init__(self, channel, reduction=4, lr_mult=1.0):
+        super().__init__()
+        self.conv1 = nn.Conv2d(channel, channel // reduction, 1)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(channel // reduction, channel, 1)
+        self.hardsigmoid = nn.Hardsigmoid()
+
+    def forward(self, x):
+        identity = x
+        x = x.mean(dim=[2, 3], keepdim=True)
+        x = self.relu(self.conv1(x))
+        x = self.hardsigmoid(self.conv2(x))
+        return x * identity
+
+
+class RepDWConv(nn.Module):
+    """Reparameterizable depthwise convolution.
+
+    Training: 3-branch (3x3 DW + 1x1 DW + identity BN)
+    Inference: fused into a single 3x3 DW Conv
+    """
+
+    def __init__(self, channels, kernel_size=3):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        padding = (kernel_size - 1) // 2
+
+        self.conv = Conv2D_BN(
+            channels, channels, kernel_size, 1, padding, groups=channels
+        )
+        self.conv1 = nn.Conv2d(channels, channels, 1, 1, 0, groups=channels, bias=False)
+        self.bn = nn.BatchNorm2d(channels)
+        nn.init.constant_(self.bn.weight, 1.0)
+        nn.init.constant_(self.bn.bias, 0.0)
+
+        self.is_repped = False
+        self.reparam_conv = None
+
+    def forward(self, x):
+        if self.is_repped:
+            return self.reparam_conv(x)
+        return self.bn(self.conv(x) + self.conv1(x) + x)
+
+    def rep(self, fuse_lab=None):
+        if self.is_repped:
+            return
+        fused = self._fuse_conv()
+        padding = (self.kernel_size - 1) // 2
+        self.reparam_conv = nn.Conv2d(
+            self.channels,
+            self.channels,
+            self.kernel_size,
+            1,
+            padding,
+            groups=self.channels,
+        )
+        self.reparam_conv.weight.data.copy_(fused.weight)
+        self.reparam_conv.bias.data.copy_(fused.bias)
+        del self.conv
+        del self.conv1
+        del self.bn
+        self.is_repped = True
+
+    @torch.no_grad()
+    def _fuse_conv(self):
+        conv = self.conv.fuse()
+        pad_size = self.kernel_size // 2
+        conv1_w = F.pad(self.conv1.weight, [pad_size, pad_size, pad_size, pad_size])
+        identity = F.pad(
+            torch.ones(
+                [self.conv1.weight.shape[0], self.conv1.weight.shape[1], 1, 1],
+                device=self.conv1.weight.device,
+                dtype=self.conv1.weight.dtype,
+            ),
+            [pad_size, pad_size, pad_size, pad_size],
+        )
+        w = conv.weight + conv1_w + identity
+        conv.weight.data.copy_(w)
+        bn = self.bn
+        scale = bn.weight / (bn.running_var + bn.eps) ** 0.5
+        conv.weight.data.copy_(conv.weight * scale[:, None, None, None])
+        conv.bias.data.copy_(bn.bias + (conv.bias - bn.running_mean) * scale)
+        return conv
+
+
+class LCNetV4Block(nn.Module):
+    """LCNetV4 block for detection and recognition.
+
+    Token mixer: RepDWConv when stride=1 and in==out, else plain Conv2D_BN DW conv.
+    Channel mixer: expand -> act -> compress (+ residual when stride=1 and in==out)
+    rep() fuses all Conv2D_BN layers (mathematically exact, no accuracy change).
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        stride,
+        dw_size,
+        use_se=False,
+        lr_mult=1.0,
+        expand_ratio=2,
+        act_type="gelu",
+    ):
+        super().__init__()
+        self.is_repped = False
+        self.has_residual = in_channels == out_channels and stride == 1
+        self.use_rep_dw = stride == 1 and in_channels == out_channels
+
+        self.token_mixer = nn.Sequential()
+        if self.use_rep_dw:
+            self.token_mixer.add_module("rep_dw", RepDWConv(in_channels, dw_size))
+        else:
+            padding = (dw_size - 1) // 2
+            self.token_mixer.add_module(
+                "dw_conv",
+                Conv2D_BN(
+                    in_channels,
+                    in_channels,
+                    dw_size,
+                    stride,
+                    padding,
+                    groups=in_channels,
+                ),
+            )
+        if use_se:
+            self.token_mixer.add_module("se", SELayer(in_channels, lr_mult=lr_mult))
+
+        hidden_channels = int(in_channels * expand_ratio)
+        compress_bn_init = 0.0 if self.has_residual else 1.0
+        self.channel_mixer = nn.Sequential()
+        self.channel_mixer.add_module(
+            "expand", Conv2D_BN(in_channels, hidden_channels, 1, 1, 0)
+        )
+        if act_type == "gelu":
+            self.channel_mixer.add_module("act", nn.GELU())
+        elif act_type == "hswish":
+            self.channel_mixer.add_module("act", nn.Hardswish())
+        elif act_type == "relu":
+            self.channel_mixer.add_module("act", nn.ReLU(inplace=True))
+        self.channel_mixer.add_module(
+            "compress",
+            Conv2D_BN(
+                hidden_channels,
+                out_channels,
+                1,
+                1,
+                0,
+                bn_weight_init=compress_bn_init,
+            ),
+        )
+
+    def forward(self, x):
+        x = self.token_mixer(x)
+        if self.has_residual:
+            return x + self.channel_mixer(x)
+        return self.channel_mixer(x)
+
+    def rep(self, fuse_lab=None):
+        if self.is_repped:
+            return
+        if self.use_rep_dw:
+            self.token_mixer.rep_dw.rep(fuse_lab=fuse_lab)
+        else:
+            self.token_mixer.dw_conv = self.token_mixer.dw_conv.fuse()
+        for name in ("expand", "compress"):
+            m = getattr(self.channel_mixer, name, None)
+            if isinstance(m, Conv2D_BN):
+                setattr(self.channel_mixer, name, m.fuse())
+        self.is_repped = True
+
+
+# ---------------------------------------------------------------------------
+# PPLCNetV4 – unified backbone for det & rec
+# ---------------------------------------------------------------------------
+
+
+class PPLCNetV4(nn.Module):
+    """Unified PPLCNetV4 backbone for text detection and recognition.
+
+    Detection (det=True):
+        model_size in {'tiny', 'small', 'medium'} -- see NET_CONFIG_DET.
+        Returns 4-level feature list [s1_out, s2_out, s3_out, s4_out].
+
+    Recognition (det=False):
+        model_size in {'tiny', 'small', 'medium'} -- see NET_CONFIG_REC.
+        Returns pooled feature tensor [B, C, 1, W].
+    """
+
+    def __init__(
+        self,
+        det=False,
+        model_size="small",
+        in_channels=3,
+        lr_mult_list=None,
+        **kwargs,
+    ):
+        super().__init__()
+        if lr_mult_list is None:
+            lr_mult_list = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+        self.det = det
+        self.is_repped = False
+
+        if det:
+            assert (
+                model_size in NET_CONFIG_DET
+            ), "det model_size must be one of {} but got '{}'".format(
+                list(NET_CONFIG_DET.keys()), model_size
+            )
+            cfg = NET_CONFIG_DET[model_size]
+            stem_mid, stem_out = cfg["stem"]
+            self.stem = StemBlock(in_channels, stem_mid, stem_out)
+
+            def make_stage(key):
+                return nn.Sequential(
+                    *[
+                        LCNetV4Block(in_c, out_c, s, k, se, expand_ratio=2)
+                        for k, in_c, out_c, s, se in cfg[key]
+                    ]
+                )
+
+            self.blocks_s1 = make_stage("blocks_s1")
+            self.blocks_s2 = make_stage("blocks_s2")
+            self.blocks_s3 = make_stage("blocks_s3")
+            self.blocks_s4 = make_stage("blocks_s4")
+            self.out_channels = [
+                cfg["blocks_s1"][-1][2],
+                cfg["blocks_s2"][-1][2],
+                cfg["blocks_s3"][-1][2],
+                cfg["blocks_s4"][-1][2],
+            ]
+
+        else:
+            assert isinstance(lr_mult_list, (list, tuple)) and len(lr_mult_list) == 6
+            assert (
+                model_size in NET_CONFIG_REC
+            ), "rec model_size must be one of {} but got '{}'".format(
+                list(NET_CONFIG_REC.keys()), model_size
+            )
+            self.lr_mult_list = lr_mult_list
+            cfg = NET_CONFIG_REC[model_size]
+            stem_mid, stem_out = cfg["stem"]
+            if cfg["stem_type"] == "branch":
+                self.conv1 = StemBlock(3, stem_mid, stem_out, lr_mult=lr_mult_list[0])
+            else:
+                self.conv1 = nn.Sequential(
+                    Conv2D_BN(3, stem_mid, 3, 2, 1),
+                    nn.GELU(),
+                    Conv2D_BN(stem_mid, stem_out, 3, 2, 1),
+                )
+
+            def make_stage(stage_name, lr_idx):
+                return nn.Sequential(
+                    *[
+                        LCNetV4Block(
+                            in_c,
+                            out_c,
+                            s,
+                            k,
+                            se,
+                            lr_mult=lr_mult_list[lr_idx],
+                            expand_ratio=2,
+                        )
+                        for k, in_c, out_c, s, se in cfg.get(stage_name, [])
+                    ]
+                )
+
+            self.blocks2 = make_stage("blocks2", 1)
+            self.blocks3 = make_stage("blocks3", 2)
+            self.blocks4 = make_stage("blocks4", 3)
+            self.blocks5 = make_stage("blocks5", 4)
+            self.blocks6 = make_stage("blocks6", 5)
+
+            for sname in reversed(
+                ["blocks2", "blocks3", "blocks4", "blocks5", "blocks6"]
+            ):
+                if cfg.get(sname):
+                    self.out_channels = cfg[sname][-1][2]
+                    break
+
+    def forward(self, x):
+        if self.det:
+            x = self.stem(x)
+            o1 = self.blocks_s1(x)
+            o2 = self.blocks_s2(o1)
+            o3 = self.blocks_s3(o2)
+            o4 = self.blocks_s4(o3)
+            return [o1, o2, o3, o4]
+        else:
+            x = self.conv1(x)
+            x = self.blocks2(x)
+            x = self.blocks3(x)
+            x = self.blocks4(x)
+            x = self.blocks5(x)
+            x = self.blocks6(x)
+            if self.training:
+                x = F.adaptive_avg_pool2d(x, [1, 40])
+            else:
+                assert x.shape[2] >= 3, f"Feature height {x.shape[2]} < pool kernel 3."
+                x = F.avg_pool2d(x, [3, 2])
+            return x
+
+    def rep(self, fuse_lab=None):
+        if self.is_repped:
+            return
+        if self.det:
+            self.stem.rep()
+            for stage in [
+                self.blocks_s1,
+                self.blocks_s2,
+                self.blocks_s3,
+                self.blocks_s4,
+            ]:
+                for block in stage:
+                    block.rep(fuse_lab=fuse_lab)
+        else:
+            if hasattr(self.conv1, "rep"):
+                self.conv1.rep()
+            for stage in [
+                self.blocks2,
+                self.blocks3,
+                self.blocks4,
+                self.blocks5,
+                self.blocks6,
+            ]:
+                for block in stage:
+                    block.rep()
+        self.is_repped = True

python/rapidocr/inference_engine/pytorch/networks/backbones/rec_pplcnet.py

@@ -0,0 +1,287 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+# copyright (c) 2024 PytorchOCR 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.
+"""
+PP-LCNet v1 backbone implementation for PyTorch.
+
+PP-LCNet: A Lightweight CPU Convolutional Neural Network
+Paper: https://arxiv.org/abs/2109.15099
+
+This backbone implements the original PP-LCNet (v1) architecture used in
+PaddleOCR's textline_ori and doc_ori models.
+The backbone includes the GAP + 1280-dim 1x1 conv + H-Swish + Dropout
+that is characteristic of PP-LCNet, allowing ClsHead to work directly.
+"""
+
+from __future__ import absolute_import, division, print_function
+
+import torch.nn as nn
+
+
+# PP-LCNet network configuration (original ImageNet style)
+# Format: blocks{stage}: [kernel_size, in_channels, out_channels, stride, use_se]
+NET_CONFIG = {
+    "blocks2": [
+        [3, 16, 32, 1, False],
+    ],
+    "blocks3": [
+        [3, 32, 64, 2, False],
+        [3, 64, 64, 1, False],
+    ],
+    "blocks4": [
+        [3, 64, 128, 2, False],
+        [3, 128, 128, 1, False],
+    ],
+    "blocks5": [
+        [3, 128, 256, 2, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+        [5, 256, 256, 1, False],
+    ],
+    "blocks6": [
+        [5, 256, 512, 2, True],
+        [5, 512, 512, 1, True],
+    ],
+}
+
+
+def make_divisible(v, divisor=8, min_value=None):
+    """Make channel divisible by divisor (same as PaddleClas)."""
+    if min_value is None:
+        min_value = divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    if new_v < 0.9 * v:
+        new_v += divisor
+    return new_v
+
+
+class DepthSepConv(nn.Module):
+    """
+    Depthwise Separable Convolution block (MobileNetV1 style).
+    PP-LCNet uses simple DepthSepConv without residual connections.
+
+    DW Conv(kxk, groups=in_ch) → BN → H-Swish → PW Conv(1x1) → BN → H-Swish
+    """
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
+        super().__init__()
+        self.dw_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=in_channels,
+            bias=False,
+        )
+        self.dw_bn = nn.BatchNorm2d(in_channels)
+        self.dw_act = nn.Hardswish(inplace=True)
+
+        self.pw_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.pw_bn = nn.BatchNorm2d(out_channels)
+        self.pw_act = nn.Hardswish(inplace=True)
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        x = self.dw_bn(x)
+        x = self.dw_act(x)
+        x = self.pw_conv(x)
+        x = self.pw_bn(x)
+        x = self.pw_act(x)
+        return x
+
+
+class SELayer(nn.Module):
+    """Squeeze-and-Excitation module (PP-LCNet uses reduction=4)."""
+
+    def __init__(self, channel, reduction=4):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv1 = nn.Conv2d(
+            in_channels=channel,
+            out_channels=channel // reduction,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(
+            in_channels=channel // reduction,
+            out_channels=channel,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=True,
+        )
+        self.hardsigmoid = nn.Hardsigmoid(inplace=True)
+
+    def forward(self, x):
+        identity = x
+        x = self.avg_pool(x)
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.hardsigmoid(x)
+        return identity * x
+
+
+class PPLCNetBlock(nn.Module):
+    """PP-LCNet basic block: DW Conv → BN → H-Swish → (SE) → PW Conv → BN → H-Swish."""
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, use_se=False):
+        super().__init__()
+        self.use_se = use_se
+        # Depthwise conv
+        self.dw_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=in_channels,
+            bias=False,
+        )
+        self.dw_bn = nn.BatchNorm2d(in_channels)
+        self.dw_act = nn.Hardswish(inplace=True)
+
+        # SE on intermediate features (after DW, before PW)
+        if use_se:
+            self.se = SELayer(in_channels, reduction=4)
+
+        # Pointwise conv
+        self.pw_conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.pw_bn = nn.BatchNorm2d(out_channels)
+        self.pw_act = nn.Hardswish(inplace=True)
+
+    def forward(self, x):
+        x = self.dw_conv(x)
+        x = self.dw_bn(x)
+        x = self.dw_act(x)
+        if self.use_se:
+            x = self.se(x)
+        x = self.pw_conv(x)
+        x = self.pw_bn(x)
+        x = self.pw_act(x)
+        return x
+
+
+class PPLCNet(nn.Module):
+    """
+    PP-LCNet v1 backbone for classification tasks.
+
+    Used in PaddleOCR's textline_ori and doc_ori models.
+    Outputs a 1280-dim feature vector (after GAP + 1x1 Conv + H-Swish + Dropout),
+    suitable for direct use with ClsHead.
+
+    Args:
+        scale (float): Channel scaling factor (0.25, 0.35, 0.5, 0.75, 1.0, etc.)
+        class_num (int): Number of output classes (NOT used here; ClsHead handles classification).
+                         Included for compatibility but ignored.
+    """
+
+    def __init__(self, scale=1.0, **kwargs):
+        super().__init__()
+        self.scale = scale
+        self.net_config = NET_CONFIG
+
+        # Stem: Conv2D(3→16*scale, k=3, s=2)
+        stem_out = make_divisible(16 * scale)
+        self.conv1 = nn.Conv2d(
+            in_channels=3,
+            out_channels=stem_out,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            bias=False,
+        )
+        self.conv1_bn = nn.BatchNorm2d(stem_out)
+        self.conv1_act = nn.Hardswish(inplace=True)
+
+        # Blocks 2-6 (stages)
+        self.blocks2 = self._make_stage(self.net_config["blocks2"])
+        self.blocks3 = self._make_stage(self.net_config["blocks3"])
+        self.blocks4 = self._make_stage(self.net_config["blocks4"])
+        self.blocks5 = self._make_stage(self.net_config["blocks5"])
+        self.blocks6 = self._make_stage(self.net_config["blocks6"])
+
+        # Last stage output channels (after scale)
+        last_out = make_divisible(512 * scale)
+
+        # PP-LCNet characteristic: GAP → 1x1 Conv(→1280) → H-Swish → Dropout
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.last_conv = nn.Conv2d(
+            in_channels=last_out,
+            out_channels=1280,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.last_act = nn.Hardswish(inplace=True)
+        self.dropout = nn.Dropout(p=0.2)
+
+        self.out_channels = 1280
+
+    def _make_stage(self, stage_config):
+        """Build a sequential stage from config list."""
+        layers = []
+        for k, in_c, out_c, s, se in stage_config:
+            layers.append(
+                PPLCNetBlock(
+                    in_channels=make_divisible(in_c * self.scale),
+                    out_channels=make_divisible(out_c * self.scale),
+                    kernel_size=k,
+                    stride=s,
+                    use_se=se,
+                )
+            )
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        # Stem
+        x = self.conv1(x)
+        x = self.conv1_bn(x)
+        x = self.conv1_act(x)
+
+        # Feature extraction stages
+        x = self.blocks2(x)
+        x = self.blocks3(x)
+        x = self.blocks4(x)
+        x = self.blocks5(x)
+        x = self.blocks6(x)
+
+        # Head projection (PP-LCNet characteristic)
+        x = self.avg_pool(x)  # → [B, C, 1, 1]
+        x = self.last_conv(x)  # → [B, 1280, 1, 1]
+        x = self.last_act(x)
+        x = self.dropout(x)
+
+        return x

python/rapidocr/inference_engine/pytorch/networks/heads/multiheadAttention.py

@@ -0,0 +1,163 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Linear
+from torch.nn.init import xavier_uniform_
+
+
+class MultiheadAttention(nn.Module):
+    """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,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_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.scaling = self.head_dim**-0.5
+        self.out_proj = Linear(embed_dim, embed_dim, bias=bias)
+        self._reset_parameters()
+        self.conv1 = torch.nn.Conv2d(
+            in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)
+        )
+        self.conv2 = torch.nn.Conv2d(
+            in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)
+        )
+        self.conv3 = torch.nn.Conv2d(
+            in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)
+        )
+
+    def _reset_parameters(self):
+        xavier_uniform_(self.out_proj.weight)
+
+    def forward(
+        self,
+        query,
+        key,
+        value,
+        key_padding_mask=None,
+        incremental_state=None,
+        attn_mask=None,
+    ):
+        """
+        Inputs of forward function
+            query: [target length, batch size, embed dim]
+            key: [sequence length, batch size, embed dim]
+            value: [sequence length, batch size, embed dim]
+            key_padding_mask: if True, mask padding based on batch size
+            incremental_state: if provided, previous time steps are cashed
+            need_weights: output attn_output_weights
+            static_kv: key and value are static
+
+        Outputs of forward function
+            attn_output: [target length, batch size, embed dim]
+            attn_output_weights: [batch size, target length, sequence length]
+        """
+        q_shape = query.shape
+        src_shape = key.shape
+        q = self._in_proj_q(query)
+        k = self._in_proj_k(key)
+        v = self._in_proj_v(value)
+        q *= self.scaling
+        # q = paddle.transpose(
+        #     paddle.reshape(
+        #         q, [q_shape[0], q_shape[1], self.num_heads, self.head_dim]),
+        #     [1, 2, 0, 3])
+        q = torch.reshape(q, (q_shape[0], q_shape[1], self.num_heads, self.head_dim))
+        q = q.permute(1, 2, 0, 3)
+        # k = paddle.transpose(
+        #     paddle.reshape(
+        #         k, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]),
+        #     [1, 2, 0, 3])
+        k = torch.reshape(k, (src_shape[0], q_shape[1], self.num_heads, self.head_dim))
+        k = k.permute(1, 2, 0, 3)
+        # v = paddle.transpose(
+        #     paddle.reshape(
+        #         v, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]),
+        #     [1, 2, 0, 3])
+        v = torch.reshape(v, (src_shape[0], q_shape[1], self.num_heads, self.head_dim))
+        v = v.permute(1, 2, 0, 3)
+        if key_padding_mask is not None:
+            assert key_padding_mask.shape[0] == q_shape[1]
+            assert key_padding_mask.shape[1] == src_shape[0]
+        attn_output_weights = torch.matmul(q, k.permute(0, 1, 3, 2))
+        if attn_mask is not None:
+            attn_mask = torch.unsqueeze(torch.unsqueeze(attn_mask, 0), 0)
+            attn_output_weights += attn_mask
+        if key_padding_mask is not None:
+            attn_output_weights = torch.reshape(
+                attn_output_weights,
+                [q_shape[1], self.num_heads, q_shape[0], src_shape[0]],
+            )
+            key = torch.unsqueeze(torch.unsqueeze(key_padding_mask, 1), 2)
+            key = key.type(torch.float32)
+            y = torch.full(
+                size=key.shape, fill_value=float("-Inf"), dtype=torch.float32
+            )
+            y = torch.where(key == 0.0, key, y)
+            attn_output_weights += y
+        attn_output_weights = F.softmax(
+            attn_output_weights.type(torch.float32),
+            dim=-1,
+            dtype=torch.float32
+            if attn_output_weights.dtype == torch.float16
+            else attn_output_weights.dtype,
+        )
+        attn_output_weights = F.dropout(
+            attn_output_weights, p=self.dropout, training=self.training
+        )
+
+        attn_output = torch.matmul(attn_output_weights, v)
+        attn_output = torch.reshape(
+            attn_output.permute(2, 0, 1, 3), [q_shape[0], q_shape[1], self.embed_dim]
+        )
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+
+    def _in_proj_q(self, query):
+        query = query.permute(1, 2, 0)
+        query = torch.unsqueeze(query, dim=2)
+        res = self.conv1(query)
+        res = torch.squeeze(res, dim=2)
+        res = res.permute(2, 0, 1)
+        return res
+
+    def _in_proj_k(self, key):
+        key = key.permute(1, 2, 0)
+        key = torch.unsqueeze(key, dim=2)
+        res = self.conv2(key)
+        res = torch.squeeze(res, dim=2)
+        res = res.permute(2, 0, 1)
+        return res
+
+    def _in_proj_v(self, value):
+        value = value.permute(1, 2, 0)  # (1, 2, 0)
+        value = torch.unsqueeze(value, dim=2)
+        res = self.conv3(value)
+        res = torch.squeeze(res, dim=2)
+        res = res.permute(2, 0, 1)
+        return res

python/rapidocr/inference_engine/pytorch/networks/heads/rec_nrtr_head.py

@@ -0,0 +1,838 @@
+import copy
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Dropout, LayerNorm
+from torch.nn import ModuleList as LayerList
+from torch.nn.init import xavier_normal_, xavier_uniform_
+
+from .multiheadAttention import MultiheadAttention
+
+
+class Transformer(nn.Module):
+    """A transformer model. User is able to modify the attributes as needed. The architechture
+    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,
+        custom_encoder=None,
+        custom_decoder=None,
+        in_channels=0,
+        out_channels=0,
+        scale_embedding=True,
+    ):
+        super(Transformer, self).__init__()
+        self.out_channels = 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 custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            if num_encoder_layers > 0:
+                encoder_layer = TransformerEncoderLayer(
+                    d_model,
+                    nhead,
+                    dim_feedforward,
+                    attention_dropout_rate,
+                    residual_dropout_rate,
+                )
+                self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers)
+            else:
+                self.encoder = None
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            decoder_layer = TransformerDecoderLayer(
+                d_model,
+                nhead,
+                dim_feedforward,
+                attention_dropout_rate,
+                residual_dropout_rate,
+            )
+            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
+
+        self._reset_parameters()
+        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=False)
+        w0 = np.random.normal(
+            0.0, d_model**-0.5, (self.out_channels, d_model)
+        ).astype(np.float32)
+        self.tgt_word_prj.weight.data = torch.from_numpy(w0)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            xavier_normal_(m.weight)
+            if m.bias is not None:
+                torch.nn.init.zeros_(m.bias)
+
+    def forward_train(self, src, tgt):
+        tgt = tgt[:, :-1]
+
+        tgt_key_padding_mask = self.generate_padding_mask(tgt)
+        tgt = self.embedding(tgt).permute(1, 0, 2)
+        tgt = self.positional_encoding(tgt)
+        tgt_mask = self.generate_square_subsequent_mask(tgt.shape[0], tgt.device)
+
+        if self.encoder is not None:
+            src = self.positional_encoding(src.permute(1, 0, 2))
+            memory = self.encoder(src)
+        else:
+            memory = src.squeeze(2).permute(2, 0, 1)
+        output = self.decoder(
+            tgt,
+            memory,
+            tgt_mask=tgt_mask,
+            memory_mask=None,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=None,
+        )
+        output = output.permute(1, 0, 2)
+        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:`(S, N, E)`.
+            - tgt: :math:`(T, N, E)`.
+        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.permute(1, 0, 2))
+            memory = self.encoder(src)
+        else:
+            memory = torch.squeeze(src, 2).permute(2, 0, 1)
+        dec_seq = torch.full((bs, 1), 2, dtype=torch.int64)
+        dec_prob = torch.full((bs, 1), 1.0, dtype=torch.float32)
+        for len_dec_seq in range(1, 25):
+            dec_seq_embed = self.embedding(dec_seq).permute(1, 0, 2)
+            dec_seq_embed = self.positional_encoding(dec_seq_embed)
+            tgt_mask = self.generate_square_subsequent_mask(dec_seq_embed.shape[0])
+            output = self.decoder(
+                dec_seq_embed,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=None,
+                tgt_key_padding_mask=None,
+                memory_key_padding_mask=None,
+            )
+            dec_output = output.permute(1, 0, 2)
+            dec_output = dec_output[:, -1, :]
+            tgt_word_prj = self.tgt_word_prj(dec_output)
+            word_prob = F.softmax(tgt_word_prj, dim=1)
+            preds_idx = word_prob.argmax(dim=1)
+            if torch.equal(
+                preds_idx, torch.full(preds_idx.shape, 3, dtype=torch.int64)
+            ):
+                break
+            preds_prob = torch.max(word_prob, dim=1).values
+            dec_seq = torch.cat([dec_seq, torch.reshape(preds_idx, (-1, 1))], dim=1)
+            dec_prob = torch.cat([dec_prob, torch.reshape(preds_prob, (-1, 1))], dim=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 = torch.tensor(active_inst_idx, dtype=torch.int64)
+            active_src_enc = collect_active_part(
+                src_enc.permute(1, 0, 2), active_inst_idx, n_prev_active_inst, n_bm
+            ).permute(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,
+            memory_key_padding_mask,
+        ):
+            """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 = torch.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, memory_key_padding_mask
+            ):
+                dec_seq = self.embedding(dec_seq).permute(1, 0, 2)
+                dec_seq = self.positional_encoding(dec_seq)
+                tgt_mask = self.generate_square_subsequent_mask(dec_seq.shape[0])
+                dec_output = self.decoder(
+                    dec_seq,
+                    enc_output,
+                    tgt_mask=tgt_mask,
+                    tgt_key_padding_mask=None,
+                    memory_key_padding_mask=memory_key_padding_mask,
+                )
+                dec_output = dec_output.permute(1, 0, 2)
+                dec_output = dec_output[:, -1, :]  # Pick the last step: (bh * bm) * d_h
+                word_prob = F.softmax(self.tgt_word_prj(dec_output), dim=1)
+                word_prob = torch.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, None)
+            # 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 torch.no_grad():
+            # -- Encode
+            if self.encoder is not None:
+                src = self.positional_encoding(images.permute(1, 0, 2))
+                src_enc = self.encoder(src)
+            else:
+                src_enc = images.squeeze(2).transpose([0, 2, 1])
+
+            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])
+            src_enc = src_enc.repeat(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, 25):
+                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,
+                    None,
+                )
+                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 [
+            torch.tensor(np.array(result_hyp), dtype=torch.int64),
+            torch.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 = torch.zeros([sz, sz], dtype=torch.float32)
+        mask_inf = torch.triu(
+            torch.full(size=[sz, sz], fill_value=float("-Inf"), dtype=torch.float32),
+            diagonal=1,
+        )
+        mask = mask + mask_inf
+        return mask
+
+    def generate_padding_mask(self, x):
+        # padding_mask = paddle.equal(x, paddle.to_tensor(0, dtype=x.dtype))
+        padding_mask = x == torch.tensor(0, dtype=x.dtype)
+        return padding_mask
+
+    def _reset_parameters(self):
+        """Initiate parameters in the transformer model."""
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                xavier_uniform_(p)
+
+
+class TransformerEncoder(nn.Module):
+    """TransformerEncoder is a stack of N encoder layers
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+    """
+
+    def __init__(self, encoder_layer, num_layers):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+
+    def forward(self, src):
+        """Pass the input through the endocder layers in turn.
+        Args:
+            src: the sequnce to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+        """
+        output = src
+
+        for i in range(self.num_layers):
+            output = self.layers[i](output, src_mask=None, src_key_padding_mask=None)
+
+        return output
+
+
+class TransformerDecoder(nn.Module):
+    """TransformerDecoder is a stack of N decoder layers
+
+    Args:
+        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+        num_layers: the number of sub-decoder-layers in the decoder (required).
+        norm: the layer normalization component (optional).
+
+    """
+
+    def __init__(self, decoder_layer, num_layers):
+        super(TransformerDecoder, self).__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+    ):
+        """Pass the inputs (and mask) through the decoder layer in turn.
+
+        Args:
+            tgt: the sequence to the decoder (required).
+            memory: the sequnce from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+        """
+        output = tgt
+        for i in range(self.num_layers):
+            output = self.layers[i](
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+            )
+
+        return output
+
+
+class TransformerEncoderLayer(nn.Module):
+    """TransformerEncoderLayer is made up of self-attn and feedforward network.
+    This standard encoder layer 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. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        attention_dropout_rate=0.0,
+        residual_dropout_rate=0.1,
+    ):
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model, nhead, dropout=attention_dropout_rate
+        )
+
+        self.conv1 = nn.Conv2d(
+            in_channels=d_model, out_channels=dim_feedforward, kernel_size=(1, 1)
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=dim_feedforward, out_channels=d_model, kernel_size=(1, 1)
+        )
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout1 = Dropout(residual_dropout_rate)
+        self.dropout2 = Dropout(residual_dropout_rate)
+
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        """Pass the input through the endocder layer.
+        Args:
+            src: the sequnce to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+        """
+        src2 = self.self_attn(
+            src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
+        )
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+
+        src = src.permute(1, 2, 0)
+        src = torch.unsqueeze(src, 2)
+        src2 = self.conv2(F.relu(self.conv1(src)))
+        src2 = torch.squeeze(src2, 2)
+        src2 = src2.permute(2, 0, 1)
+        src = torch.squeeze(src, 2)
+        src = src.permute(2, 0, 1)
+
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+
+class TransformerDecoderLayer(nn.Module):
+    """TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
+    This standard decoder layer 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. Users may modify or implement
+    in a different way during application.
+
+    Args:
+        d_model: the number of expected features in the input (required).
+        nhead: the number of heads in the multiheadattention models (required).
+        dim_feedforward: the dimension of the feedforward network model (default=2048).
+        dropout: the dropout value (default=0.1).
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        attention_dropout_rate=0.0,
+        residual_dropout_rate=0.1,
+    ):
+        super(TransformerDecoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model, nhead, dropout=attention_dropout_rate
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model, nhead, dropout=attention_dropout_rate
+        )
+
+        self.conv1 = nn.Conv2d(
+            in_channels=d_model, out_channels=dim_feedforward, kernel_size=(1, 1)
+        )
+        self.conv2 = nn.Conv2d(
+            in_channels=dim_feedforward, out_channels=d_model, kernel_size=(1, 1)
+        )
+
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.norm3 = LayerNorm(d_model)
+        self.dropout1 = Dropout(residual_dropout_rate)
+        self.dropout2 = Dropout(residual_dropout_rate)
+        self.dropout3 = Dropout(residual_dropout_rate)
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask=None,
+        memory_mask=None,
+        tgt_key_padding_mask=None,
+        memory_key_padding_mask=None,
+    ):
+        """Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequnce from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+
+        """
+        tgt2 = self.self_attn(
+            tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(
+            tgt,
+            memory,
+            memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+
+        # default
+        tgt = tgt.permute(1, 2, 0)
+        tgt = torch.unsqueeze(tgt, 2)
+        tgt2 = self.conv2(F.relu(self.conv1(tgt)))
+        tgt2 = torch.squeeze(tgt2, 2)
+        tgt2 = tgt2.permute(2, 0, 1)
+        tgt = torch.squeeze(tgt, 2)
+        tgt = tgt.permute(2, 0, 1)
+
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+
+def _get_clones(module, N):
+    return LayerList([copy.deepcopy(module) for i in range(N)])
+
+
+class PositionalEncoding(nn.Module):
+    """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 = torch.zeros([max_len, dim])
+        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, dim, 2).type(torch.float32) * (-math.log(10000.0) / dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = torch.unsqueeze(pe, 0)
+        pe = pe.permute(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 + self.pe[: x.shape[0], :]
+        return self.dropout(x)
+
+
+class PositionalEncoding_2d(nn.Module):
+    """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 = torch.zeros([max_len, dim])
+        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, dim, 2).type(torch.float32) * (-math.log(10000.0) / dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = torch.unsqueeze(pe, 0).permute(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 = w_pe.permute(1, 2, 0)
+        w_pe = torch.unsqueeze(w_pe, 2)
+
+        h_pe = self.pe[: x.shape[-2], :]
+        w2 = self.linear2(self.avg_pool_2(x).squeeze()).unsqueeze(0)
+        h_pe = h_pe * w2
+        h_pe = h_pe.permute(1, 2, 0)
+        h_pe = torch.unsqueeze(h_pe, 3)
+
+        x = x + w_pe + h_pe
+        x = torch.reshape(x, [x.shape[0], x.shape[1], x.shape[2] * x.shape[3]]).permute(
+            2, 0, 1
+        )
+
+        return self.dropout(x)
+
+
+class Embeddings(nn.Module):
+    def __init__(self, d_model, vocab, padding_idx, scale_embedding):
+        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.data = torch.from_numpy(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 = torch.zeros((size,), dtype=torch.float32)
+        self.all_scores = []
+        # The backpointers at each time-step.
+        self.prev_ks = []
+        # The outputs at each time-step.
+        self.next_ys = [torch.full((size,), 0, dtype=torch.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, torch.tensor(
+            [i for i in range(int(self.scores.shape[0]))], dtype=torch.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 = torch.tensor(hyps, dtype=torch.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]))

python/rapidocr/inference_engine/pytorch/networks/transforms/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__all__ = ["build_transform"]
+
+
+def build_transform(config):
+    from .tps import TPS
+    from .stn import STN_ON
+    from .tsrn import TSRN
+    from .tbsrn import TBSRN
+
+    support_dict = ["TPS", "STN_ON", "TSRN", "TBSRN"]
+
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "transform only support {}".format(support_dict)
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

python/rapidocr/inference_engine/pytorch/networks/transforms/stn.py

@@ -0,0 +1,121 @@
+"""
+This code is refer from:
+https://github.com/ayumiymk/aster.pytorch/blob/master/lib/models/stn_head.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+import numpy as np
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+
+
+def conv3x3_block(in_channels, out_channels, stride=1):
+    n = 3 * 3 * out_channels
+    w = math.sqrt(2.0 / n)
+    conv_layer = nn.Conv2d(
+        in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True
+    )
+    block = nn.Sequential(conv_layer, nn.BatchNorm2d(out_channels), nn.ReLU())
+    return block
+
+
+class STN(nn.Module):
+    def __init__(self, in_channels, num_ctrlpoints, activation="none"):
+        super(STN, self).__init__()
+        self.in_channels = in_channels
+        self.num_ctrlpoints = num_ctrlpoints
+        self.activation = activation
+        self.stn_convnet = nn.Sequential(
+            conv3x3_block(in_channels, 32),  # 32x64
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(32, 64),  # 16x32
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(64, 128),  # 8*16
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(128, 256),  # 4*8
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(256, 256),  # 2*4,
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            conv3x3_block(256, 256),
+        )  # 1*2
+        self.stn_fc1 = nn.Sequential(
+            nn.Linear(2 * 256, 512, bias=True),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+        )
+        fc2_bias = self.init_stn()
+        self.stn_fc2 = nn.Linear(512, num_ctrlpoints * 2, bias=True)
+
+    def init_stn(self):
+        margin = 0.01
+        sampling_num_per_side = int(self.num_ctrlpoints / 2)
+        ctrl_pts_x = np.linspace(margin, 1.0 - margin, sampling_num_per_side)
+        ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
+        ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1 - margin)
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0).astype(
+            np.float32
+        )
+        if self.activation == "none":
+            pass
+        elif self.activation == "sigmoid":
+            ctrl_points = -np.log(1.0 / ctrl_points - 1.0)
+        ctrl_points = torch.Tensor(ctrl_points)
+        # fc2_bias = ctrl_points.view(-1)
+        fc2_bias = torch.reshape(
+            ctrl_points, shape=[ctrl_points.shape[0] * ctrl_points.shape[1]]
+        )
+        return fc2_bias
+
+    def forward(self, x):
+        x = self.stn_convnet(x)
+        batch_size, _, h, w = x.shape
+        # x = x.view(batch_size, -1)
+        x = torch.reshape(x, shape=(batch_size, -1))
+        img_feat = self.stn_fc1(x)
+        x = self.stn_fc2(0.1 * img_feat)
+        if self.activation == "sigmoid":
+            x = F.sigmoid(x)
+        # x = x.view(-1, self.num_ctrlpoints, 2)
+        x = torch.reshape(x, shape=[-1, self.num_ctrlpoints, 2])
+        return img_feat, x
+
+
+class STN_ON(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        tps_inputsize,
+        tps_outputsize,
+        num_control_points,
+        tps_margins,
+        stn_activation,
+    ):
+        super(STN_ON, self).__init__()
+        self.tps = TPSSpatialTransformer(
+            output_image_size=tuple(tps_outputsize),
+            num_control_points=num_control_points,
+            margins=tuple(tps_margins),
+        )
+        self.stn_head = STN(
+            in_channels=in_channels,
+            num_ctrlpoints=num_control_points,
+            activation=stn_activation,
+        )
+        self.tps_inputsize = tps_inputsize
+        self.out_channels = in_channels
+
+    def forward(self, image):
+        stn_input = torch.nn.functional.interpolate(
+            image, self.tps_inputsize, mode="bilinear", align_corners=True
+        )
+        stn_img_feat, ctrl_points = self.stn_head(stn_input)
+        x, _ = self.tps(image, ctrl_points)
+        return x

python/rapidocr/inference_engine/pytorch/networks/transforms/tbsrn.py

@@ -0,0 +1,305 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/scene-text-telescope/model/tbsrn.py
+"""
+
+import math
+import warnings
+import numpy as np
+import torch
+from torch import nn
+import string
+
+warnings.filterwarnings("ignore")
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+from .stn import STN as STNHead
+from .tsrn import GruBlock, mish, UpsampleBLock
+from pytorchocr.modeling.heads.sr_rensnet_transformer import (
+    Transformer,
+    LayerNorm,
+    PositionwiseFeedForward,
+    MultiHeadedAttention,
+)
+
+
+def positionalencoding2d(d_model, height, width):
+    """
+    :param d_model: dimension of the model
+    :param height: height of the positions
+    :param width: width of the positions
+    :return: d_model*height*width position matrix
+    """
+    if d_model % 4 != 0:
+        raise ValueError(
+            "Cannot use sin/cos positional encoding with "
+            "odd dimension (got dim={:d})".format(d_model)
+        )
+    pe = torch.zeros([d_model, height, width])
+    # Each dimension use half of d_model
+    d_model = int(d_model / 2)
+    div_term = torch.exp(torch.arange(0.0, d_model, 2) * -(math.log(10000.0) / d_model))
+    pos_w = torch.arange(0.0, width, dtype=torch.float32).unsqueeze(1)
+    pos_h = torch.arange(0.0, height, dtype=torch.float32).unsqueeze(1)
+
+    pe[0:d_model:2, :, :] = (
+        torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[1:d_model:2, :, :] = (
+        torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
+    )
+    pe[d_model::2, :, :] = (
+        torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+    pe[d_model + 1 :: 2, :, :] = (
+        torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
+    )
+
+    return pe
+
+
+class FeatureEnhancer(nn.Module):
+    def __init__(self):
+        super(FeatureEnhancer, self).__init__()
+
+        self.multihead = MultiHeadedAttention(h=4, d_model=128, dropout=0.1)
+        self.mul_layernorm1 = LayerNorm(features=128)
+
+        self.pff = PositionwiseFeedForward(128, 128)
+        self.mul_layernorm3 = LayerNorm(features=128)
+
+        self.linear = nn.Linear(128, 64)
+
+    def forward(self, conv_feature):
+        """
+        text : (batch, seq_len, embedding_size)
+        global_info: (batch, embedding_size, 1, 1)
+        conv_feature: (batch, channel, H, W)
+        """
+        batch = conv_feature.shape[0]
+        if torch.cuda.is_available():
+            position2d = (
+                positionalencoding2d(64, 16, 64)
+                .float()
+                .cuda()
+                .unsqueeze(0)
+                .reshape([1, 64, 1024])
+            )
+        else:
+            position2d = (
+                positionalencoding2d(64, 16, 64)
+                .float()
+                .unsqueeze(0)
+                .reshape([1, 64, 1024])
+            )
+        position2d = position2d.repeat(batch, 1, 1)
+        conv_feature = torch.cat(
+            [conv_feature, position2d], 1
+        )  # batch, 128(64+64), 32, 128
+        result = conv_feature.permute(0, 2, 1).contiguous()
+        origin_result = result
+        result = self.mul_layernorm1(
+            origin_result + self.multihead(result, result, result, mask=None)[0]
+        )
+        origin_result = result
+        result = self.mul_layernorm3(origin_result + self.pff(result))
+        result = self.linear(result)
+        return result.permute(0, 2, 1).contiguous()
+
+
+def str_filt(str_, voc_type):
+    alpha_dict = {
+        "digit": string.digits,
+        "lower": string.digits + string.ascii_lowercase,
+        "upper": string.digits + string.ascii_letters,
+        "all": string.digits + string.ascii_letters + string.punctuation,
+    }
+    if voc_type == "lower":
+        str_ = str_.lower()
+    for char in str_:
+        if char not in alpha_dict[voc_type]:
+            str_ = str_.replace(char, "")
+    str_ = str_.lower()
+    return str_
+
+
+class TBSRN(nn.Module):
+    def __init__(
+        self,
+        in_channels=3,
+        scale_factor=2,
+        width=128,
+        height=32,
+        STN=True,
+        srb_nums=5,
+        mask=False,
+        hidden_units=32,
+        infer_mode=False,
+    ):
+        super(TBSRN, self).__init__()
+        in_planes = 3
+        if mask:
+            in_planes = 4
+        assert math.log(scale_factor, 2) % 1 == 0
+        upsample_block_num = int(math.log(scale_factor, 2))
+        self.block1 = nn.Sequential(
+            nn.Conv2d(in_planes, 2 * hidden_units, kernel_size=9, padding=4),
+            nn.PReLU()
+            # nn.ReLU()
+        )
+        self.srb_nums = srb_nums
+        for i in range(srb_nums):
+            setattr(self, "block%d" % (i + 2), RecurrentResidualBlock(2 * hidden_units))
+
+        setattr(
+            self,
+            "block%d" % (srb_nums + 2),
+            nn.Sequential(
+                nn.Conv2d(2 * hidden_units, 2 * hidden_units, kernel_size=3, padding=1),
+                nn.BatchNorm2d(2 * hidden_units),
+            ),
+        )
+
+        # self.non_local = NonLocalBlock2D(64, 64)
+        block_ = [UpsampleBLock(2 * hidden_units, 2) for _ in range(upsample_block_num)]
+        block_.append(nn.Conv2d(2 * hidden_units, in_planes, kernel_size=9, padding=4))
+        setattr(self, "block%d" % (srb_nums + 3), nn.Sequential(*block_))
+        self.tps_inputsize = [height // scale_factor, width // scale_factor]
+        tps_outputsize = [height // scale_factor, width // scale_factor]
+        num_control_points = 20
+        tps_margins = [0.05, 0.05]
+        self.stn = STN
+        self.out_channels = in_channels
+        if self.stn:
+            self.tps = TPSSpatialTransformer(
+                output_image_size=tuple(tps_outputsize),
+                num_control_points=num_control_points,
+                margins=tuple(tps_margins),
+            )
+
+            self.stn_head = STNHead(
+                in_channels=in_planes,
+                num_ctrlpoints=num_control_points,
+                activation="none",
+            )
+        self.infer_mode = infer_mode
+
+        self.english_alphabet = (
+            "-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+        )
+        self.english_dict = {}
+        for index in range(len(self.english_alphabet)):
+            self.english_dict[self.english_alphabet[index]] = index
+        transformer = Transformer(alphabet="-0123456789abcdefghijklmnopqrstuvwxyz")
+        self.transformer = transformer
+        for param in self.transformer.parameters():
+            param.trainable = False
+
+    def label_encoder(self, label):
+        batch = len(label)
+
+        length = [len(i) for i in label]
+        length_tensor = torch.Tensor(length).type(torch.int64)
+
+        max_length = max(length)
+        input_tensor = np.zeros((batch, max_length))
+        for i in range(batch):
+            for j in range(length[i] - 1):
+                input_tensor[i][j + 1] = self.english_dict[label[i][j]]
+
+        text_gt = []
+        for i in label:
+            for j in i:
+                text_gt.append(self.english_dict[j])
+        text_gt = torch.Tensor(text_gt).type(torch.int64)
+
+        input_tensor = torch.Tensor(input_tensor).type(torch.int64)
+        return length_tensor, input_tensor, text_gt
+
+    def forward(self, x):
+        output = {}
+        if self.infer_mode:
+            output["lr_img"] = x
+            y = x
+        else:
+            output["lr_img"] = x[0]
+            output["hr_img"] = x[1]
+            y = x[0]
+        if self.stn and self.training:
+            _, ctrl_points_x = self.stn_head(y)
+            y, _ = self.tps(y, ctrl_points_x)
+        block = {"1": self.block1(y)}
+        for i in range(self.srb_nums + 1):
+            block[str(i + 2)] = getattr(self, "block%d" % (i + 2))(block[str(i + 1)])
+
+        block[str(self.srb_nums + 3)] = getattr(self, "block%d" % (self.srb_nums + 3))(
+            (block["1"] + block[str(self.srb_nums + 2)])
+        )
+
+        sr_img = torch.tanh(block[str(self.srb_nums + 3)])
+        output["sr_img"] = sr_img
+
+        if self.training:
+            hr_img = x[1]
+
+            # add transformer
+            label = [str_filt(i, "lower") + "-" for i in x[2]]
+            length_tensor, input_tensor, text_gt = self.label_encoder(label)
+            hr_pred, word_attention_map_gt, hr_correct_list = self.transformer(
+                hr_img, length_tensor, input_tensor
+            )
+            sr_pred, word_attention_map_pred, sr_correct_list = self.transformer(
+                sr_img, length_tensor, input_tensor
+            )
+            output["hr_img"] = hr_img
+            output["hr_pred"] = hr_pred
+            output["text_gt"] = text_gt
+            output["word_attention_map_gt"] = word_attention_map_gt
+            output["sr_pred"] = sr_pred
+            output["word_attention_map_pred"] = word_attention_map_pred
+
+        return output
+
+
+class RecurrentResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(RecurrentResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.gru1 = GruBlock(channels, channels)
+        # self.prelu = nn.ReLU()
+        self.prelu = mish()
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.gru2 = GruBlock(channels, channels)
+        self.feature_enhancer = FeatureEnhancer()
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = self.bn1(residual)
+        residual = self.prelu(residual)
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+
+        size = residual.shape
+        residual = residual.reshape([size[0], size[1], -1])
+        residual = self.feature_enhancer(residual)
+        residual = residual.reshape([size[0], size[1], size[2], size[3]])
+        return x + residual

python/rapidocr/inference_engine/pytorch/networks/transforms/tps.py

@@ -0,0 +1,292 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pytorchocr.modeling.common import Activation
+
+# import paddle
+# from paddle import nn, ParamAttr
+# from paddle.nn import functional as F
+import numpy as np
+
+
+class ConvBNLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        groups=1,
+        act=None,
+        name=None,
+    ):
+        super(ConvBNLayer, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            bias=False,
+        )
+        bn_name = "bn_" + name
+        self.bn = nn.BatchNorm2d(
+            out_channels,
+        )
+        self.act = act
+        if act is not None:
+            self._act = Activation(act)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        if self.act is not None:
+            x = self._act(x)
+        return x
+
+
+class LocalizationNetwork(nn.Module):
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(LocalizationNetwork, self).__init__()
+        self.F = num_fiducial
+        F = num_fiducial
+        if model_name == "large":
+            num_filters_list = [64, 128, 256, 512]
+            fc_dim = 256
+        else:
+            num_filters_list = [16, 32, 64, 128]
+            fc_dim = 64
+
+        # self.block_list = []
+        self.block_list = nn.Sequential()
+        for fno in range(0, len(num_filters_list)):
+            num_filters = num_filters_list[fno]
+            name = "loc_conv%d" % fno
+            # conv = self.add_sublayer(
+            #     name,
+            #     ConvBNLayer(
+            #         in_channels=in_channels,
+            #         out_channels=num_filters,
+            #         kernel_size=3,
+            #         act='relu',
+            #         name=name))
+            conv = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=num_filters,
+                kernel_size=3,
+                act="relu",
+                name=name,
+            )
+            # self.block_list.append(conv)
+            self.block_list.add_module(name, conv)
+            if fno == len(num_filters_list) - 1:
+                pool = nn.AdaptiveAvgPool2d(1)
+            else:
+                # pool = nn.MaxPool2D(kernel_size=2, stride=2, padding=0)
+                pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+            in_channels = num_filters
+            # self.block_list.append(pool)
+            self.block_list.add_module("{}_pool".format(name), pool)
+        name = "loc_fc1"
+        stdv = 1.0 / math.sqrt(num_filters_list[-1] * 1.0)
+        self.fc1 = nn.Linear(
+            in_channels,
+            fc_dim,
+            bias=True,
+        )
+
+        # Init fc2 in LocalizationNetwork
+        initial_bias = self.get_initial_fiducials()
+        initial_bias = initial_bias.reshape(-1)
+        name = "loc_fc2"
+        self.fc2 = nn.Linear(fc_dim, F * 2, bias=True)
+        self.out_channels = F * 2
+
+    def forward(self, x):
+        """
+        Estimating parameters of geometric transformation
+        Args:
+            image: input
+        Return:
+            batch_C_prime: the matrix of the geometric transformation
+        """
+        B = x.shape[0]
+        i = 0
+        for block in self.block_list:
+            x = block(x)
+        x = x.squeeze(dim=2).squeeze(dim=2)
+        x = self.fc1(x)
+
+        x = F.relu(x)
+        x = self.fc2(x)
+        x = x.reshape(shape=[-1, self.F, 2])
+        return x
+
+    def get_initial_fiducials(self):
+        """see RARE paper Fig. 6 (a)"""
+        F = self.F
+        ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2))
+        ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2))
+        ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2))
+        ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+        ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+        initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+        return initial_bias
+
+
+class GridGenerator(nn.Module):
+    def __init__(self, in_channels, num_fiducial):
+        super(GridGenerator, self).__init__()
+        self.eps = 1e-6
+        self.F = num_fiducial
+
+        name = "ex_fc"
+        self.fc = nn.Linear(in_channels, 6, bias=True)
+
+    def forward(self, batch_C_prime, I_r_size):
+        """
+        Generate the grid for the grid_sampler.
+        Args:
+            batch_C_prime: the matrix of the geometric transformation
+            I_r_size: the shape of the input image
+        Return:
+            batch_P_prime: the grid for the grid_sampler
+        """
+        C = self.build_C_paddle()
+        P = self.build_P_paddle(I_r_size)
+
+        inv_delta_C_tensor = self.build_inv_delta_C_paddle(C).type(torch.float32)
+        P_hat_tensor = self.build_P_hat_paddle(C, torch.as_tensor(P)).type(
+            torch.float32
+        )
+
+        inv_delta_C_tensor.stop_gradient = True
+        P_hat_tensor.stop_gradient = True
+
+        batch_C_ex_part_tensor = self.get_expand_tensor(batch_C_prime)
+
+        batch_C_ex_part_tensor.stop_gradient = True
+
+        batch_C_prime_with_zeros = torch.cat(
+            [batch_C_prime, batch_C_ex_part_tensor], dim=1
+        )
+        inv_delta_C_tensor = inv_delta_C_tensor.to(batch_C_prime_with_zeros.device)
+        batch_T = torch.matmul(inv_delta_C_tensor, batch_C_prime_with_zeros)
+        P_hat_tensor = P_hat_tensor.to(batch_T.device)
+        batch_P_prime = torch.matmul(P_hat_tensor, batch_T)
+        return batch_P_prime
+
+    def build_C_paddle(self):
+        """Return coordinates of fiducial points in I_r; C"""
+        F = self.F
+        ctrl_pts_x = torch.linspace(-1.0, 1.0, int(F / 2), dtype=torch.float64)
+        ctrl_pts_y_top = -1 * torch.ones([int(F / 2)], dtype=torch.float64)
+        ctrl_pts_y_bottom = torch.ones([int(F / 2)], dtype=torch.float64)
+        ctrl_pts_top = torch.stack([ctrl_pts_x, ctrl_pts_y_top], dim=1)
+        ctrl_pts_bottom = torch.stack([ctrl_pts_x, ctrl_pts_y_bottom], dim=1)
+        C = torch.cat([ctrl_pts_top, ctrl_pts_bottom], dim=0)
+        return C  # F x 2
+
+    def build_P_paddle(self, I_r_size):
+        I_r_height, I_r_width = I_r_size
+        I_r_grid_x = (
+            torch.arange(-I_r_width, I_r_width, 2, dtype=torch.float64) + 1.0
+        ) / torch.as_tensor(np.array([I_r_width]).astype(np.float64))
+
+        I_r_grid_y = (
+            torch.arange(-I_r_height, I_r_height, 2, dtype=torch.float64) + 1.0
+        ) / torch.as_tensor(np.array([I_r_height]).astype(np.float64))
+
+        # P: self.I_r_width x self.I_r_height x 2
+        P = torch.stack(torch.meshgrid([I_r_grid_x, I_r_grid_y]), dim=2)
+        # P = paddle.transpose(P, perm=[1, 0, 2])
+        P = P.permute(1, 0, 2)
+        # n (= self.I_r_width x self.I_r_height) x 2
+        return P.reshape([-1, 2])
+
+    def build_inv_delta_C_paddle(self, C):
+        """Return inv_delta_C which is needed to calculate T"""
+        F = self.F
+        hat_C = torch.zeros((F, F), dtype=torch.float64)  # F x F
+        for i in range(0, F):
+            for j in range(i, F):
+                if i == j:
+                    hat_C[i, j] = 1
+                else:
+                    r = torch.norm(C[i] - C[j])
+                    hat_C[i, j] = r
+                    hat_C[j, i] = r
+        hat_C = (hat_C**2) * torch.log(hat_C)
+        delta_C = torch.cat(  # F+3 x F+3
+            [
+                torch.cat(
+                    [torch.ones((F, 1), dtype=torch.float64), C, hat_C], dim=1
+                ),  # F x F+3
+                torch.cat(
+                    [torch.zeros((2, 3), dtype=torch.float64), C.permute(1, 0)], dim=1
+                ),  # 2 x F+3
+                torch.cat(
+                    [
+                        torch.zeros((1, 3), dtype=torch.float64),
+                        torch.ones((1, F), dtype=torch.float64),
+                    ],
+                    dim=1,
+                ),  # 1 x F+3
+            ],
+            dim=0,
+        )
+        inv_delta_C = torch.inverse(delta_C)
+        return inv_delta_C  # F+3 x F+3
+
+    def build_P_hat_paddle(self, C, P):
+        F = self.F
+        eps = self.eps
+        n = P.shape[0]  # n (= self.I_r_width x self.I_r_height)
+        # P_tile: n x 2 -> n x 1 x 2 -> n x F x 2
+        # P_tile = paddle.tile(paddle.unsqueeze(P, axis=1), (1, F, 1))
+        P_tile = torch.unsqueeze(P, dim=1).repeat(1, F, 1)
+        C_tile = torch.unsqueeze(C, dim=0)  # 1 x F x 2
+        P_diff = P_tile - C_tile  # n x F x 2
+        # rbf_norm: n x F
+        rbf_norm = torch.norm(P_diff, p=2, dim=2, keepdim=False)
+
+        # rbf: n x F
+        # rbf = torch.mul(
+        #     torch.square(rbf_norm), torch.log(rbf_norm + eps))
+        rbf = torch.mul(rbf_norm**2, torch.log(rbf_norm + eps))
+        P_hat = torch.cat([torch.ones((n, 1), dtype=torch.float64), P, rbf], dim=1)
+        return P_hat  # n x F+3
+
+    def get_expand_tensor(self, batch_C_prime):
+        B, H, C = batch_C_prime.shape
+        batch_C_prime = batch_C_prime.reshape([B, H * C])
+        batch_C_ex_part_tensor = self.fc(batch_C_prime)
+        batch_C_ex_part_tensor = batch_C_ex_part_tensor.reshape([-1, 3, 2])
+        return batch_C_ex_part_tensor
+
+
+class TPS(nn.Module):
+    def __init__(self, in_channels, num_fiducial, loc_lr, model_name):
+        super(TPS, self).__init__()
+        self.loc_net = LocalizationNetwork(
+            in_channels, num_fiducial, loc_lr, model_name
+        )
+        self.grid_generator = GridGenerator(self.loc_net.out_channels, num_fiducial)
+        self.out_channels = in_channels
+
+    def forward(self, image):
+        image.stop_gradient = False
+        batch_C_prime = self.loc_net(image)
+        batch_P_prime = self.grid_generator(batch_C_prime, image.shape[2:])
+        batch_P_prime = batch_P_prime.reshape([-1, image.shape[2], image.shape[3], 2])
+        if torch.__version__ < "1.3.0":
+            batch_I_r = F.grid_sample(image, grid=batch_P_prime)
+        else:
+            batch_I_r = F.grid_sample(image, grid=batch_P_prime, align_corners=True)
+        return batch_I_r

python/rapidocr/inference_engine/pytorch/networks/transforms/tps_spatial_transformer.py

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

python/rapidocr/inference_engine/pytorch/networks/transforms/tsrn.py

@@ -0,0 +1,220 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/text-gestalt/model/tsrn.py
+"""
+
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+import warnings
+import math
+
+warnings.filterwarnings("ignore")
+
+from .tps_spatial_transformer import TPSSpatialTransformer
+from .stn import STN as STN_model
+from pytorchocr.modeling.heads.sr_rensnet_transformer import Transformer
+
+
+class TSRN(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        scale_factor=2,
+        width=128,
+        height=32,
+        STN=False,
+        srb_nums=5,
+        mask=False,
+        hidden_units=32,
+        infer_mode=False,
+        **kwargs,
+    ):
+        super(TSRN, self).__init__()
+        in_planes = 3
+        if mask:
+            in_planes = 4
+        assert math.log(scale_factor, 2) % 1 == 0
+        upsample_block_num = int(math.log(scale_factor, 2))
+        self.block1 = nn.Sequential(
+            nn.Conv2d(in_planes, 2 * hidden_units, kernel_size=9, padding=4), nn.PReLU()
+        )
+        self.srb_nums = srb_nums
+        for i in range(srb_nums):
+            setattr(self, "block%d" % (i + 2), RecurrentResidualBlock(2 * hidden_units))
+
+        setattr(
+            self,
+            "block%d" % (srb_nums + 2),
+            nn.Sequential(
+                nn.Conv2d(2 * hidden_units, 2 * hidden_units, kernel_size=3, padding=1),
+                nn.BatchNorm2d(2 * hidden_units),
+            ),
+        )
+
+        block_ = [UpsampleBLock(2 * hidden_units, 2) for _ in range(upsample_block_num)]
+        block_.append(nn.Conv2d(2 * hidden_units, in_planes, kernel_size=9, padding=4))
+        setattr(self, "block%d" % (srb_nums + 3), nn.Sequential(*block_))
+        self.tps_inputsize = [height // scale_factor, width // scale_factor]
+        tps_outputsize = [height // scale_factor, width // scale_factor]
+        num_control_points = 20
+        tps_margins = [0.05, 0.05]
+        self.stn = STN
+        if self.stn:
+            self.tps = TPSSpatialTransformer(
+                output_image_size=tuple(tps_outputsize),
+                num_control_points=num_control_points,
+                margins=tuple(tps_margins),
+            )
+
+            self.stn_head = STN_model(
+                in_channels=in_planes,
+                num_ctrlpoints=num_control_points,
+                activation="none",
+            )
+        self.out_channels = in_channels
+
+        self.r34_transformer = Transformer()
+        for param in self.r34_transformer.parameters():
+            param.trainable = False
+        self.infer_mode = infer_mode
+
+    def forward(self, x):
+        output = {}
+        if self.infer_mode:
+            output["lr_img"] = x
+            y = x
+        else:
+            output["lr_img"] = x[0]
+            output["hr_img"] = x[1]
+            y = x[0]
+        if self.stn and self.training:
+            _, ctrl_points_x = self.stn_head(y)
+            y, _ = self.tps(y, ctrl_points_x)
+        block = {"1": self.block1(y)}
+        for i in range(self.srb_nums + 1):
+            block[str(i + 2)] = getattr(self, "block%d" % (i + 2))(block[str(i + 1)])
+
+        block[str(self.srb_nums + 3)] = getattr(self, "block%d" % (self.srb_nums + 3))(
+            (block["1"] + block[str(self.srb_nums + 2)])
+        )
+
+        sr_img = torch.tanh(block[str(self.srb_nums + 3)])
+
+        output["sr_img"] = sr_img
+
+        if self.training:
+            hr_img = x[1]
+            length = x[2]
+            input_tensor = x[3]
+
+            # add transformer
+            sr_pred, word_attention_map_pred, _ = self.r34_transformer(
+                sr_img, length, input_tensor
+            )
+
+            hr_pred, word_attention_map_gt, _ = self.r34_transformer(
+                hr_img, length, input_tensor
+            )
+
+            output["hr_img"] = hr_img
+            output["hr_pred"] = hr_pred
+            output["word_attention_map_gt"] = word_attention_map_gt
+            output["sr_pred"] = sr_pred
+            output["word_attention_map_pred"] = word_attention_map_pred
+
+        return output
+
+
+class RecurrentResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(RecurrentResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(channels)
+        self.gru1 = GruBlock(channels, channels)
+        self.prelu = mish()
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(channels)
+        self.gru2 = GruBlock(channels, channels)
+
+    def forward(self, x):
+        residual = self.conv1(x)
+        residual = self.bn1(residual)
+        residual = self.prelu(residual)
+        residual = self.conv2(residual)
+        residual = self.bn2(residual)
+        residual = (
+            self.gru1(residual.permute(0, 1, 3, 2).contiguous())
+            .permute(0, 1, 3, 2)
+            .contiguous()
+        )
+
+        return self.gru2(x + residual).contiguous()
+
+
+class UpsampleBLock(nn.Module):
+    def __init__(self, in_channels, up_scale):
+        super(UpsampleBLock, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, in_channels * up_scale**2, kernel_size=3, padding=1
+        )
+
+        self.pixel_shuffle = nn.PixelShuffle(up_scale)
+        self.prelu = mish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.pixel_shuffle(x)
+        x = self.prelu(x)
+        return x
+
+
+class mish(nn.Module):
+    def __init__(
+        self,
+    ):
+        super(mish, self).__init__()
+        self.activated = True
+
+    def forward(self, x):
+        if self.activated:
+            x = x * (torch.tanh(F.softplus(x)))
+        return x
+
+
+class GruBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(GruBlock, self).__init__()
+        assert out_channels % 2 == 0
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
+        self.gru = nn.GRU(
+            out_channels,
+            out_channels // 2,
+            bidirectional=True,
+            batch_first=True,
+        )
+
+    def forward(self, x):
+        # x: b, c, w, h
+        x = self.conv1(x)
+        x = x.permute(0, 2, 3, 1).contiguous()  # b, w, h, c
+        batch_size, w, h, c = x.size()
+        x = x.view(batch_size * w, h, c)  # b*w, h, c
+        x, _ = self.gru(x)
+        x = x.view(batch_size, w, h, c)
+        x = x.permute(0, 3, 1, 2).contiguous()
+        return x