Upload 28 files

train_old.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eed45bcd25593dca2576c20721a57a449a6b557faf40336bcd7690b2a82eb2e1
+size 89572737

vgg-seq2seq.yaml

@@ -0,0 +1,90 @@
+project: vietocr_new
+name: Train
+
+device: cuda:0
+
+# change to list chars of your dataset or use default vietnamese chars
+vocab: 'aAàÀảẢãÃáÁạẠăĂằẰẳẲẵẴắẮặẶâÂầẦẩẨẫẪấẤậẬbBcCdDđĐeEèÈẻẺẽẼéÉẹẸêÊềỀểỂễỄếẾệỆfFgGhHiIìÌỉỈĩĨíÍịỊjJkKlLmMnNoOòÒỏỎõÕóÓọỌôÔồỒổỔỗỖốỐộỘơƠờỜởỞỡỠớỚợỢpPqQrRsStTuUùÙủỦũŨúÚụỤưƯừỪửỬữỮứỨựỰvVwWxXyYỳỲỷỶỹỸýÝỵỴzZ0123456789!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
+
+seq_modeling: seq2seq
+transformer:  
+  encoder_hidden: 256
+  decoder_hidden: 256
+  img_channel: 256
+  decoder_embedded: 256
+  dropout: 0.1
+
+optimizer:
+  max_lr: 0.001 
+  pct_start: 0.1
+
+trainer:
+  batch_size: 128
+  print_every: 100
+  valid_every: 500
+  test_every: 500
+  iters: 10000
+  # where to save our model for prediction
+  export: weights/train_model.pth
+  checkpoint: ./checkpoint/checkpoint_model.pth
+  log: ./train.log
+  # null to disable compuate accuracy, or change to number of sample to enable validiation while training
+  metrics: 49228
+  test_metrics: 28918
+  pretrained: false
+
+dataset:    
+  # path to image
+  data_root: /mnt/disk3/CGGANv2
+  # path to annotation
+  train_annotation: datasets/labels/train.txt
+  valid_annotation: datasets/labels/valid.txt
+  test_annotation: datasets/labels/test.txt
+  # path to lmdb datasets
+  train_lmdb: datasets/lmdb/train
+  valid_lmdb: datasets/lmdb/valid
+  test_lmdb: datasets/lmdb/test
+
+  # resize image to 32 height, larger height will increase accuracy
+  image_height: 32
+  image_min_width: 32
+  image_max_width: 512
+
+dataloader:
+  num_workers: 12
+  pin_memory: true
+
+aug:
+  image_aug: false
+  masked_language_model: false
+
+predictor:
+  # disable or enable beamsearch while prediction, use beamsearch will be slower
+  beamsearch: false
+
+quiet: false
+
+# for train
+pretrain: https://vocr.vn/data/vietocr/vgg_seq2seq.pth
+
+# url or local path (for predict)
+weights: https://vocr.vn/data/vietocr/vgg_seq2seq.pth
+
+backbone: vgg19_bn
+cnn:
+  # pooling stride size
+  ss:
+    - [2, 2]
+    - [2, 2]
+    - [2, 1]
+    - [2, 1]
+    - [1, 1]         
+  # pooling kernel size 
+  ks:
+    - [2, 2]
+    - [2, 2]
+    - [2, 1]
+    - [2, 1]
+    - [1, 1]
+  # dim of ouput feature map
+  hidden: 256

vietocr/model/backbone/cnn.py

@@ -0,0 +1,28 @@
+import torch
+from torch import nn
+
+import vietocr.model.backbone.vgg as vgg
+from vietocr.model.backbone.resnet import Resnet50
+
+class CNN(nn.Module):
+    def __init__(self, backbone, **kwargs):
+        super(CNN, self).__init__()
+
+        if backbone == 'vgg11_bn':
+            self.model = vgg.vgg11_bn(**kwargs)
+        elif backbone == 'vgg19_bn':
+            self.model = vgg.vgg19_bn(**kwargs)
+        elif backbone == 'resnet50':
+            self.model = Resnet50(**kwargs)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def freeze(self):
+        for name, param in self.model.features.named_parameters():
+            if name != 'last_conv_1x1':
+                param.requires_grad = False
+
+    def unfreeze(self):
+        for param in self.model.features.parameters():
+            param.requires_grad = True

vietocr/model/backbone/resnet.py

@@ -0,0 +1,140 @@
+import torch
+from torch import nn
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = self._conv3x3(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = self._conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def _conv3x3(self, in_planes, out_planes, stride=1):
+        "3x3 convolution with padding"
+        return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                         padding=1, bias=False)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+
+        return out
+    
+class ResNet(nn.Module):
+
+    def __init__(self, input_channel, output_channel, block, layers):
+        super(ResNet, self).__init__()
+
+        self.output_channel_block = [int(output_channel / 4), int(output_channel / 2), output_channel, output_channel]
+
+        self.inplanes = int(output_channel / 8)
+        self.conv0_1 = nn.Conv2d(input_channel, int(output_channel / 16),
+                                 kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn0_1 = nn.BatchNorm2d(int(output_channel / 16))
+        self.conv0_2 = nn.Conv2d(int(output_channel / 16), self.inplanes,
+                                 kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn0_2 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.layer1 = self._make_layer(block, self.output_channel_block[0], layers[0])
+        self.conv1 = nn.Conv2d(self.output_channel_block[0], self.output_channel_block[
+                               0], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.output_channel_block[0])
+
+        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
+        self.layer2 = self._make_layer(block, self.output_channel_block[1], layers[1], stride=1)
+        self.conv2 = nn.Conv2d(self.output_channel_block[1], self.output_channel_block[
+                               1], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(self.output_channel_block[1])
+
+        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=(2, 1), padding=(0, 1))
+        self.layer3 = self._make_layer(block, self.output_channel_block[2], layers[2], stride=1)
+        self.conv3 = nn.Conv2d(self.output_channel_block[2], self.output_channel_block[
+                               2], kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.output_channel_block[2])
+
+        self.layer4 = self._make_layer(block, self.output_channel_block[3], layers[3], stride=1)
+        self.conv4_1 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
+                                 3], kernel_size=2, stride=(2, 1), padding=(0, 1), bias=False)
+        self.bn4_1 = nn.BatchNorm2d(self.output_channel_block[3])
+        self.conv4_2 = nn.Conv2d(self.output_channel_block[3], self.output_channel_block[
+                                 3], kernel_size=2, stride=1, padding=0, bias=False)
+        self.bn4_2 = nn.BatchNorm2d(self.output_channel_block[3])
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv0_1(x)
+        x = self.bn0_1(x)
+        x = self.relu(x)
+        x = self.conv0_2(x)
+        x = self.bn0_2(x)
+        x = self.relu(x)
+
+        x = self.maxpool1(x)
+        x = self.layer1(x)
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.maxpool2(x)
+        x = self.layer2(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.maxpool3(x)
+        x = self.layer3(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.layer4(x)
+        x = self.conv4_1(x)
+        x = self.bn4_1(x)
+        x = self.relu(x)
+        x = self.conv4_2(x)
+        x = self.bn4_2(x)
+        conv = self.relu(x)
+        
+        conv = conv.transpose(-1, -2)
+        conv = conv.flatten(2)
+        conv = conv.permute(-1, 0, 1)
+
+        return conv
+
+def Resnet50(ss, hidden):
+    return ResNet(3, hidden, BasicBlock, [1, 2, 5, 3])
+

vietocr/model/backbone/vgg.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+from torchvision import models
+from einops import rearrange
+from torchvision.models._utils import IntermediateLayerGetter
+
+
+class Vgg(nn.Module):
+    def __init__(self, name, ss, ks, hidden, pretrained=True, dropout=0.5):
+        super(Vgg, self).__init__()
+
+        if name == 'vgg11_bn':
+            cnn = models.vgg11_bn(weights='DEFAULT')
+        elif name == 'vgg19_bn':
+            cnn = models.vgg19_bn(weights='DEFAULT')
+
+        pool_idx = 0
+        
+        for i, layer in enumerate(cnn.features):
+            if isinstance(layer, torch.nn.MaxPool2d):        
+                cnn.features[i] = torch.nn.AvgPool2d(kernel_size=ks[pool_idx], stride=ss[pool_idx], padding=0)
+                pool_idx += 1
+ 
+        self.features = cnn.features
+        self.dropout = nn.Dropout(dropout)
+        self.last_conv_1x1 = nn.Conv2d(512, hidden, 1)
+
+    def forward(self, x):
+        """
+        Shape: 
+            - x: (N, C, H, W)
+            - output: (W, N, C)
+        """
+
+        conv = self.features(x)
+        conv = self.dropout(conv)
+        conv = self.last_conv_1x1(conv)
+
+#        conv = rearrange(conv, 'b d h w -> b d (w h)')
+        conv = conv.transpose(-1, -2)
+        conv = conv.flatten(2)
+        conv = conv.permute(-1, 0, 1)
+        return conv
+
+def vgg11_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
+    return Vgg('vgg11_bn', ss, ks, hidden, pretrained, dropout)
+
+def vgg19_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
+    return Vgg('vgg19_bn', ss, ks, hidden, pretrained, dropout)
+   

vietocr/model/seqmodel/convseq2seq.py

@@ -0,0 +1,324 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+class Encoder(nn.Module):
+    def __init__(self, 
+                 emb_dim, 
+                 hid_dim, 
+                 n_layers, 
+                 kernel_size, 
+                 dropout, 
+                 device,
+                 max_length = 512):
+        super().__init__()
+        
+        assert kernel_size % 2 == 1, "Kernel size must be odd!"
+        
+        self.device = device
+        
+        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)
+        
+#         self.tok_embedding = nn.Embedding(input_dim, emb_dim)
+        self.pos_embedding = nn.Embedding(max_length, emb_dim)
+        
+        self.emb2hid = nn.Linear(emb_dim, hid_dim)
+        self.hid2emb = nn.Linear(hid_dim, emb_dim)
+        
+        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, 
+                                              out_channels = 2 * hid_dim, 
+                                              kernel_size = kernel_size, 
+                                              padding = (kernel_size - 1) // 2)
+                                    for _ in range(n_layers)])
+        
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, src):
+        
+        #src = [batch size, src len]
+        
+        src = src.transpose(0, 1)
+        
+        batch_size = src.shape[0]
+        src_len = src.shape[1]
+        device = src.device
+        
+        #create position tensor
+        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(device)
+        
+        #pos = [0, 1, 2, 3, ..., src len - 1]
+        
+        #pos = [batch size, src len]
+        
+        #embed tokens and positions
+        
+#         tok_embedded = self.tok_embedding(src)
+        tok_embedded = src
+    
+        pos_embedded = self.pos_embedding(pos)
+        
+        #tok_embedded = pos_embedded = [batch size, src len, emb dim]
+        
+        #combine embeddings by elementwise summing
+        embedded = self.dropout(tok_embedded + pos_embedded)
+        
+        #embedded = [batch size, src len, emb dim]
+        
+        #pass embedded through linear layer to convert from emb dim to hid dim
+        conv_input = self.emb2hid(embedded)
+        
+        #conv_input = [batch size, src len, hid dim]
+        
+        #permute for convolutional layer
+        conv_input = conv_input.permute(0, 2, 1) 
+        
+        #conv_input = [batch size, hid dim, src len]
+        
+        #begin convolutional blocks...
+        
+        for i, conv in enumerate(self.convs):
+        
+            #pass through convolutional layer
+            conved = conv(self.dropout(conv_input))
+
+            #conved = [batch size, 2 * hid dim, src len]
+
+            #pass through GLU activation function
+            conved = F.glu(conved, dim = 1)
+
+            #conved = [batch size, hid dim, src len]
+            
+            #apply residual connection
+            conved = (conved + conv_input) * self.scale
+
+            #conved = [batch size, hid dim, src len]
+            
+            #set conv_input to conved for next loop iteration
+            conv_input = conved
+        
+        #...end convolutional blocks
+        
+        #permute and convert back to emb dim
+        conved = self.hid2emb(conved.permute(0, 2, 1))
+        
+        #conved = [batch size, src len, emb dim]
+        
+        #elementwise sum output (conved) and input (embedded) to be used for attention
+        combined = (conved + embedded) * self.scale
+        
+        #combined = [batch size, src len, emb dim]
+        
+        return conved, combined
+
+class Decoder(nn.Module):
+    def __init__(self, 
+                 output_dim, 
+                 emb_dim, 
+                 hid_dim, 
+                 n_layers, 
+                 kernel_size, 
+                 dropout, 
+                 trg_pad_idx, 
+                 device,
+                 max_length = 512):
+        super().__init__()
+        
+        self.kernel_size = kernel_size
+        self.trg_pad_idx = trg_pad_idx
+        self.device = device
+        
+        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)
+        
+        self.tok_embedding = nn.Embedding(output_dim, emb_dim)
+        self.pos_embedding = nn.Embedding(max_length, emb_dim)
+        
+        self.emb2hid = nn.Linear(emb_dim, hid_dim)
+        self.hid2emb = nn.Linear(hid_dim, emb_dim)
+        
+        self.attn_hid2emb = nn.Linear(hid_dim, emb_dim)
+        self.attn_emb2hid = nn.Linear(emb_dim, hid_dim)
+        
+        self.fc_out = nn.Linear(emb_dim, output_dim)
+        
+        self.convs = nn.ModuleList([nn.Conv1d(in_channels = hid_dim, 
+                                              out_channels = 2 * hid_dim, 
+                                              kernel_size = kernel_size)
+                                    for _ in range(n_layers)])
+        
+        self.dropout = nn.Dropout(dropout)
+      
+    def calculate_attention(self, embedded, conved, encoder_conved, encoder_combined):
+        
+        #embedded = [batch size, trg len, emb dim]
+        #conved = [batch size, hid dim, trg len]
+        #encoder_conved = encoder_combined = [batch size, src len, emb dim]
+        
+        #permute and convert back to emb dim
+        conved_emb = self.attn_hid2emb(conved.permute(0, 2, 1))
+        
+        #conved_emb = [batch size, trg len, emb dim]
+        
+        combined = (conved_emb + embedded) * self.scale
+        
+        #combined = [batch size, trg len, emb dim]
+                
+        energy = torch.matmul(combined, encoder_conved.permute(0, 2, 1))
+        
+        #energy = [batch size, trg len, src len]
+        
+        attention = F.softmax(energy, dim=2)
+        
+        #attention = [batch size, trg len, src len]
+            
+        attended_encoding = torch.matmul(attention, encoder_combined)
+        
+        #attended_encoding = [batch size, trg len, emd dim]
+        
+        #convert from emb dim -> hid dim
+        attended_encoding = self.attn_emb2hid(attended_encoding)
+        
+        #attended_encoding = [batch size, trg len, hid dim]
+        
+        #apply residual connection
+        attended_combined = (conved + attended_encoding.permute(0, 2, 1)) * self.scale
+        
+        #attended_combined = [batch size, hid dim, trg len]
+        
+        return attention, attended_combined
+        
+    def forward(self, trg, encoder_conved, encoder_combined):
+        
+        #trg = [batch size, trg len]
+        #encoder_conved = encoder_combined = [batch size, src len, emb dim]
+        trg = trg.transpose(0, 1)
+        
+        batch_size = trg.shape[0]
+        trg_len = trg.shape[1]
+        device = trg.device
+        
+        #create position tensor
+        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(device)
+        
+        #pos = [batch size, trg len]
+        
+        #embed tokens and positions
+        tok_embedded = self.tok_embedding(trg)
+        pos_embedded = self.pos_embedding(pos)
+        
+        #tok_embedded = [batch size, trg len, emb dim]
+        #pos_embedded = [batch size, trg len, emb dim]
+        
+        #combine embeddings by elementwise summing
+        embedded = self.dropout(tok_embedded + pos_embedded)
+        
+        #embedded = [batch size, trg len, emb dim]
+        
+        #pass embedded through linear layer to go through emb dim -> hid dim
+        conv_input = self.emb2hid(embedded)
+        
+        #conv_input = [batch size, trg len, hid dim]
+        
+        #permute for convolutional layer
+        conv_input = conv_input.permute(0, 2, 1) 
+        
+        #conv_input = [batch size, hid dim, trg len]
+        
+        batch_size = conv_input.shape[0]
+        hid_dim = conv_input.shape[1]
+        
+        for i, conv in enumerate(self.convs):
+        
+            #apply dropout
+            conv_input = self.dropout(conv_input)
+        
+            #need to pad so decoder can't "cheat"
+            padding = torch.zeros(batch_size, 
+                                  hid_dim, 
+                                  self.kernel_size - 1).fill_(self.trg_pad_idx).to(device)
+                
+            padded_conv_input = torch.cat((padding, conv_input), dim = 2)
+        
+            #padded_conv_input = [batch size, hid dim, trg len + kernel size - 1]
+        
+            #pass through convolutional layer
+            conved = conv(padded_conv_input)
+
+            #conved = [batch size, 2 * hid dim, trg len]
+            
+            #pass through GLU activation function
+            conved = F.glu(conved, dim = 1)
+
+            #conved = [batch size, hid dim, trg len]
+            
+            #calculate attention
+            attention, conved = self.calculate_attention(embedded, 
+                                                         conved, 
+                                                         encoder_conved, 
+                                                         encoder_combined)
+            
+            #attention = [batch size, trg len, src len]
+            
+            #apply residual connection
+            conved = (conved + conv_input) * self.scale
+            
+            #conved = [batch size, hid dim, trg len]
+            
+            #set conv_input to conved for next loop iteration
+            conv_input = conved
+            
+        conved = self.hid2emb(conved.permute(0, 2, 1))
+         
+        #conved = [batch size, trg len, emb dim]
+            
+        output = self.fc_out(self.dropout(conved))
+        
+        #output = [batch size, trg len, output dim]
+            
+        return output, attention
+
+class ConvSeq2Seq(nn.Module):
+    def __init__(self, vocab_size, emb_dim, hid_dim, enc_layers, dec_layers, enc_kernel_size, dec_kernel_size, enc_max_length, dec_max_length, dropout, pad_idx, device):
+        super().__init__()
+        
+        enc = Encoder(emb_dim, hid_dim, enc_layers, enc_kernel_size, dropout, device, enc_max_length)
+        dec = Decoder(vocab_size, emb_dim, hid_dim, dec_layers, dec_kernel_size, dropout, pad_idx, device, dec_max_length)
+
+        self.encoder = enc
+        self.decoder = dec
+
+    def forward_encoder(self, src):
+        encoder_conved, encoder_combined = self.encoder(src)
+        
+        return encoder_conved, encoder_combined
+
+    def forward_decoder(self, trg, memory):
+        encoder_conved, encoder_combined = memory
+        output, attention = self.decoder(trg, encoder_conved, encoder_combined)
+        
+        return output, (encoder_conved, encoder_combined) 
+
+    def forward(self, src, trg):
+        
+        #src = [batch size, src len]
+        #trg = [batch size, trg len - 1] (<eos> token sliced off the end)
+           
+        #calculate z^u (encoder_conved) and (z^u + e) (encoder_combined)
+        #encoder_conved is output from final encoder conv. block
+        #encoder_combined is encoder_conved plus (elementwise) src embedding plus 
+        #  positional embeddings 
+        encoder_conved, encoder_combined = self.encoder(src)
+            
+        #encoder_conved = [batch size, src len, emb dim]
+        #encoder_combined = [batch size, src len, emb dim]
+        
+        #calculate predictions of next words
+        #output is a batch of predictions for each word in the trg sentence
+        #attention a batch of attention scores across the src sentence for 
+        #  each word in the trg sentence
+        output, attention = self.decoder(trg, encoder_conved, encoder_combined)
+        
+        #output = [batch size, trg len - 1, output dim]
+        #attention = [batch size, trg len - 1, src len]
+        
+        return output#, attention

vietocr/model/seqmodel/seq2seq.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+class Encoder(nn.Module):
+    def __init__(self, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
+        super().__init__()
+                
+        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
+        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, src):
+        """
+        src: src_len x batch_size x img_channel
+        outputs: src_len x batch_size x hid_dim 
+        hidden: batch_size x hid_dim
+        """
+
+        embedded = self.dropout(src)
+        
+        outputs, hidden = self.rnn(embedded)
+                                 
+        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
+        
+        return outputs, hidden
+
+class Attention(nn.Module):
+    def __init__(self, enc_hid_dim, dec_hid_dim):
+        super().__init__()
+        
+        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
+        self.v = nn.Linear(dec_hid_dim, 1, bias = False)
+        
+    def forward(self, hidden, encoder_outputs):
+        """
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim,
+        outputs: batch_size x src_len
+        """
+        
+        batch_size = encoder_outputs.shape[1]
+        src_len = encoder_outputs.shape[0]
+        
+        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
+  
+        encoder_outputs = encoder_outputs.permute(1, 0, 2)
+        
+        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim = 2))) 
+        
+        attention = self.v(energy).squeeze(2)
+        
+        return F.softmax(attention, dim = 1)
+
+class Decoder(nn.Module):
+    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
+        super().__init__()
+
+        self.output_dim = output_dim
+        self.attention = attention
+        
+        self.embedding = nn.Embedding(output_dim, emb_dim)
+        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
+        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, input, hidden, encoder_outputs):
+        """
+        inputs: batch_size
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim
+        """
+             
+        input = input.unsqueeze(0)
+        
+        embedded = self.dropout(self.embedding(input))
+        
+        a = self.attention(hidden, encoder_outputs)
+                
+        a = a.unsqueeze(1)
+        
+        encoder_outputs = encoder_outputs.permute(1, 0, 2)
+        
+        weighted = torch.bmm(a, encoder_outputs)
+        
+        weighted = weighted.permute(1, 0, 2)
+        
+        rnn_input = torch.cat((embedded, weighted), dim = 2)
+        
+        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
+        
+        assert (output == hidden).all()
+        
+        embedded = embedded.squeeze(0)
+        output = output.squeeze(0)
+        weighted = weighted.squeeze(0)
+        
+        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))
+        
+        return prediction, hidden.squeeze(0), a.squeeze(1)
+
+class Seq2Seq(nn.Module):
+    def __init__(self, vocab_size, encoder_hidden, decoder_hidden, img_channel, decoder_embedded, dropout=0.1):
+        super().__init__()
+        
+        attn = Attention(encoder_hidden, decoder_hidden)
+        
+        self.encoder = Encoder(img_channel, encoder_hidden, decoder_hidden, dropout)
+        self.decoder = Decoder(vocab_size, decoder_embedded, encoder_hidden, decoder_hidden, dropout, attn)
+        
+    def forward_encoder(self, src):       
+        """
+        src: timestep x batch_size x channel
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim
+        """
+
+        encoder_outputs, hidden = self.encoder(src)
+
+        return (hidden, encoder_outputs)
+
+    def forward_decoder(self, tgt, memory):
+        """
+        tgt: timestep x batch_size 
+        hidden: batch_size x hid_dim
+        encouder: src_len x batch_size x hid_dim
+        output: batch_size x 1 x vocab_size
+        """
+        
+        tgt = tgt[-1]
+        hidden, encoder_outputs = memory
+        output, hidden, _ = self.decoder(tgt, hidden, encoder_outputs)
+        output = output.unsqueeze(1)
+        
+        return output, (hidden, encoder_outputs)
+
+    def forward(self, src, trg):
+        """
+        src: time_step x batch_size
+        trg: time_step x batch_size
+        outputs: batch_size x time_step x vocab_size
+        """
+
+        batch_size = src.shape[1]
+        trg_len = trg.shape[0]
+        trg_vocab_size = self.decoder.output_dim
+        device = src.device
+
+        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(device)
+        encoder_outputs, hidden = self.encoder(src)
+                
+        for t in range(trg_len):
+            input = trg[t] 
+            output, hidden, _ = self.decoder(input, hidden, encoder_outputs)
+            
+            outputs[t] = output
+            
+        outputs = outputs.transpose(0, 1).contiguous()
+
+        return outputs
+
+    def expand_memory(self, memory, beam_size):
+        hidden, encoder_outputs = memory
+        hidden = hidden.repeat(beam_size, 1)
+        encoder_outputs = encoder_outputs.repeat(1, beam_size, 1)
+
+        return (hidden, encoder_outputs)
+    
+    def get_memory(self, memory, i):
+        hidden, encoder_outputs = memory
+        hidden = hidden[[i]]
+        encoder_outputs = encoder_outputs[:, [i],:]
+
+        return (hidden, encoder_outputs)

vietocr/model/seqmodel/transformer.py

@@ -0,0 +1,124 @@
+from einops import rearrange
+from torchvision import models
+import math
+import torch
+from torch import nn
+
+class LanguageTransformer(nn.Module):
+    def __init__(self, vocab_size, 
+                 d_model, nhead, 
+                 num_encoder_layers, num_decoder_layers, 
+                 dim_feedforward, max_seq_length, 
+                 pos_dropout, trans_dropout):
+        super().__init__()
+        
+        self.d_model = d_model
+        self.embed_tgt = nn.Embedding(vocab_size, d_model)
+        self.pos_enc = PositionalEncoding(d_model, pos_dropout, max_seq_length)
+#        self.learned_pos_enc = LearnedPositionalEncoding(d_model, pos_dropout, max_seq_length)
+
+        self.transformer = nn.Transformer(d_model, nhead, 
+                                          num_encoder_layers, num_decoder_layers, 
+                                          dim_feedforward, trans_dropout)
+        
+        self.fc = nn.Linear(d_model, vocab_size)
+        
+    def forward(self, src, tgt, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        """
+        Shape:
+            - src: (W, N, C)
+            - tgt: (T, N) 
+            - src_key_padding_mask: (N, S)
+            - tgt_key_padding_mask: (N, T)
+            - memory_key_padding_mask: (N, S)
+            - output: (N, T, E)
+            
+        """
+        tgt_mask = self.gen_nopeek_mask(tgt.shape[0]).to(src.device)
+        
+        src = self.pos_enc(src*math.sqrt(self.d_model))
+#        src = self.learned_pos_enc(src*math.sqrt(self.d_model))
+
+        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
+        
+        output = self.transformer(src, tgt, tgt_mask=tgt_mask, src_key_padding_mask=src_key_padding_mask,
+                                  tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask)
+#        output = rearrange(output, 't n e -> n t e')
+        output = output.transpose(0, 1)
+        return self.fc(output)
+
+    def gen_nopeek_mask(self, length):
+        mask = (torch.triu(torch.ones(length, length)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+
+        return mask
+    
+    def forward_encoder(self, src):
+        src = self.pos_enc(src*math.sqrt(self.d_model))
+        memory = self.transformer.encoder(src)
+        return memory
+    
+    def forward_decoder(self, tgt, memory):
+        tgt_mask = self.gen_nopeek_mask(tgt.shape[0]).to(tgt.device)
+        tgt = self.pos_enc(self.embed_tgt(tgt) * math.sqrt(self.d_model))
+        
+        output = self.transformer.decoder(tgt, memory, tgt_mask=tgt_mask)
+#        output = rearrange(output, 't n e -> n t e')
+        output = output.transpose(0, 1)
+
+        return self.fc(output), memory
+    
+    def expand_memory(self, memory, beam_size):
+        memory = memory.repeat(1, beam_size, 1)
+        return memory
+    
+    def get_memory(self, memory, i):
+        memory = memory[:, [i], :]
+        return memory
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=100):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[:x.size(0), :]
+
+        return self.dropout(x)
+ 
+class LearnedPositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=100):
+        super(LearnedPositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        self.pos_embed = nn.Embedding(max_len, d_model)
+        self.layernorm = LayerNorm(d_model)
+
+    def forward(self, x):
+        seq_len = x.size(0)
+        pos = torch.arange(seq_len, dtype=torch.long, device=x.device)
+        pos = pos.unsqueeze(-1).expand(x.size()[:2])
+        x = x + self.pos_embed(pos)
+        return self.dropout(self.layernorm(x))
+
+class LayerNorm(nn.Module):
+    "A layernorm module in the TF style (epsilon inside the square root)."
+    def __init__(self, d_model, variance_epsilon=1e-12):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.ones(d_model))
+        self.beta  = nn.Parameter(torch.zeros(d_model))
+        self.variance_epsilon = variance_epsilon
+
+    def forward(self, x):
+        u = x.mean(-1, keepdim=True)
+        s = (x - u).pow(2).mean(-1, keepdim=True)
+        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
+        return self.gamma * x + self.beta  

vietocr/model/transformerocr.py

@@ -0,0 +1,44 @@
+from vietocr.model.backbone.cnn import CNN
+from vietocr.model.seqmodel.transformer import LanguageTransformer
+from vietocr.model.seqmodel.seq2seq import Seq2Seq
+from vietocr.model.seqmodel.convseq2seq import ConvSeq2Seq
+from torch import nn
+
+class VietOCR(nn.Module):
+    def __init__(self, vocab_size,
+                 backbone,
+                 cnn_args, 
+                 transformer_args, seq_modeling='transformer'):
+        
+        super(VietOCR, self).__init__()
+        
+        self.cnn = CNN(backbone, **cnn_args)
+        self.seq_modeling = seq_modeling
+
+        if seq_modeling == 'transformer':
+            self.transformer = LanguageTransformer(vocab_size, **transformer_args)
+        elif seq_modeling == 'seq2seq':
+            self.transformer = Seq2Seq(vocab_size, **transformer_args)
+        elif seq_modeling == 'convseq2seq':
+            self.transformer = ConvSeq2Seq(vocab_size, **transformer_args)
+        else:
+            raise('Not Support Seq Model')
+
+    def forward(self, img, tgt_input, tgt_key_padding_mask):
+        """
+        Shape:
+            - img: (N, C, H, W)
+            - tgt_input: (T, N)
+            - tgt_key_padding_mask: (N, T)
+            - output: b t v
+        """
+        src = self.cnn(img)
+
+        if self.seq_modeling == 'transformer':
+            outputs = self.transformer(src, tgt_input, tgt_key_padding_mask=tgt_key_padding_mask)
+        elif self.seq_modeling == 'seq2seq':
+            outputs = self.transformer(src, tgt_input)
+        elif self.seq_modeling == 'convseq2seq': 
+            outputs = self.transformer(src, tgt_input)
+        return outputs
+

vietocr/model/vocab.py

@@ -0,0 +1,36 @@
+class Vocab():
+    def __init__(self, chars):
+        self.pad = 0
+        self.go = 1
+        self.eos = 2
+        self.mask_token = 3
+
+        self.chars = chars
+
+        self.c2i = {c:i+4 for i, c in enumerate(chars)}
+
+        self.i2c = {i+4:c for i, c in enumerate(chars)}
+        
+        self.i2c[0] = '<pad>'
+        self.i2c[1] = '<sos>'
+        self.i2c[2] = '<eos>'
+        self.i2c[3] = '*'
+
+    def encode(self, chars):
+        return [self.go] + [self.c2i[c] for c in chars] + [self.eos]
+    
+    def decode(self, ids):
+        first = 1 if self.go in ids else 0
+        last = ids.index(self.eos) if self.eos in ids else None
+        sent = ''.join([self.i2c[i] for i in ids[first:last]])
+        return sent
+    
+    def __len__(self):
+        return len(self.c2i) + 4
+    
+    def batch_decode(self, arr):
+        texts = [self.decode(ids) for ids in arr]
+        return texts
+
+    def __str__(self):
+        return self.chars

vietocr/translate.py

@@ -0,0 +1,62 @@
+import torch
+import numpy as np
+import math
+from PIL import Image
+from torch.nn.functional import softmax
+
+def translate(img, model, max_seq_length=128, sos_token=1, eos_token=2):
+    "data: BxCXHxW"
+    model.eval()
+
+    with torch.no_grad():
+        src = model.cnn(img)
+        memory = model.transformer.forward_encoder(src)
+
+        translated_sentence = [[sos_token]*len(img)]
+
+        max_length = 0
+
+        while max_length <= max_seq_length and not all(np.any(np.asarray(translated_sentence).T==eos_token, axis=1)):
+            tgt_inp = torch.LongTensor(translated_sentence)
+            
+            output, memory = model.transformer.forward_decoder(tgt_inp, memory)
+            output = softmax(output, dim=-1)
+
+            _, indices  = torch.topk(output, 5)
+            
+            indices = indices[:, -1, 0]
+            indices = indices.tolist()
+
+            translated_sentence.append(indices)   
+            max_length += 1
+
+        translated_sentence = np.asarray(translated_sentence).T
+    
+    return translated_sentence
+
+def resize(w, h, expected_height, image_min_width, image_max_width):
+    new_w = int(expected_height * float(w) / float(h))
+    round_to = 10
+    new_w = math.ceil(new_w/round_to)*round_to
+    new_w = max(new_w, image_min_width)
+    new_w = min(new_w, image_max_width)
+
+    return new_w, expected_height
+
+def process_image(image, image_height, image_min_width, image_max_width):
+    img = image.convert('RGB')
+
+    w, h = img.size
+    new_w, image_height = resize(w, h, image_height, image_min_width, image_max_width)
+
+    img = img.resize((new_w, image_height), Image.Resampling.LANCZOS)
+
+    img = np.asarray(img).transpose(2,0, 1)
+    img = img/255
+    return img
+
+def process_input(image, image_height, image_min_width, image_max_width):
+    img = process_image(image, image_height, image_min_width, image_max_width)
+    img = img[np.newaxis, ...]
+    img = torch.FloatTensor(img)
+    return img