openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/common.py

@@ -0,0 +1,238 @@
+import torch
+import torch.nn as nn
+
+
+class GELU(nn.Module):
+
+    def __init__(self, inplace=True):
+        super(GELU, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        return torch.nn.functional.gelu(x)
+
+
+class Swish(nn.Module):
+
+    def __init__(self, inplace=True):
+        super(Swish, self).__init__()
+        self.inplace = inplace
+
+    def forward(self, x):
+        if self.inplace:
+            x.mul_(torch.sigmoid(x))
+            return x
+        else:
+            return x * torch.sigmoid(x)
+
+
+class Activation(nn.Module):
+
+    def __init__(self, act_type, inplace=True):
+        super(Activation, self).__init__()
+        act_type = act_type.lower()
+        if act_type == 'relu':
+            self.act = nn.ReLU(inplace=inplace)
+        elif act_type == 'relu6':
+            self.act = nn.ReLU6(inplace=inplace)
+        elif act_type == 'sigmoid':
+            self.act = nn.Sigmoid()
+        elif act_type == 'hard_sigmoid':
+            self.act = nn.Hardsigmoid(inplace)
+        elif act_type == 'hard_swish':
+            self.act = nn.Hardswish(inplace=inplace)
+        elif act_type == 'leakyrelu':
+            self.act = nn.LeakyReLU(inplace=inplace)
+        elif act_type == 'gelu':
+            self.act = GELU(inplace=inplace)
+        elif act_type == 'swish':
+            self.act = Swish(inplace=inplace)
+        else:
+            raise NotImplementedError
+
+    def forward(self, inputs):
+        return self.act(inputs)
+
+
+def drop_path(x,
+              drop_prob: float = 0.0,
+              training: bool = False,
+              scale_by_keep: bool = True):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of
+    residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0], ) + (1, ) * (
+        x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of
+    residual blocks)."""
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f'drop_prob={round(self.drop_prob,3):0.3f}'
+
+
+class Identity(nn.Module):
+
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, input):
+        return input
+
+
+class Mlp(nn.Module):
+
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads=8,
+                 qkv_bias=False,
+                 qk_scale=None,
+                 attn_drop=0.0,
+                 proj_drop=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (self.qkv(x).reshape(B, N, 3, self.num_heads,
+                                   C // self.num_heads).permute(2, 0, 3, 1, 4))
+        q, k, v = qkv[0], qkv[1], qkv[
+            2]  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim,
+                       hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer,
+                       drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding."""
+
+    def __init__(self,
+                 img_size=[32, 128],
+                 patch_size=[4, 4],
+                 in_chans=3,
+                 embed_dim=768):
+        super().__init__()
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] //
+                                                        patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans,
+                              embed_dim,
+                              kernel_size=patch_size,
+                              stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x

openocr/openrec/modeling/decoders/__init__.py

@@ -0,0 +1,109 @@
+import torch.nn as nn
+
+__all__ = ['build_decoder']
+
+
+def build_decoder(config):
+    # rec decoder
+    from .abinet_decoder import ABINetDecoder
+    from .aster_decoder import ASTERDecoder
+    from .cdistnet_decoder import CDistNetDecoder
+    from .cppd_decoder import CPPDDecoder
+    from .rctc_decoder import RCTCDecoder
+    from .ctc_decoder import CTCDecoder
+    from .dan_decoder import DANDecoder
+    from .igtr_decoder import IGTRDecoder
+    from .lister_decoder import LISTERDecoder
+    from .lpv_decoder import LPVDecoder
+    from .mgp_decoder import MGPDecoder
+    from .nrtr_decoder import NRTRDecoder
+    from .parseq_decoder import PARSeqDecoder
+    from .robustscanner_decoder import RobustScannerDecoder
+    from .sar_decoder import SARDecoder
+    from .smtr_decoder import SMTRDecoder
+    from .smtr_decoder_nattn import SMTRDecoderNumAttn
+    from .srn_decoder import SRNDecoder
+    from .visionlan_decoder import VisionLANDecoder
+    from .matrn_decoder import MATRNDecoder
+    from .cam_decoder import CAMDecoder
+    from .ote_decoder import OTEDecoder
+    from .bus_decoder import BUSDecoder
+
+    support_dict = [
+        'CTCDecoder', 'NRTRDecoder', 'CPPDDecoder', 'ABINetDecoder',
+        'CDistNetDecoder', 'VisionLANDecoder', 'PARSeqDecoder', 'IGTRDecoder',
+        'SMTRDecoder', 'LPVDecoder', 'SARDecoder', 'RobustScannerDecoder',
+        'SRNDecoder', 'ASTERDecoder', 'RCTCDecoder', 'LISTERDecoder',
+        'GTCDecoder', 'SMTRDecoderNumAttn', 'MATRNDecoder', 'MGPDecoder',
+        'DANDecoder', 'CAMDecoder', 'OTEDecoder', 'BUSDecoder'
+    ]
+
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception(
+        'decoder only support {}'.format(support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class
+
+
+class GTCDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 gtc_decoder,
+                 ctc_decoder,
+                 detach=True,
+                 infer_gtc=False,
+                 out_channels=0,
+                 **kwargs):
+        super(GTCDecoder, self).__init__()
+        self.detach = detach
+        self.infer_gtc = infer_gtc
+        if infer_gtc:
+            gtc_decoder['out_channels'] = out_channels[0]
+            ctc_decoder['out_channels'] = out_channels[1]
+            gtc_decoder['in_channels'] = in_channels
+            ctc_decoder['in_channels'] = in_channels
+            self.gtc_decoder = build_decoder(gtc_decoder)
+        else:
+            ctc_decoder['in_channels'] = in_channels
+            ctc_decoder['out_channels'] = out_channels
+        self.ctc_decoder = build_decoder(ctc_decoder)
+
+    def forward(self, x, data=None):
+        ctc_pred = self.ctc_decoder(x.detach() if self.detach else x,
+                                    data=data)
+        if self.training or self.infer_gtc:
+            gtc_pred = self.gtc_decoder(x.flatten(2).transpose(1, 2),
+                                        data=data)
+            return {'gtc_pred': gtc_pred, 'ctc_pred': ctc_pred}
+        else:
+            return ctc_pred
+
+
+class GTCDecoderTwo(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 gtc_decoder,
+                 ctc_decoder,
+                 infer_gtc=False,
+                 out_channels=0,
+                 **kwargs):
+        super(GTCDecoderTwo, self).__init__()
+        self.infer_gtc = infer_gtc
+        gtc_decoder['out_channels'] = out_channels[0]
+        ctc_decoder['out_channels'] = out_channels[1]
+        gtc_decoder['in_channels'] = in_channels
+        ctc_decoder['in_channels'] = in_channels
+        self.gtc_decoder = build_decoder(gtc_decoder)
+        self.ctc_decoder = build_decoder(ctc_decoder)
+
+    def forward(self, x, data=None):
+        x_ctc, x_gtc = x
+        ctc_pred = self.ctc_decoder(x_ctc, data=data)
+        if self.training or self.infer_gtc:
+            gtc_pred = self.gtc_decoder(x_gtc.flatten(2).transpose(1, 2),
+                                        data=data)
+            return {'gtc_pred': gtc_pred, 'ctc_pred': ctc_pred}
+        else:
+            return ctc_pred

openocr/openrec/modeling/decoders/abinet_decoder.py

@@ -0,0 +1,283 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from openrec.modeling.decoders.nrtr_decoder import PositionalEncoding, TransformerBlock
+
+
+class BCNLanguage(nn.Module):
+
+    def __init__(
+        self,
+        d_model=512,
+        nhead=8,
+        num_layers=4,
+        dim_feedforward=2048,
+        dropout=0.0,
+        max_length=25,
+        detach=True,
+        num_classes=37,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.detach = detach
+        self.max_length = max_length + 1
+
+        self.proj = nn.Linear(num_classes, d_model, False)
+        self.token_encoder = PositionalEncoding(dropout=0.1,
+                                                dim=d_model,
+                                                max_len=self.max_length)
+        self.pos_encoder = PositionalEncoding(dropout=0,
+                                              dim=d_model,
+                                              max_len=self.max_length)
+        self.decoder = nn.ModuleList([
+            TransformerBlock(
+                d_model=d_model,
+                nhead=nhead,
+                dim_feedforward=dim_feedforward,
+                attention_dropout_rate=dropout,
+                residual_dropout_rate=dropout,
+                with_self_attn=False,
+                with_cross_attn=True,
+            ) for i in range(num_layers)
+        ])
+
+        self.cls = nn.Linear(d_model, num_classes)
+
+    def forward(self, tokens, lengths):
+        """
+        Args:
+            tokens: (N, T, C) where T is length, N is batch size and C is classes number
+            lengths: (N,)
+        """
+        if self.detach:
+            tokens = tokens.detach()
+        embed = self.proj(tokens)  # (N, T, E)
+        embed = self.token_encoder(embed)  # (N, T, E)
+        mask = _get_mask(lengths, self.max_length)  # (N, 1, T, T)
+        zeros = embed.new_zeros(*embed.shape)
+        qeury = self.pos_encoder(zeros)
+        for decoder_layer in self.decoder:
+            qeury = decoder_layer(qeury, embed, cross_mask=mask)
+        output = qeury  # (N, T, E)
+
+        logits = self.cls(output)  # (N, T, C)
+        return output, logits
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c), nn.ReLU(True))
+
+
+class DecoderUpsample(nn.Module):
+
+    def __init__(self, in_c, out_c, k=3, s=1, p=1, mode='nearest') -> None:
+        super().__init__()
+        self.align_corners = None if mode == 'nearest' else True
+        self.mode = mode
+        # nn.Upsample(size=size, scale_factor=scale_factor, mode=mode, align_corners=align_corners),
+        self.w = nn.Sequential(
+            nn.Conv2d(in_c, out_c, k, s, p),
+            nn.BatchNorm2d(out_c),
+            nn.ReLU(True),
+        )
+
+    def forward(self, x, size):
+        x = F.interpolate(x,
+                          size=size,
+                          mode=self.mode,
+                          align_corners=self.align_corners)
+        return self.w(x)
+
+
+class PositionAttention(nn.Module):
+
+    def __init__(self,
+                 max_length,
+                 in_channels=512,
+                 num_channels=64,
+                 mode='nearest',
+                 **kwargs):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+        )
+        self.k_decoder = nn.ModuleList([
+            DecoderUpsample(num_channels, num_channels, mode=mode),
+            DecoderUpsample(num_channels, num_channels, mode=mode),
+            DecoderUpsample(num_channels, num_channels, mode=mode),
+            DecoderUpsample(num_channels, in_channels, mode=mode),
+        ])
+
+        self.pos_encoder = PositionalEncoding(dropout=0,
+                                              dim=in_channels,
+                                              max_len=max_length)
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x, query=None):
+        N, E, H, W = x.size()
+        k, v = x, x  # (N, E, H, W)
+
+        # calculate key vector
+        features = []
+        size_decoder = []
+        for i in range(0, len(self.k_encoder)):
+            size_decoder.append(k.shape[2:])
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k, size=size_decoder[-(i + 1)])
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k, size=size_decoder[0])  # (N, E, H, W)
+        # calculate query vector
+        # TODO q=f(q,k)
+        zeros = x.new_zeros(
+            (N, self.max_length, E)) if query is None else query  # (N, T, E)
+        q = self.pos_encoder(zeros)  # (N, T, E)
+        q = self.project(q)  # (N, T, E)
+
+        # calculate attention
+        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
+        attn_scores = attn_scores / (E**0.5)
+        attn_scores = F.softmax(attn_scores, dim=-1)
+
+        # (N, E, H, W) -> (N, H, W, E) -> (N, (H*W), E)
+        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
+        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)
+        return attn_vecs, attn_scores.view(N, -1, H, W)
+
+
+class ABINetDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 nhead=8,
+                 num_layers=3,
+                 dim_feedforward=2048,
+                 dropout=0.1,
+                 max_length=25,
+                 iter_size=3,
+                 **kwargs):
+        super().__init__()
+        self.max_length = max_length + 1
+        d_model = in_channels
+        self.pos_encoder = PositionalEncoding(dropout=0.1, dim=d_model)
+        self.encoder = nn.ModuleList([
+            TransformerBlock(
+                d_model=d_model,
+                nhead=nhead,
+                dim_feedforward=dim_feedforward,
+                attention_dropout_rate=dropout,
+                residual_dropout_rate=dropout,
+                with_self_attn=True,
+                with_cross_attn=False,
+            ) for _ in range(num_layers)
+        ])
+        self.decoder = PositionAttention(
+            max_length=self.max_length,  # additional stop token
+            in_channels=d_model,
+            num_channels=d_model // 8,
+            mode='nearest',
+        )
+        self.out_channels = out_channels
+        self.cls = nn.Linear(d_model, self.out_channels)
+        self.iter_size = iter_size
+        if iter_size > 0:
+            self.language = BCNLanguage(
+                d_model=d_model,
+                nhead=nhead,
+                num_layers=4,
+                dim_feedforward=dim_feedforward,
+                dropout=dropout,
+                max_length=max_length,
+                num_classes=self.out_channels,
+            )
+            # alignment
+            self.w_att_align = nn.Linear(2 * d_model, d_model)
+            self.cls_align = nn.Linear(d_model, self.out_channels)
+
+    def forward(self, x, data=None):
+        # bs, c, h, w
+        x = x.permute([0, 2, 3, 1])  # bs, h, w, c
+        _, H, W, C = x.shape
+        # assert H % 8 == 0 and W % 16 == 0, 'The height and width should be multiples of 8 and 16.'
+        feature = x.flatten(1, 2)  # bs, h*w, c
+        feature = self.pos_encoder(feature)  # bs, h*w, c
+        for encoder_layer in self.encoder:
+            feature = encoder_layer(feature)
+        # bs, h*w, c
+        feature = feature.reshape([-1, H, W, C]).permute(0, 3, 1,
+                                                         2)  # bs, c, h, w
+        v_feature, _ = self.decoder(feature)  # (bs[N], T, E)
+        vis_logits = self.cls(v_feature)  # (bs[N], T, E)
+        align_lengths = _get_length(vis_logits)
+        align_logits = vis_logits
+        all_l_res, all_a_res = [], []
+        for _ in range(self.iter_size):
+            tokens = F.softmax(align_logits, dim=-1)
+            lengths = torch.clamp(
+                align_lengths, 2,
+                self.max_length)  # TODO: move to language model
+            l_feature, l_logits = self.language(tokens, lengths)
+
+            # alignment
+            all_l_res.append(l_logits)
+            fuse = torch.cat((l_feature, v_feature), -1)
+            f_att = torch.sigmoid(self.w_att_align(fuse))
+            output = f_att * v_feature + (1 - f_att) * l_feature
+            align_logits = self.cls_align(output)
+
+            align_lengths = _get_length(align_logits)
+            all_a_res.append(align_logits)
+        if self.training:
+            return {
+                'align': all_a_res,
+                'lang': all_l_res,
+                'vision': vis_logits
+            }
+        else:
+            return F.softmax(align_logits, -1)
+
+
+def _get_length(logit):
+    """Greed decoder to obtain length from logit."""
+    out = logit.argmax(dim=-1) == 0
+    non_zero_mask = out.int() != 0
+    mask_max_values, mask_max_indices = torch.max(non_zero_mask.int(), dim=-1)
+    mask_max_indices[mask_max_values == 0] = -1
+    out = mask_max_indices + 1
+    return out
+
+
+def _get_mask(length, max_length):
+    """Generate a square mask for the sequence.
+
+    The masked positions are filled with float('-inf'). Unmasked positions are
+    filled with float(0.0).
+    """
+    length = length.unsqueeze(-1)
+    N = length.size(0)
+    grid = torch.arange(0, max_length, device=length.device).unsqueeze(0)
+    zero_mask = torch.zeros([N, max_length],
+                            dtype=torch.float32,
+                            device=length.device)
+    inf_mask = torch.full([N, max_length],
+                          float('-inf'),
+                          dtype=torch.float32,
+                          device=length.device)
+    diag_mask = torch.diag(
+        torch.full([max_length],
+                   float('-inf'),
+                   dtype=torch.float32,
+                   device=length.device),
+        diagonal=0,
+    )
+    mask = torch.where(grid >= length, inf_mask, zero_mask)
+    mask = mask.unsqueeze(1) + diag_mask
+    return mask.unsqueeze(1)