openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/decoders/srn_decoder.py

@@ -0,0 +1,283 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .nrtr_decoder import Embeddings, TransformerBlock
+
+
+class PVAM(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 char_num,
+                 max_text_length,
+                 num_heads,
+                 hidden_dims,
+                 dropout_rate=0):
+        super(PVAM, self).__init__()
+        self.char_num = char_num
+        self.max_length = max_text_length
+        self.num_heads = num_heads
+        self.hidden_dims = hidden_dims
+        self.dropout_rate = dropout_rate
+        #TODO
+        self.emb = nn.Embedding(num_embeddings=256,
+                                embedding_dim=hidden_dims,
+                                sparse=False)
+        self.drop_out = nn.Dropout(dropout_rate)
+        self.feat_emb = nn.Linear(in_channels, in_channels)
+        self.token_emb = nn.Embedding(max_text_length, in_channels)
+        self.score = nn.Linear(in_channels, 1, bias=False)
+
+    def feat_pos_mix(self, conv_features, encoder_word_pos, dropout_rate):
+        #b h*w c
+        pos_emb = self.emb(encoder_word_pos)
+        # pos_emb = pos_emb.detach()
+        enc_input = conv_features + pos_emb
+
+        if dropout_rate:
+            enc_input = self.drop_out(enc_input)
+
+        return enc_input
+
+    def forward(self, inputs):
+        b, c, h, w = inputs.shape
+        conv_features = inputs.view(-1, c, h * w)
+        conv_features = conv_features.permute(0, 2, 1).contiguous()
+        # b h*w c
+
+        # transformer encoder
+        b, t, c = conv_features.shape
+
+        encoder_feat_pos = torch.arange(t, dtype=torch.long).to(inputs.device)
+
+        enc_inputs = self.feat_pos_mix(conv_features, encoder_feat_pos,
+                                       self.dropout_rate)
+
+        inputs = self.feat_emb(enc_inputs)  # feat emb
+
+        inputs = inputs.unsqueeze(1).expand(-1, self.max_length, -1, -1)
+
+        # b maxlen h*w c
+
+        tokens_pos = torch.arange(self.max_length,
+                                  dtype=torch.long).to(inputs.device)
+        tokens_pos = tokens_pos.unsqueeze(0).expand(b, -1)
+
+        tokens_pos_emd = self.token_emb(tokens_pos)
+        tokens_pos_emd = tokens_pos_emd.unsqueeze(2).expand(-1, -1, t, -1)
+        # b maxlen h*w c
+
+        attention_weight = torch.tanh(tokens_pos_emd + inputs)
+
+        attention_weight = torch.squeeze(self.score(attention_weight),
+                                         -1)  #b,25,256
+
+        attention_weight = F.softmax(attention_weight, dim=-1)  #b,25,256
+
+        pvam_features = torch.matmul(attention_weight, enc_inputs)
+
+        return pvam_features
+
+
+class GSRM(nn.Module):
+
+    def __init__(self,
+                 in_channel,
+                 char_num,
+                 max_len,
+                 num_heads,
+                 hidden_dims,
+                 num_layers,
+                 dropout_rate=0,
+                 attention_dropout=0.1):
+        super(GSRM, self).__init__()
+        self.char_num = char_num
+        self.max_len = max_len
+        self.num_heads = num_heads
+
+        self.cls_op = nn.Linear(in_channel, self.char_num)
+        self.cls_final = nn.Linear(in_channel, self.char_num)
+
+        self.word_emb = Embeddings(d_model=hidden_dims, vocab=char_num)
+        self.pos_emb = nn.Embedding(char_num, hidden_dims)
+        self.dropout_rate = dropout_rate
+        self.emb_drop_out = nn.Dropout(dropout_rate)
+
+        self.forward_self_attn = nn.ModuleList([
+            TransformerBlock(
+                d_model=hidden_dims,
+                nhead=num_heads,
+                attention_dropout_rate=attention_dropout,
+                residual_dropout_rate=0.1,
+                dim_feedforward=hidden_dims,
+                with_self_attn=True,
+                with_cross_attn=False,
+            ) for i in range(num_layers)
+        ])
+
+        self.backward_self_attn = nn.ModuleList([
+            TransformerBlock(
+                d_model=hidden_dims,
+                nhead=num_heads,
+                attention_dropout_rate=attention_dropout,
+                residual_dropout_rate=0.1,
+                dim_feedforward=hidden_dims,
+                with_self_attn=True,
+                with_cross_attn=False,
+            ) for i in range(num_layers)
+        ])
+
+    def _pos_emb(self, word_seq, pos, dropoutrate):
+        """
+        word_Seq: bsz len
+        pos: bsz len
+        """
+        word_emb_seq = self.word_emb(word_seq)
+        pos_emb_seq = self.pos_emb(pos)
+        # pos_emb_seq = pos_emb_seq.detach()
+
+        input_mix = word_emb_seq + pos_emb_seq
+        if dropoutrate > 0:
+            input_mix = self.emb_drop_out(input_mix)
+
+        return input_mix
+
+    def forward(self, inputs):
+
+        bos_idx = self.char_num - 2
+        eos_idx = self.char_num - 1
+        b, t, c = inputs.size()  #b,25,512
+        inputs = inputs.view(-1, c)
+        cls_res = self.cls_op(inputs)  #b,25,n_class
+
+        word_pred_PVAM = F.softmax(cls_res, dim=-1).argmax(-1)
+        word_pred_PVAM = word_pred_PVAM.view(-1, t, 1)
+        #b 25 1
+        word1 = F.pad(word_pred_PVAM, [0, 0, 1, 0], 'constant', value=bos_idx)
+        word_forward = word1[:, :-1, :].squeeze(-1)
+
+        word_backward = word_pred_PVAM.squeeze(-1)
+
+        #mask
+        attn_mask_forward = torch.triu(
+            torch.full((self.max_len, self.max_len),
+                       dtype=torch.float32,
+                       fill_value=-torch.inf),
+            diagonal=1,
+        ).to(inputs.device)
+        attn_mask_forward = attn_mask_forward.unsqueeze(0).expand(
+            self.num_heads, -1, -1)
+        attn_mask_backward = torch.tril(
+            torch.full((self.max_len, self.max_len),
+                       dtype=torch.float32,
+                       fill_value=-torch.inf),
+            diagonal=-1,
+        ).to(inputs.device)
+        attn_mask_backward = attn_mask_backward.unsqueeze(0).expand(
+            self.num_heads, -1, -1)
+
+        #B,25
+
+        pos = torch.arange(self.max_len, dtype=torch.long).to(inputs.device)
+        pos = pos.unsqueeze(0).expand(b, -1)  #b,25
+
+        word_front_mix = self._pos_emb(word_forward, pos, self.dropout_rate)
+        word_backward_mix = self._pos_emb(word_backward, pos,
+                                          self.dropout_rate)
+        # b 25 emb_dim
+
+        for attn_layer in self.forward_self_attn:
+            word_front_mix = attn_layer(word_front_mix,
+                                        self_mask=attn_mask_forward)
+
+        for attn_layer in self.backward_self_attn:
+            word_backward_mix = attn_layer(word_backward_mix,
+                                           self_mask=attn_mask_backward)
+
+        #b,25,emb_dim
+        eos_emd = self.word_emb(torch.full(
+            (1, ), eos_idx).to(inputs.device)).expand(b, 1, -1)
+        word_backward_mix = torch.cat((word_backward_mix, eos_emd), dim=1)
+        word_backward_mix = word_backward_mix[:, 1:, ]
+
+        gsrm_features = word_front_mix + word_backward_mix
+
+        gsrm_out = self.cls_final(gsrm_features)
+        # torch.matmul(gsrm_features,
+        #                         self.word_emb.embedding.weight.permute(1, 0))
+
+        b, t, c = gsrm_out.size()
+        #b,25,n_class
+        gsrm_out = gsrm_out.view(-1, c).contiguous()
+
+        return gsrm_features, cls_res, gsrm_out
+
+
+class VSFD(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super(VSFD, self).__init__()
+        self.char_num = out_channels
+        self.fc0 = nn.Linear(in_channels * 2, in_channels)
+        self.fc1 = nn.Linear(in_channels, self.char_num)
+
+    def forward(self, pvam_feature, gsrm_feature):
+        _, t, c1 = pvam_feature.size()
+        _, t, c2 = gsrm_feature.size()
+        combine_featurs = torch.cat([pvam_feature, gsrm_feature], dim=-1)
+        combine_featurs = combine_featurs.view(-1, c1 + c2).contiguous()
+        atten = self.fc0(combine_featurs)
+        atten = torch.sigmoid(atten)
+        atten = atten.view(-1, t, c1)
+        combine_featurs = atten * pvam_feature + (1 - atten) * gsrm_feature
+        combine_featurs = combine_featurs.view(-1, c1).contiguous()
+        out = self.fc1(combine_featurs)
+        return out
+
+
+class SRNDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 hidden_dims,
+                 num_decoder_layers=4,
+                 max_text_length=25,
+                 num_heads=8,
+                 **kwargs):
+        super(SRNDecoder, self).__init__()
+
+        self.max_text_length = max_text_length
+        self.num_heads = num_heads
+
+        self.pvam = PVAM(in_channels=in_channels,
+                         char_num=out_channels,
+                         max_text_length=max_text_length,
+                         num_heads=num_heads,
+                         hidden_dims=hidden_dims,
+                         dropout_rate=0.1)
+
+        self.gsrm = GSRM(in_channel=in_channels,
+                         char_num=out_channels,
+                         max_len=max_text_length,
+                         num_heads=num_heads,
+                         num_layers=num_decoder_layers,
+                         hidden_dims=hidden_dims)
+
+        self.vsfd = VSFD(in_channels=in_channels, out_channels=out_channels)
+
+    def forward(self, feat, data=None):
+        # feat [B,512,8,32]
+
+        pvam_feature = self.pvam(feat)
+
+        gsrm_features, pvam_preds, gsrm_preds = self.gsrm(pvam_feature)
+
+        vsfd_preds = self.vsfd(pvam_feature, gsrm_features)
+
+        if not self.training:
+            preds = F.softmax(vsfd_preds, dim=-1)
+            return preds
+
+        return [pvam_preds, gsrm_preds, vsfd_preds]

openocr/openrec/modeling/decoders/visionlan_decoder.py

@@ -0,0 +1,321 @@
+import torch
+import torch.nn as nn
+
+from openrec.modeling.decoders.nrtr_decoder import PositionalEncoding, TransformerBlock
+
+
+class Transformer_Encoder(nn.Module):
+
+    def __init__(
+        self,
+        n_layers=3,
+        n_head=8,
+        d_model=512,
+        d_inner=2048,
+        dropout=0.1,
+        n_position=256,
+    ):
+
+        super(Transformer_Encoder, self).__init__()
+        self.pe = PositionalEncoding(dropout=dropout,
+                                     dim=d_model,
+                                     max_len=n_position)
+        self.layer_stack = nn.ModuleList([
+            TransformerBlock(d_model, n_head, d_inner) for _ in range(n_layers)
+        ])
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+
+    def forward(self, enc_output, src_mask):
+        enc_output = self.pe(enc_output)  # position embeding
+        for enc_layer in self.layer_stack:
+            enc_output = enc_layer(enc_output, self_mask=src_mask)
+        enc_output = self.layer_norm(enc_output)
+        return enc_output
+
+
+class PP_layer(nn.Module):
+
+    def __init__(self, n_dim=512, N_max_character=25, n_position=256):
+        super(PP_layer, self).__init__()
+        self.character_len = N_max_character
+        self.f0_embedding = nn.Embedding(N_max_character, n_dim)
+        self.w0 = nn.Linear(N_max_character, n_position)
+        self.wv = nn.Linear(n_dim, n_dim)
+        self.we = nn.Linear(n_dim, N_max_character)
+        self.active = nn.Tanh()
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, enc_output):
+        reading_order = torch.arange(self.character_len,
+                                     dtype=torch.long,
+                                     device=enc_output.device)
+        reading_order = reading_order.unsqueeze(0).expand(
+            enc_output.shape[0], -1)  # (S,) -> (B, S)
+        reading_order = self.f0_embedding(reading_order)  # b,25,512
+        # calculate attention
+
+        t = self.w0(reading_order.transpose(1, 2))  # b,512,256
+        t = self.active(t.transpose(1, 2) + self.wv(enc_output))  # b,256,512
+        t = self.we(t)  # b,256,25
+        t = self.softmax(t.transpose(1, 2))  # b,25,256
+        g_output = torch.bmm(t, enc_output)  # b,25,512
+        return g_output
+
+
+class Prediction(nn.Module):
+
+    def __init__(
+        self,
+        n_dim=512,
+        n_class=37,
+        N_max_character=25,
+        n_position=256,
+    ):
+        super(Prediction, self).__init__()
+        self.pp = PP_layer(n_dim=n_dim,
+                           N_max_character=N_max_character,
+                           n_position=n_position)
+        self.pp_share = PP_layer(n_dim=n_dim,
+                                 N_max_character=N_max_character,
+                                 n_position=n_position)
+        self.w_vrm = nn.Linear(n_dim, n_class)  # output layer
+        self.w_share = nn.Linear(n_dim, n_class)  # output layer
+        self.nclass = n_class
+
+    def forward(self, cnn_feature, f_res, f_sub, is_Train=False, use_mlm=True):
+        if is_Train:
+            if not use_mlm:
+                g_output = self.pp(cnn_feature)  # b,25,512
+                g_output = self.w_vrm(g_output)
+                f_res = 0
+                f_sub = 0
+                return g_output, f_res, f_sub
+            g_output = self.pp(cnn_feature)  # b,25,512
+            f_res = self.pp_share(f_res)
+            f_sub = self.pp_share(f_sub)
+            g_output = self.w_vrm(g_output)
+            f_res = self.w_share(f_res)
+            f_sub = self.w_share(f_sub)
+            return g_output, f_res, f_sub
+        else:
+            g_output = self.pp(cnn_feature)  # b,25,512
+            g_output = self.w_vrm(g_output)
+            return g_output
+
+
+class MLM(nn.Module):
+    """Architecture of MLM."""
+
+    def __init__(
+        self,
+        n_dim=512,
+        n_position=256,
+        n_head=8,
+        dim_feedforward=2048,
+        max_text_length=25,
+    ):
+        super(MLM, self).__init__()
+        self.MLM_SequenceModeling_mask = Transformer_Encoder(
+            n_layers=2,
+            n_head=n_head,
+            d_model=n_dim,
+            d_inner=dim_feedforward,
+            n_position=n_position,
+        )
+        self.MLM_SequenceModeling_WCL = Transformer_Encoder(
+            n_layers=1,
+            n_head=n_head,
+            d_model=n_dim,
+            d_inner=dim_feedforward,
+            n_position=n_position,
+        )
+        self.pos_embedding = nn.Embedding(max_text_length, n_dim)
+        self.w0_linear = nn.Linear(1, n_position)
+        self.wv = nn.Linear(n_dim, n_dim)
+        self.active = nn.Tanh()
+        self.we = nn.Linear(n_dim, 1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, input, label_pos):
+        # transformer unit for generating mask_c
+        feature_v_seq = self.MLM_SequenceModeling_mask(input, src_mask=None)
+        # position embedding layer
+        pos_emb = self.pos_embedding(label_pos.long())
+        pos_emb = self.w0_linear(torch.unsqueeze(pos_emb,
+                                                 dim=2)).transpose(1, 2)
+        # fusion position embedding with features V & generate mask_c
+        att_map_sub = self.active(pos_emb + self.wv(feature_v_seq))
+        att_map_sub = self.we(att_map_sub)  # b,256,1
+        att_map_sub = self.sigmoid(att_map_sub.transpose(1, 2))  # b,1,256
+        # WCL
+        # generate inputs for WCL
+        f_res = input * (1 - att_map_sub.transpose(1, 2)
+                         )  # second path with remaining string
+        f_sub = input * (att_map_sub.transpose(1, 2)
+                         )  # first path with occluded character
+        # transformer units in WCL
+        f_res = self.MLM_SequenceModeling_WCL(f_res, src_mask=None)
+        f_sub = self.MLM_SequenceModeling_WCL(f_sub, src_mask=None)
+        return f_res, f_sub, att_map_sub
+
+
+class MLM_VRM(nn.Module):
+
+    def __init__(
+        self,
+        n_layers=3,
+        n_position=256,
+        n_dim=512,
+        n_head=8,
+        dim_feedforward=2048,
+        max_text_length=25,
+        nclass=37,
+    ):
+        super(MLM_VRM, self).__init__()
+        self.MLM = MLM(
+            n_dim=n_dim,
+            n_position=n_position,
+            n_head=n_head,
+            dim_feedforward=dim_feedforward,
+            max_text_length=max_text_length,
+        )
+        self.SequenceModeling = Transformer_Encoder(
+            n_layers=n_layers,
+            n_head=n_head,
+            d_model=n_dim,
+            d_inner=dim_feedforward,
+            n_position=n_position,
+        )
+        self.Prediction = Prediction(
+            n_dim=n_dim,
+            n_position=n_position,
+            N_max_character=max_text_length + 1,
+            n_class=nclass,
+        )  # N_max_character = 1 eos + 25 characters
+        self.nclass = nclass
+        self.max_text_length = max_text_length
+
+    def forward(self, input, label_pos, training_step, is_Train=False):
+        nT = self.max_text_length
+
+        b, c, h, w = input.shape
+        input = input.reshape(b, c, -1)
+        input = input.transpose(1, 2)
+
+        if is_Train:
+            if training_step == 'LF_1':
+                f_res = 0
+                f_sub = 0
+                input = self.SequenceModeling(input, src_mask=None)
+                text_pre, text_rem, text_mas = self.Prediction(input,
+                                                               f_res,
+                                                               f_sub,
+                                                               is_Train=True,
+                                                               use_mlm=False)
+                return text_pre, text_pre, text_pre
+            elif training_step == 'LF_2':
+                # MLM
+                f_res, f_sub, mask_c = self.MLM(input, label_pos)
+                input = self.SequenceModeling(input, src_mask=None)
+                text_pre, text_rem, text_mas = self.Prediction(input,
+                                                               f_res,
+                                                               f_sub,
+                                                               is_Train=True)
+                return text_pre, text_rem, text_mas
+            elif training_step == 'LA':
+                # MLM
+                f_res, f_sub, mask_c = self.MLM(input, label_pos)
+                # use the mask_c (1 for occluded character and 0 for remaining characters) to occlude input
+                # ratio controls the occluded number in a batch
+                ratio = 2
+                character_mask = torch.zeros_like(mask_c)
+                character_mask[0:b // ratio, :, :] = mask_c[0:b // ratio, :, :]
+                input = input * (1 - character_mask.transpose(1, 2))
+                # VRM
+                # transformer unit for VRM
+                input = self.SequenceModeling(input, src_mask=None)
+                # prediction layer for MLM and VSR
+                text_pre, text_rem, text_mas = self.Prediction(input,
+                                                               f_res,
+                                                               f_sub,
+                                                               is_Train=True)
+                return text_pre, text_rem, text_mas
+        else:  # VRM is only used in the testing stage
+            f_res = 0
+            f_sub = 0
+            contextual_feature = self.SequenceModeling(input, src_mask=None)
+            C = self.Prediction(contextual_feature,
+                                f_res,
+                                f_sub,
+                                is_Train=False,
+                                use_mlm=False)
+            C = C.transpose(1, 0)  # (25, b, 38))
+            out_res = torch.zeros(nT, b, self.nclass).type_as(input.data)
+
+            out_length = torch.zeros(b).type_as(input.data)
+            now_step = 0
+            while 0 in out_length and now_step < nT:
+                tmp_result = C[now_step, :, :]
+                out_res[now_step] = tmp_result
+                tmp_result = tmp_result.topk(1)[1].squeeze(dim=1)
+                for j in range(b):
+                    if out_length[j] == 0 and tmp_result[j] == 0:
+                        out_length[j] = now_step + 1
+                now_step += 1
+            for j in range(0, b):
+                if int(out_length[j]) == 0:
+                    out_length[j] = nT
+            start = 0
+            output = torch.zeros(int(out_length.sum()),
+                                 self.nclass).type_as(input.data)
+            for i in range(0, b):
+                cur_length = int(out_length[i])
+                output[start:start + cur_length] = out_res[0:cur_length, i, :]
+                start += cur_length
+
+            return output, out_length
+
+
+class VisionLANDecoder(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        n_head=None,
+        training_step='LA',
+        n_layers=3,
+        n_position=256,
+        max_text_length=25,
+    ):
+        super(VisionLANDecoder, self).__init__()
+        self.training_step = training_step
+        n_dim = in_channels
+        dim_feedforward = n_dim * 4
+        n_head = n_head if n_head is not None else n_dim // 32
+
+        self.MLM_VRM = MLM_VRM(
+            n_layers=n_layers,
+            n_position=n_position,
+            n_dim=n_dim,
+            n_head=n_head,
+            dim_feedforward=dim_feedforward,
+            max_text_length=max_text_length,
+            nclass=out_channels + 1,
+        )
+
+    def forward(self, x, data=None):
+        # MLM + VRM
+        if self.training:
+            label_pos = data[-2]
+            text_pre, text_rem, text_mas = self.MLM_VRM(x,
+                                                        label_pos,
+                                                        self.training_step,
+                                                        is_Train=True)
+            return text_pre, text_rem, text_mas
+        else:
+            output, out_length = self.MLM_VRM(x,
+                                              None,
+                                              self.training_step,
+                                              is_Train=False)
+            return output, out_length

openocr/openrec/modeling/encoders/__init__.py

@@ -0,0 +1,39 @@
+__all__ = ['build_encoder']
+
+
+def build_encoder(config):
+    # from .rec_mobilenet_v3 import MobileNetV3
+    from .focalsvtr import FocalSVTR
+    from .rec_hgnet import PPHGNet_small
+    from .rec_lcnetv3 import PPLCNetV3
+    from .rec_mv1_enhance import MobileNetV1Enhance
+    from .rec_nrtr_mtb import MTB
+    from .rec_resnet_31 import ResNet31
+    from .rec_resnet_45 import ResNet45
+    from .rec_resnet_fpn import ResNet_FPN
+    from .rec_resnet_vd import ResNet
+    from .resnet31_rnn import ResNet_ASTER
+    from .svtrnet import SVTRNet
+    from .svtrnet2dpos import SVTRNet2DPos
+    from .svtrv2 import SVTRv2
+    from .svtrv2_lnconv import SVTRv2LNConv
+    from .svtrv2_lnconv_two33 import SVTRv2LNConvTwo33
+    from .vit import ViT
+    from .cam_encoder import CAMEncoder
+    from .convnextv2 import ConvNeXtV2
+    from .autostr_encoder import AutoSTREncoder
+    from .nrtr_encoder import NRTREncoder
+    from .repvit import RepSVTREncoder
+    support_dict = [
+        'MobileNetV1Enhance', 'ResNet31', 'MobileNetV3', 'PPLCNetV3',
+        'PPHGNet_small', 'ResNet', 'MTB', 'SVTRNet', 'ResNet45', 'ViT',
+        'SVTRNet2DPos', 'SVTRv2', 'FocalSVTR', 'ResNet_FPN', 'ResNet_ASTER',
+        'SVTRv2LNConv', 'SVTRv2LNConvTwo33', 'CAMEncoder', 'ConvNeXtV2',
+        'AutoSTREncoder', 'NRTREncoder', 'RepSVTREncoder'
+    ]
+
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception(
+        'when encoder of rec model only support {}'.format(support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class