openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/modeling/decoders/dan_decoder.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CAM(nn.Module):
+    '''
+    Convolutional Alignment Module
+    '''
+
+    # Current version only supports input whose size is a power of 2, such as 32, 64, 128 etc.
+    # You can adapt it to any input size by changing the padding or stride.
+    def __init__(self,
+                 channels_list=[64, 128, 256, 512],
+                 strides_list=[[2, 2], [1, 1], [1, 1]],
+                 in_shape=[8, 32],
+                 maxT=25,
+                 depth=4,
+                 num_channels=128):
+        super(CAM, self).__init__()
+        # cascade multiscale features
+        fpn = []
+        for i in range(1, len(channels_list)):
+            fpn.append(
+                nn.Sequential(
+                    nn.Conv2d(channels_list[i - 1], channels_list[i], (3, 3),
+                              (strides_list[i - 1][0], strides_list[i - 1][1]),
+                              1), nn.BatchNorm2d(channels_list[i]),
+                    nn.ReLU(True)))
+        self.fpn = nn.Sequential(*fpn)
+        # convolutional alignment
+        # convs
+        assert depth % 2 == 0, 'the depth of CAM must be a even number.'
+        # in_shape = scales[-1]
+        strides = []
+        conv_ksizes = []
+        deconv_ksizes = []
+        h, w = in_shape[0], in_shape[1]
+        for i in range(0, int(depth / 2)):
+            stride = [2] if 2**(depth / 2 - i) <= h else [1]
+            stride = stride + [2] if 2**(depth / 2 - i) <= w else stride + [1]
+            strides.append(stride)
+            conv_ksizes.append([3, 3])
+            deconv_ksizes.append([_**2 for _ in stride])
+        convs = [
+            nn.Sequential(
+                nn.Conv2d(channels_list[-1], num_channels,
+                          tuple(conv_ksizes[0]), tuple(strides[0]),
+                          (int((conv_ksizes[0][0] - 1) / 2),
+                           int((conv_ksizes[0][1] - 1) / 2))),
+                nn.BatchNorm2d(num_channels), nn.ReLU(True))
+        ]
+        for i in range(1, int(depth / 2)):
+            convs.append(
+                nn.Sequential(
+                    nn.Conv2d(num_channels, num_channels,
+                              tuple(conv_ksizes[i]), tuple(strides[i]),
+                              (int((conv_ksizes[i][0] - 1) / 2),
+                               int((conv_ksizes[i][1] - 1) / 2))),
+                    nn.BatchNorm2d(num_channels), nn.ReLU(True)))
+        self.convs = nn.Sequential(*convs)
+        # deconvs
+        deconvs = []
+        for i in range(1, int(depth / 2)):
+            deconvs.append(
+                nn.Sequential(
+                    nn.ConvTranspose2d(
+                        num_channels, num_channels,
+                        tuple(deconv_ksizes[int(depth / 2) - i]),
+                        tuple(strides[int(depth / 2) - i]),
+                        (int(deconv_ksizes[int(depth / 2) - i][0] / 4.),
+                         int(deconv_ksizes[int(depth / 2) - i][1] / 4.))),
+                    nn.BatchNorm2d(num_channels), nn.ReLU(True)))
+        deconvs.append(
+            nn.Sequential(
+                nn.ConvTranspose2d(num_channels, maxT, tuple(deconv_ksizes[0]),
+                                   tuple(strides[0]),
+                                   (int(deconv_ksizes[0][0] / 4.),
+                                    int(deconv_ksizes[0][1] / 4.))),
+                nn.Sigmoid()))
+        self.deconvs = nn.Sequential(*deconvs)
+
+    def forward(self, input):
+        x = input[0]
+        for i in range(0, len(self.fpn)):
+            # print(self.fpn[i](x).shape, input[i+1].shape)
+            x = self.fpn[i](x) + input[i + 1]
+        conv_feats = []
+        for i in range(0, len(self.convs)):
+            x = self.convs[i](x)
+            conv_feats.append(x)
+        for i in range(0, len(self.deconvs) - 1):
+            x = self.deconvs[i](x)
+            x = x + conv_feats[len(conv_feats) - 2 - i]
+        x = self.deconvs[-1](x)
+        return x
+
+
+class CAMSimp(nn.Module):
+
+    def __init__(self, maxT=25, num_channels=128):
+        super(CAMSimp, self).__init__()
+        self.conv = nn.Sequential(nn.Conv2d(num_channels, maxT, 1, 1, 0),
+                                  nn.Sigmoid())
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+class DANDecoder(nn.Module):
+    '''
+    Decoupled Text Decoder
+    '''
+
+    def __init__(self,
+                 out_channels,
+                 in_channels,
+                 use_cam=True,
+                 max_len=25,
+                 channels_list=[64, 128, 256, 512],
+                 strides_list=[[2, 2], [1, 1], [1, 1]],
+                 in_shape=[8, 32],
+                 depth=4,
+                 dropout=0.3,
+                 **kwargs):
+        super(DANDecoder, self).__init__()
+        self.eos = 0
+        self.bos = out_channels - 2
+        self.ignore_index = out_channels - 1
+        nchannel = in_channels
+        self.nchannel = in_channels
+        self.use_cam = use_cam
+        if use_cam:
+            self.cam = CAM(channels_list=channels_list,
+                           strides_list=strides_list,
+                           in_shape=in_shape,
+                           maxT=max_len + 1,
+                           depth=depth,
+                           num_channels=nchannel)
+        else:
+            self.cam = CAMSimp(maxT=max_len + 1, num_channels=nchannel)
+        self.pre_lstm = nn.LSTM(nchannel,
+                                int(nchannel / 2),
+                                bidirectional=True)
+        self.rnn = nn.GRUCell(nchannel * 2, nchannel)
+        self.generator = nn.Sequential(nn.Dropout(p=dropout),
+                                       nn.Linear(nchannel, out_channels - 2))
+        self.char_embeddings = nn.Embedding(out_channels,
+                                            embedding_dim=in_channels,
+                                            padding_idx=out_channels - 1)
+
+    def forward(self, inputs, data=None):
+        A = self.cam(inputs)
+        if isinstance(inputs, list):
+            feature = inputs[-1]
+        else:
+            feature = inputs
+        nB, nC, nH, nW = feature.shape
+        nT = A.shape[1]
+        # Normalize
+        A = A / A.view(nB, nT, -1).sum(2).view(nB, nT, 1, 1)
+        # weighted sum
+        C = feature.view(nB, 1, nC, nH, nW) * A.view(nB, nT, 1, nH, nW)
+        C = C.view(nB, nT, nC, -1).sum(3).transpose(1, 0)  # T, B, C
+        C, _ = self.pre_lstm(C)  # T, B, C
+        C = F.dropout(C, p=0.3, training=self.training)
+        if self.training:
+            text = data[0]
+            text_length = data[-1]
+            nsteps = int(text_length.max())
+            gru_res = torch.zeros_like(C)
+            hidden = torch.zeros(nB, self.nchannel).type_as(C.data)
+            prev_emb = self.char_embeddings(text[:, 0])
+            for i in range(0, nsteps + 1):
+                hidden = self.rnn(torch.cat((C[i, :, :], prev_emb), dim=1),
+                                  hidden)
+                gru_res[i, :, :] = hidden
+                prev_emb = self.char_embeddings(text[:, i + 1])
+            gru_res = self.generator(gru_res)
+            return gru_res[:nsteps + 1, :, :].transpose(1, 0)
+        else:
+            gru_res = torch.zeros_like(C)
+            hidden = torch.zeros(nB, self.nchannel).type_as(C.data)
+            prev_emb = self.char_embeddings(
+                torch.zeros(nB, dtype=torch.int64, device=feature.device) +
+                self.bos)
+            dec_seq = torch.full((nB, nT),
+                                self.ignore_index,
+                                dtype=torch.int64,
+                                device=feature.get_device())
+            
+            for i in range(0, nT):
+                hidden = self.rnn(torch.cat((C[i, :, :], prev_emb), dim=1),
+                                  hidden)
+                gru_res[i, :, :] = hidden
+                mid_res = self.generator(hidden).argmax(-1)
+                dec_seq[:, i] = mid_res.squeeze(0)
+                if (dec_seq == self.eos).any(dim=-1).all():
+                    break
+                prev_emb = self.char_embeddings(mid_res)
+            gru_res = self.generator(gru_res)
+            return F.softmax(gru_res.transpose(1, 0), -1)