openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/decoders/aster_decoder.py

@@ -0,0 +1,170 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.nn import init
+
+
+class Embedding(nn.Module):
+
+    def __init__(self, in_timestep, in_planes, mid_dim=4096, embed_dim=300):
+        super(Embedding, self).__init__()
+        self.in_timestep = in_timestep
+        self.in_planes = in_planes
+        self.embed_dim = embed_dim
+        self.mid_dim = mid_dim
+        self.eEmbed = nn.Linear(
+            in_timestep * in_planes,
+            self.embed_dim)  # Embed encoder output to a word-embedding like
+
+    def forward(self, x):
+        x = x.flatten(1)
+        x = self.eEmbed(x)
+        return x
+
+
+class Attn_Rnn_Block(nn.Module):
+
+    def __init__(self, featdim, hiddendim, embedding_dim, out_channels,
+                 attndim):
+        super(Attn_Rnn_Block, self).__init__()
+
+        self.attndim = attndim
+        self.embedding_dim = embedding_dim
+        self.feat_embed = nn.Linear(featdim, attndim)
+        self.hidden_embed = nn.Linear(hiddendim, attndim)
+        self.attnfeat_embed = nn.Linear(attndim, 1)
+        self.gru = nn.GRU(input_size=featdim + self.embedding_dim,
+                          hidden_size=hiddendim,
+                          batch_first=True)
+        self.fc = nn.Linear(hiddendim, out_channels)
+        self.init_weights()
+
+    def init_weights(self):
+        init.normal_(self.hidden_embed.weight, std=0.01)
+        init.constant_(self.hidden_embed.bias, 0)
+        init.normal_(self.attnfeat_embed.weight, std=0.01)
+        init.constant_(self.attnfeat_embed.bias, 0)
+
+    def _attn(self, feat, h_state):
+        b, t, _ = feat.shape
+        feat = self.feat_embed(feat)
+        h_state = self.hidden_embed(h_state.squeeze(0)).unsqueeze(1)
+        h_state = h_state.expand(b, t, self.attndim)
+        sumTanh = torch.tanh(feat + h_state)
+        attn_w = self.attnfeat_embed(sumTanh).squeeze(-1)
+        attn_w = F.softmax(attn_w, dim=1).unsqueeze(1)
+        # [B,1,25]
+        return attn_w
+
+    def forward(self, feat, h_state, label_input):
+
+        attn_w = self._attn(feat, h_state)
+
+        attn_feat = attn_w @ feat
+        attn_feat = attn_feat.squeeze(1)
+
+        output, h_state = self.gru(
+            torch.cat([label_input, attn_feat], 1).unsqueeze(1), h_state)
+        pred = self.fc(output)
+
+        return pred, h_state
+
+
+class ASTERDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 embedding_dim=256,
+                 hiddendim=256,
+                 attndim=256,
+                 max_len=25,
+                 seed=False,
+                 time_step=32,
+                 **kwargs):
+        super(ASTERDecoder, self).__init__()
+        self.num_classes = out_channels
+        self.bos = out_channels - 2
+        self.eos = 0
+        self.padding_idx = out_channels - 1
+        self.seed = seed
+        if seed:
+            self.embeder = Embedding(
+                in_timestep=time_step,
+                in_planes=in_channels,
+            )
+        self.word_embedding = nn.Embedding(self.num_classes,
+                                           embedding_dim,
+                                           padding_idx=self.padding_idx)
+
+        self.attndim = attndim
+        self.hiddendim = hiddendim
+        self.max_seq_len = max_len + 1
+
+        self.featdim = in_channels
+
+        self.attn_rnn_block = Attn_Rnn_Block(
+            featdim=self.featdim,
+            hiddendim=hiddendim,
+            embedding_dim=embedding_dim,
+            out_channels=out_channels - 2,
+            attndim=attndim,
+        )
+        self.embed_fc = nn.Linear(300, self.hiddendim)
+
+    def get_initial_state(self, embed, tile_times=1):
+        assert embed.shape[1] == 300
+        state = self.embed_fc(embed)  # N * sDim
+        if tile_times != 1:
+            state = state.unsqueeze(1)
+            trans_state = state.transpose(0, 1)
+            state = trans_state.tile([tile_times, 1, 1])
+            trans_state = state.transpose(0, 1)
+            state = trans_state.reshape(-1, self.hiddendim)
+        state = state.unsqueeze(0)  # 1 * N * sDim
+        return state
+
+    def forward(self, feat, data=None):
+        # b,25,512
+        b = feat.size(0)
+        if self.seed:
+            embedding_vectors = self.embeder(feat)
+            h_state = self.get_initial_state(embedding_vectors)
+        else:
+            h_state = torch.zeros(1, b, self.hiddendim).to(feat.device)
+        outputs = []
+        if self.training:
+            label = data[0]
+            label_embedding = self.word_embedding(label)  # [B,25,256]
+            tokens = label_embedding[:, 0, :]
+            max_len = data[1].max() + 1
+        else:
+            tokens = torch.full([b, 1],
+                                self.bos,
+                                device=feat.device,
+                                dtype=torch.long)
+            tokens = self.word_embedding(tokens.squeeze(1))
+            max_len = self.max_seq_len
+        pred, h_state = self.attn_rnn_block(feat, h_state, tokens)
+        outputs.append(pred)
+
+        dec_seq = torch.full((feat.shape[0], max_len),
+                             self.padding_idx,
+                             dtype=torch.int64,
+                             device=feat.get_device())
+        dec_seq[:, :1] = torch.argmax(pred, dim=-1)
+        for i in range(1, max_len):
+            if not self.training:
+                max_idx = torch.argmax(pred, dim=-1).squeeze(1)
+                tokens = self.word_embedding(max_idx)
+                dec_seq[:, i] = max_idx
+                if (dec_seq == self.eos).any(dim=-1).all():
+                    break
+            else:
+                tokens = label_embedding[:, i, :]
+            pred, h_state = self.attn_rnn_block(feat, h_state, tokens)
+            outputs.append(pred)
+        preds = torch.cat(outputs, 1)
+        if self.seed and self.training:
+            return [embedding_vectors, preds]
+        return preds if self.training else F.softmax(preds, -1)

openocr/openrec/modeling/decoders/bus_decoder.py

@@ -0,0 +1,133 @@
+"""This code is refer from:
+https://github.com/jjwei66/BUSNet
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .nrtr_decoder import PositionalEncoding, TransformerBlock
+from .abinet_decoder import _get_mask, _get_length
+
+
+class BUSDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 nhead=8,
+                 num_layers=4,
+                 dim_feedforward=2048,
+                 dropout=0.1,
+                 max_length=25,
+                 ignore_index=100,
+                 pretraining=False,
+                 detach=True):
+        super().__init__()
+        d_model = in_channels
+        self.ignore_index = ignore_index
+        self.pretraining = pretraining
+        self.d_model = d_model
+        self.detach = detach
+        self.max_length = max_length + 1  # additional stop token
+        self.out_channels = out_channels
+        # --------------------------------------------------------------------------
+        # decoder specifics
+        self.proj = nn.Linear(out_channels, d_model, False)
+        self.token_encoder = PositionalEncoding(dropout=0.1,
+                                                dim=d_model,
+                                                max_len=self.max_length)
+        self.pos_encoder = PositionalEncoding(dropout=0.1,
+                                              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)
+        ])
+
+        v_mask = torch.empty((1, 1, d_model))
+        l_mask = torch.empty((1, 1, d_model))
+        self.v_mask = nn.Parameter(v_mask)
+        self.l_mask = nn.Parameter(l_mask)
+        torch.nn.init.uniform_(self.v_mask, -0.001, 0.001)
+        torch.nn.init.uniform_(self.l_mask, -0.001, 0.001)
+
+        v_embeding = torch.empty((1, 1, d_model))
+        l_embeding = torch.empty((1, 1, d_model))
+        self.v_embeding = nn.Parameter(v_embeding)
+        self.l_embeding = nn.Parameter(l_embeding)
+        torch.nn.init.uniform_(self.v_embeding, -0.001, 0.001)
+        torch.nn.init.uniform_(self.l_embeding, -0.001, 0.001)
+        self.cls = nn.Linear(d_model, out_channels)
+
+    def forward_decoder(self, q, x, mask=None):
+        for decoder_layer in self.decoder:
+            q = decoder_layer(q, x, cross_mask=mask)
+        output = q  # (N, T, E)
+        logits = self.cls(output)  # (N, T, C)
+        return logits
+
+    def forward(self, img_feat, data=None):
+        """
+        Args:
+            tokens: (N, T, C) where T is length, N is batch size and C is classes number
+            lengths: (N,)
+        """
+        img_feat = img_feat + self.v_embeding
+        B, L, C = img_feat.shape
+
+        # --------------------------------------------------------------------------
+        # decoder procedure
+        T = self.max_length
+        zeros = img_feat.new_zeros((B, T, C))
+        zeros_len = img_feat.new_zeros(B)
+        query = self.pos_encoder(zeros)
+
+        # 1. vision decode
+        v_embed = torch.cat((img_feat, self.l_mask.repeat(B, T, 1)),
+                            dim=1)  # v
+        padding_mask = _get_mask(
+            self.max_length + zeros_len,
+            self.max_length)  # 对tokens长度以外的padding # B, maxlen maxlen
+        v_mask = torch.zeros((1, 1, self.max_length, L),
+                             device=img_feat.device).tile([B, 1, 1,
+                                                           1])  # maxlen L
+        mask = torch.cat((v_mask, padding_mask), 3)
+        v_logits = self.forward_decoder(query, v_embed, mask=mask)
+
+        # 2. language decode
+        if self.training and self.pretraining:
+            tgt = torch.where(data[0] == self.ignore_index, 0, data[0])
+            tokens = F.one_hot(tgt, num_classes=self.out_channels)
+            tokens = tokens.float()
+            lengths = data[-1]
+        else:
+            tokens = torch.softmax(v_logits, dim=-1)
+            lengths = _get_length(v_logits)
+            tokens = tokens.detach()
+        token_embed = self.proj(tokens)  # (N, T, E)
+        token_embed = self.token_encoder(token_embed)  # (T, N, E)
+        token_embed = token_embed + self.l_embeding
+
+        padding_mask = _get_mask(lengths,
+                                 self.max_length)  # 对tokens长度以外的padding
+        mask = torch.cat((v_mask, padding_mask), 3)
+        l_embed = torch.cat((self.v_mask.repeat(B, L, 1), token_embed), dim=1)
+        l_logits = self.forward_decoder(query, l_embed, mask=mask)
+
+        # 3. vision language decode
+        vl_embed = torch.cat((img_feat, token_embed), dim=1)
+        vl_logits = self.forward_decoder(query, vl_embed, mask=mask)
+
+        if self.training:
+            return {'align': [vl_logits], 'lang': l_logits, 'vision': v_logits}
+        else:
+            return F.softmax(vl_logits, -1)

openocr/openrec/modeling/decoders/cam_decoder.py

@@ -0,0 +1,43 @@
+import torch.nn as nn
+
+from .nrtr_decoder import NRTRDecoder
+
+
+class CAMDecoder(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        nhead=None,
+        num_encoder_layers=6,
+        beam_size=0,
+        num_decoder_layers=6,
+        max_len=25,
+        attention_dropout_rate=0.0,
+        residual_dropout_rate=0.1,
+        scale_embedding=True,
+    ):
+        super().__init__()
+
+        self.decoder = NRTRDecoder(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            nhead=nhead,
+            num_encoder_layers=num_encoder_layers,
+            beam_size=beam_size,
+            num_decoder_layers=num_decoder_layers,
+            max_len=max_len,
+            attention_dropout_rate=attention_dropout_rate,
+            residual_dropout_rate=residual_dropout_rate,
+            scale_embedding=scale_embedding,
+        )
+
+    def forward(self, x, data=None):
+        dec_in = x['refined_feat']
+        dec_output = self.decoder(dec_in, data=data)
+        x['rec_output'] = dec_output
+        if self.training:
+            return x
+        else:
+            return dec_output

openocr/openrec/modeling/decoders/cdistnet_decoder.py

@@ -0,0 +1,334 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from openrec.modeling.decoders.nrtr_decoder import Embeddings, PositionalEncoding, TransformerBlock  # , Beam
+from openrec.modeling.decoders.visionlan_decoder import Transformer_Encoder
+
+
+def generate_square_subsequent_mask(sz):
+    r"""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.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+    mask = (mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(
+        mask == 1, float(0.0)))
+    return mask
+
+
+class SEM_Pre(nn.Module):
+
+    def __init__(
+        self,
+        d_model=512,
+        dst_vocab_size=40,
+        residual_dropout_rate=0.1,
+    ):
+        super(SEM_Pre, self).__init__()
+        self.embedding = Embeddings(d_model=d_model, vocab=dst_vocab_size)
+
+        self.positional_encoding = PositionalEncoding(
+            dropout=residual_dropout_rate,
+            dim=d_model,
+        )
+
+    def forward(self, tgt):
+        tgt = self.embedding(tgt)
+        tgt = self.positional_encoding(tgt)
+        tgt_mask = generate_square_subsequent_mask(tgt.shape[1]).to(tgt.device)
+        return tgt, tgt_mask
+
+
+class POS_Pre(nn.Module):
+
+    def __init__(
+        self,
+        d_model=512,
+    ):
+        super(POS_Pre, self).__init__()
+        self.pos_encoding = PositionalEncoding(
+            dropout=0.1,
+            dim=d_model,
+        )
+        self.linear1 = nn.Linear(d_model, d_model)
+        self.linear2 = nn.Linear(d_model, d_model)
+
+        self.norm2 = nn.LayerNorm(d_model)
+
+    def forward(self, tgt):
+        pos = tgt.new_zeros(*tgt.shape)
+        pos = self.pos_encoding(pos)
+
+        pos2 = self.linear2(F.relu(self.linear1(pos)))
+        pos = self.norm2(pos + pos2)
+        return pos
+
+
+class DSF(nn.Module):
+
+    def __init__(self, d_model, fusion_num):
+        super(DSF, self).__init__()
+        self.w_att = nn.Linear(fusion_num * d_model, d_model)
+
+    def forward(self, l_feature, v_feature):
+        """
+        Args:
+            l_feature: (N, T, E) where T is length, N is batch size and d is dim of model
+            v_feature: (N, T, E) shape the same as l_feature
+            l_lengths: (N,)
+            v_lengths: (N,)
+        """
+        f = torch.cat((l_feature, v_feature), dim=2)
+        f_att = torch.sigmoid(self.w_att(f))
+        output = f_att * v_feature + (1 - f_att) * l_feature
+
+        return output
+
+
+class MDCDP(nn.Module):
+    r"""
+    Multi-Domain CharacterDistance Perception
+    """
+
+    def __init__(self, d_model, n_head, d_inner, num_layers):
+        super(MDCDP, self).__init__()
+
+        self.num_layers = num_layers
+
+        # step 1 SAE
+        self.layers_pos = nn.ModuleList([
+            TransformerBlock(d_model, n_head, d_inner)
+            for _ in range(num_layers)
+        ])
+
+        # step 2 CBI:
+        self.layers2 = nn.ModuleList([
+            TransformerBlock(
+                d_model,
+                n_head,
+                d_inner,
+                with_self_attn=False,
+                with_cross_attn=True,
+            ) for _ in range(num_layers)
+        ])
+        self.layers3 = nn.ModuleList([
+            TransformerBlock(
+                d_model,
+                n_head,
+                d_inner,
+                with_self_attn=False,
+                with_cross_attn=True,
+            ) for _ in range(num_layers)
+        ])
+
+        # step 3 :DSF
+        self.dynamic_shared_fusion = DSF(d_model, 2)
+
+    def forward(
+        self,
+        sem,
+        vis,
+        pos,
+        tgt_mask=None,
+        memory_mask=None,
+    ):
+
+        for i in range(self.num_layers):
+            # ----------step 1 -----------: SAE: Self-Attention Enhancement
+            pos = self.layers_pos[i](pos, self_mask=tgt_mask)
+
+            # ----------step 2 -----------: CBI: Cross-Branch Interaction
+
+            # CBI-V
+            pos_vis = self.layers2[i](
+                pos,
+                vis,
+                cross_mask=memory_mask,
+            )
+
+            # CBI-S
+            pos_sem = self.layers3[i](
+                pos,
+                sem,
+                cross_mask=tgt_mask,
+            )
+
+            # ----------step 3 -----------: DSF: Dynamic Shared Fusion
+            pos = self.dynamic_shared_fusion(pos_vis, pos_sem)
+
+        output = pos
+        return output
+
+
+class ConvBnRelu(nn.Module):
+    # adapt padding for kernel_size change
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        conv=nn.Conv2d,
+        stride=2,
+        inplace=True,
+    ):
+        super().__init__()
+        p_size = [int(k // 2) for k in kernel_size]
+        # p_size = int(kernel_size//2)
+        self.conv = conv(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=p_size,
+        )
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=inplace)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class CDistNetDecoder(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 n_head=None,
+                 num_encoder_blocks=3,
+                 num_decoder_blocks=3,
+                 beam_size=0,
+                 max_len=25,
+                 residual_dropout_rate=0.1,
+                 add_conv=False,
+                 **kwargs):
+        super(CDistNetDecoder, self).__init__()
+        dst_vocab_size = out_channels
+        self.ignore_index = dst_vocab_size - 1
+        self.bos = dst_vocab_size - 2
+        self.eos = 0
+        self.beam_size = beam_size
+        self.max_len = max_len
+        self.add_conv = add_conv
+        d_model = in_channels
+        dim_feedforward = d_model * 4
+        n_head = n_head if n_head is not None else d_model // 32
+
+        if add_conv:
+            self.convbnrelu = ConvBnRelu(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                kernel_size=(1, 3),
+                stride=(1, 2),
+            )
+        if num_encoder_blocks > 0:
+            self.positional_encoding = PositionalEncoding(
+                dropout=0.1,
+                dim=d_model,
+            )
+            self.trans_encoder = Transformer_Encoder(
+                n_layers=num_encoder_blocks,
+                n_head=n_head,
+                d_model=d_model,
+                d_inner=dim_feedforward,
+            )
+        else:
+            self.trans_encoder = None
+        self.semantic_branch = SEM_Pre(
+            d_model=d_model,
+            dst_vocab_size=dst_vocab_size,
+            residual_dropout_rate=residual_dropout_rate,
+        )
+        self.positional_branch = POS_Pre(d_model=d_model)
+
+        self.mdcdp = MDCDP(d_model, n_head, dim_feedforward // 2,
+                           num_decoder_blocks)
+        self._reset_parameters()
+
+        self.tgt_word_prj = nn.Linear(
+            d_model, dst_vocab_size - 2,
+            bias=False)  # We don't predict <bos> nor <pad>
+        self.tgt_word_prj.weight.data.normal_(mean=0.0, std=d_model**-0.5)
+
+    def forward(self, x, data=None):
+        if self.add_conv:
+            x = self.convbnrelu(x)
+            # x = rearrange(x, "b c h w -> b (w h) c")
+        x = x.flatten(2).transpose(1, 2)
+        if self.trans_encoder is not None:
+            x = self.positional_encoding(x)
+            vis_feat = self.trans_encoder(x, src_mask=None)
+        else:
+            vis_feat = x
+        if self.training:
+            max_len = data[1].max()
+            tgt = data[0][:, :1 + max_len]
+            res = self.forward_train(vis_feat, tgt)
+        else:
+            if self.beam_size > 0:
+                res = self.forward_beam(vis_feat)
+            else:
+                res = self.forward_test(vis_feat)
+        return res
+
+    def forward_train(self, vis_feat, tgt):
+        sem_feat, sem_mask = self.semantic_branch(tgt)
+        pos_feat = self.positional_branch(sem_feat)
+        output = self.mdcdp(
+            sem_feat,
+            vis_feat,
+            pos_feat,
+            tgt_mask=sem_mask,
+            memory_mask=None,
+        )
+
+        logit = self.tgt_word_prj(output)
+        return logit
+
+    def forward_test(self, vis_feat):
+        bs = vis_feat.size(0)
+
+        dec_seq = torch.full(
+            (bs, self.max_len + 1),
+            self.ignore_index,
+            dtype=torch.int64,
+            device=vis_feat.device,
+        )
+        dec_seq[:, 0] = self.bos
+        logits = []
+        for len_dec_seq in range(0, self.max_len):
+            sem_feat, sem_mask = self.semantic_branch(dec_seq[:, :len_dec_seq +
+                                                              1])
+            pos_feat = self.positional_branch(sem_feat)
+            output = self.mdcdp(
+                sem_feat,
+                vis_feat,
+                pos_feat,
+                tgt_mask=sem_mask,
+                memory_mask=None,
+            )
+
+            dec_output = output[:, -1:, :]
+
+            word_prob = F.softmax(self.tgt_word_prj(dec_output), dim=-1)
+            logits.append(word_prob)
+            if len_dec_seq < self.max_len:
+                # greedy decode. add the next token index to the target input
+                dec_seq[:, len_dec_seq + 1] = word_prob.squeeze(1).argmax(-1)
+                # Efficient batch decoding: If all output words have at least one EOS token, end decoding.
+                if (dec_seq == self.eos).any(dim=-1).all():
+                    break
+        logits = torch.cat(logits, dim=1)
+        return logits
+
+    def forward_beam(self, x):
+        """Translation work in one batch."""
+        # to do
+
+    def _reset_parameters(self):
+        r"""Initiate parameters in the transformer model."""
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)