Upload 38 files

vietocr/.gitignore

@@ -0,0 +1,140 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+my_train_vietocr.py/

vietocr/config.yml

@@ -0,0 +1,69 @@
+aug:
+  image_aug: true
+  masked_language_model: true
+backbone: vgg19_bn
+cnn:
+  hidden: 256
+  ks:
+  - - 2
+    - 2
+  - - 2
+    - 2
+  - - 2
+    - 1
+  - - 2
+    - 1
+  - - 1
+    - 1
+  pretrained: true
+  ss:
+  - - 2
+    - 2
+  - - 2
+    - 2
+  - - 2
+    - 1
+  - - 2
+    - 1
+  - - 1
+    - 1
+dataloader:
+  num_workers: 0
+  pin_memory: true
+dataset:
+  data_root: /Users/bmd1905/Desktop/dataset/data_line
+  image_height: 32
+  image_max_width: 512
+  image_min_width: 32
+  name: hw
+  train_annotation: train_line_annotation.txt
+  valid_annotation: test_line_annotation.txt
+device: mps
+optimizer:
+  max_lr: 0.0003
+  pct_start: 0.1
+predictor:
+  beamsearch: false
+pretrain: https://vocr.vn/data/vietocr/vgg_transformer.pth
+quiet: false
+seq_modeling: transformer
+trainer:
+  batch_size: 1
+  checkpoint: ./checkpoint/transformerocr_checkpoint.pth
+  export: ./weights/seq2seq_test_local.pth
+  iters: 100
+  log: ./train.log
+  metrics: 100
+  print_every: 1
+  valid_every: 10
+transformer:
+  d_model: 256
+  dim_feedforward: 2048
+  max_seq_length: 1024
+  nhead: 8
+  num_decoder_layers: 6
+  num_encoder_layers: 6
+  pos_dropout: 0.1
+  trans_dropout: 0.1
+vocab: 'aAàÀảẢãÃáÁạẠăĂằẰẳẲẵẴắẮặẶâÂầẦẩẨẫẪấẤậẬbBcCdDđĐeEèÈẻẺẽẼéÉẹẸêÊềỀểỂễỄếẾệỆfFgGhHiIìÌỉỈĩĨíÍịỊjJkKlLmMnNoOòÒỏỎõÕóÓọỌôÔồỒổỔỗỖốỐộỘơƠờỜởỞỡỠớỚợỢpPqQrRsStTuUùÙủỦũŨúÚụỤưƯừỪửỬữỮứỨựỰvVwWxXyYỳỲỷỶỹỸýÝỵỴzZ0123456789!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
+weights: https://vocr.vn/data/vietocr/vgg_transformer.pth

vietocr/config/base.yml

@@ -0,0 +1,58 @@
+# change to list chars of your dataset or use default vietnamese chars
+vocab: 'aAàÀảẢãÃáÁạẠăĂằẰẳẲẵẴắẮặẶâÂầẦẩẨẫẪấẤậẬbBcCdDđĐeEèÈẻẺẽẼéÉẹẸêÊềỀểỂễỄếẾệỆfFgGhHiIìÌỉỈĩĨíÍịỊjJkKlLmMnNoOòÒỏỎõÕóÓọỌôÔồỒổỔỗỖốỐộỘơƠờỜởỞỡỠớỚợỢpPqQrRsStTuUùÙủỦũŨúÚụỤưƯừỪửỬữỮứỨựỰvVwWxXyYỳỲỷỶỹỸýÝỵỴzZ0123456789!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
+
+# cpu, cuda, cuda:0
+device: cuda:0
+
+seq_modeling: transformer
+transformer:  
+    d_model: 256
+    nhead: 8
+    num_encoder_layers: 6
+    num_decoder_layers: 6
+    dim_feedforward: 2048
+    max_seq_length: 1024
+    pos_dropout: 0.1
+    trans_dropout: 0.1
+
+optimizer:
+    max_lr: 0.0003 
+    pct_start: 0.1
+
+trainer:
+    batch_size: 32
+    print_every: 200
+    valid_every: 4000
+    iters: 100000
+    # where to save our model for prediction
+    export: ./weights/transformerocr.pth
+    checkpoint: ./checkpoint/transformerocr_checkpoint.pth
+    log: ./train.log
+    # null to disable compuate accuracy, or change to number of sample to enable validiation while training
+    metrics: null
+
+dataset:    
+    # name of your dataset
+    name: data
+    # path to annotation and image
+    data_root: ./img/
+    train_annotation: annotation_train.txt
+    valid_annotation: annotation_val_small.txt
+    # resize image to 32 height, larger height will increase accuracy
+    image_height: 32
+    image_min_width: 32
+    image_max_width: 512
+
+dataloader:
+    num_workers: 3
+    pin_memory: True
+
+aug:
+    image_aug: true
+    masked_language_model: true
+
+predictor:
+    # disable or enable beamsearch while prediction, use beamsearch will be slower
+    beamsearch: False
+
+quiet: False 

vietocr/config/resnet-transformer.yml

@@ -0,0 +1,16 @@
+pretrain: 
+    id_or_url: 13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA
+    md5: 7068030afe2e8fc639d0e1e2c25612b3
+    cached: /tmp/tranformerorc.pth
+
+weights: https://drive.google.com/uc?id=12dTOZ9VP7ZVzwQgVvqBWz5JO5RXXW5NY
+
+backbone: resnet50
+cnn:
+    ss:
+        - [2, 2]
+        - [2, 1]
+        - [2, 1]
+        - [2, 1]
+        - [1, 1]          
+    hidden: 256

vietocr/config/resnet_fpn_transformer.yml

@@ -0,0 +1,9 @@
+pretrain: 
+    id_or_url: 13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA
+    md5: 7068030afe2e8fc639d0e1e2c25612b3
+    cached: /tmp/tranformerorc.pth
+
+weights: https://drive.google.com/uc?id=12dTOZ9VP7ZVzwQgVvqBWz5JO5RXXW5NY
+
+backbone: resnet50_fpn
+cnn: {}

vietocr/config/vgg-convseq2seq.yml

@@ -0,0 +1,40 @@
+pretrain: 
+    id_or_url: 13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA
+    md5: fbefa85079ad9001a71eb1bf47a93785
+    cached: /tmp/tranformerorc.pth
+
+# url or local path
+weights: https://drive.google.com/uc?id=13327Y1tz1ohsm5YZMyXVMPIOjoOA0OaA
+
+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
+
+seq_modeling: convseq2seq
+transformer:
+    emb_dim: 256
+    hid_dim: 512
+    enc_layers: 10
+    dec_layers: 10
+    enc_kernel_size: 3
+    dec_kernel_size: 3
+    dropout: 0.1
+    pad_idx: 0
+    device: cuda:1
+    enc_max_length: 512
+    dec_max_length: 512

vietocr/config/vgg-seq2seq.yml

@@ -0,0 +1,37 @@
+# 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
+
+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

vietocr/config/vgg-transformer.yml

@@ -0,0 +1,26 @@
+# for training
+pretrain: https://vocr.vn/data/vietocr/vgg_transformer.pth
+
+# url or local path (predict)
+weights: https://vocr.vn/data/vietocr/vgg_transformer.pth
+
+backbone: vgg19_bn
+cnn:
+    pretrained: True
+    # 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/vietocr/loader/aug.py

@@ -0,0 +1,48 @@
+from PIL import Image
+import numpy as np
+
+from imgaug import augmenters as iaa
+import imgaug as ia
+
+class ImgAugTransform:
+  def __init__(self):
+    sometimes = lambda aug: iaa.Sometimes(0.3, aug)
+
+    self.aug = iaa.Sequential(iaa.SomeOf((1, 5), 
+        [
+        # blur
+
+        sometimes(iaa.OneOf([iaa.GaussianBlur(sigma=(0, 1.0)),
+                            iaa.MotionBlur(k=3)])),
+        
+        # color
+        sometimes(iaa.AddToHueAndSaturation(value=(-10, 10), per_channel=True)),
+        sometimes(iaa.SigmoidContrast(gain=(3, 10), cutoff=(0.4, 0.6), per_channel=True)),
+        sometimes(iaa.Invert(0.25, per_channel=0.5)),
+        sometimes(iaa.Solarize(0.5, threshold=(32, 128))),
+        sometimes(iaa.Dropout2d(p=0.5)),
+        sometimes(iaa.Multiply((0.5, 1.5), per_channel=0.5)),
+        sometimes(iaa.Add((-40, 40), per_channel=0.5)),
+
+        sometimes(iaa.JpegCompression(compression=(5, 80))),
+        
+        # distort
+        sometimes(iaa.Crop(percent=(0.01, 0.05), sample_independently=True)),
+        sometimes(iaa.PerspectiveTransform(scale=(0.01, 0.01))),
+        sometimes(iaa.Affine(scale=(0.7, 1.3), translate_percent=(-0.1, 0.1), 
+#                            rotate=(-5, 5), shear=(-5, 5), 
+                            order=[0, 1], cval=(0, 255), 
+                            mode=ia.ALL)),
+        sometimes(iaa.PiecewiseAffine(scale=(0.01, 0.01))),
+        sometimes(iaa.OneOf([iaa.Dropout(p=(0, 0.1)),
+                            iaa.CoarseDropout(p=(0, 0.1), size_percent=(0.02, 0.25))])),
+
+    ],
+        random_order=True),
+    random_order=True)
+      
+  def __call__(self, img):
+    img = np.array(img)
+    img = self.aug.augment_image(img)
+    img = Image.fromarray(img)
+    return img

vietocr/vietocr/loader/dataloader.py

@@ -0,0 +1,205 @@
+import sys
+import os
+import random
+from PIL import Image
+from PIL import ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+from collections import defaultdict
+import numpy as np
+import torch
+import lmdb
+import six
+import time
+from tqdm import tqdm
+
+from torch.utils.data import Dataset
+from torch.utils.data.sampler import Sampler
+from vietocr.vietocr.tool.translate import process_image
+from vietocr.vietocr.tool.create_dataset import createDataset
+from vietocr.vietocr.tool.translate import resize
+
+class OCRDataset(Dataset):
+    def __init__(self, lmdb_path, root_dir, annotation_path, vocab, image_height=32, image_min_width=32, image_max_width=512, transform=None):
+        self.root_dir = root_dir
+        self.annotation_path = os.path.join(root_dir, annotation_path)
+        self.vocab = vocab
+        self.transform = transform
+
+        self.image_height = image_height
+        self.image_min_width = image_min_width
+        self.image_max_width = image_max_width
+
+        self.lmdb_path =  lmdb_path
+
+        if os.path.isdir(self.lmdb_path):
+            print('{} exists. Remove folder if you want to create new dataset'.format(self.lmdb_path))
+            sys.stdout.flush()
+        else:
+            createDataset(self.lmdb_path, root_dir, annotation_path)
+        
+        self.env = lmdb.open(
+            self.lmdb_path,
+            max_readers=8,
+            readonly=True,
+            lock=False,
+            readahead=False,
+            meminit=False)
+        self.txn = self.env.begin(write=False)
+
+        nSamples = int(self.txn.get('num-samples'.encode()))
+        self.nSamples = nSamples
+
+        self.build_cluster_indices()
+
+    def build_cluster_indices(self):
+        self.cluster_indices = defaultdict(list)
+
+        pbar = tqdm(range(self.__len__()), 
+                desc='{} build cluster'.format(self.lmdb_path), 
+                ncols = 100, position=0, leave=True) 
+
+        for i in pbar:
+            bucket = self.get_bucket(i)
+            self.cluster_indices[bucket].append(i)
+
+    
+    def get_bucket(self, idx):
+        key = 'dim-%09d'%idx
+
+        dim_img = self.txn.get(key.encode())
+        dim_img = np.fromstring(dim_img, dtype=np.int32)
+        imgH, imgW = dim_img
+
+        new_w, image_height = resize(imgW, imgH, self.image_height, self.image_min_width, self.image_max_width)
+
+        return new_w
+
+    def read_buffer(self, idx):
+        img_file = 'image-%09d'%idx
+        label_file = 'label-%09d'%idx
+        path_file = 'path-%09d'%idx
+        
+        imgbuf = self.txn.get(img_file.encode())
+        
+        label = self.txn.get(label_file.encode()).decode()
+        img_path = self.txn.get(path_file.encode()).decode()
+
+        buf = six.BytesIO()
+        buf.write(imgbuf)
+        buf.seek(0)
+    
+        return buf, label, img_path
+
+    def read_data(self, idx):
+        buf, label, img_path = self.read_buffer(idx) 
+
+        img = Image.open(buf).convert('RGB')        
+       
+        if self.transform:
+            img = self.transform(img)
+
+        img_bw = process_image(img, self.image_height, self.image_min_width, self.image_max_width)
+            
+        word = self.vocab.encode(label)
+
+        return img_bw, word, img_path
+
+    def __getitem__(self, idx):
+        img, word, img_path = self.read_data(idx)
+        
+        img_path = os.path.join(self.root_dir, img_path)
+        
+        sample = {'img': img, 'word': word, 'img_path': img_path}
+
+        return sample
+
+    def __len__(self):
+        return self.nSamples
+
+class ClusterRandomSampler(Sampler):
+    
+    def __init__(self, data_source, batch_size, shuffle=True):
+        self.data_source = data_source
+        self.batch_size = batch_size
+        self.shuffle = shuffle        
+
+    def flatten_list(self, lst):
+        return [item for sublist in lst for item in sublist]
+
+    def __iter__(self):
+        batch_lists = []
+        for cluster, cluster_indices in self.data_source.cluster_indices.items():
+            if self.shuffle:
+                random.shuffle(cluster_indices)
+
+            batches = [cluster_indices[i:i + self.batch_size] for i in range(0, len(cluster_indices), self.batch_size)]
+            batches = [_ for _ in batches if len(_) == self.batch_size]
+            if self.shuffle:
+                random.shuffle(batches)
+
+            batch_lists.append(batches)
+
+        lst = self.flatten_list(batch_lists)
+        if self.shuffle:
+            random.shuffle(lst)
+
+        lst = self.flatten_list(lst)
+
+        return iter(lst)
+
+    def __len__(self):
+        return len(self.data_source)
+
+class Collator(object):
+    def __init__(self, masked_language_model=True):
+        self.masked_language_model = masked_language_model
+
+    def __call__(self, batch):
+        filenames = []
+        img = []
+        target_weights = []
+        tgt_input = []
+        max_label_len = max(len(sample['word']) for sample in batch)
+        for sample in batch:
+            img.append(sample['img'])
+            filenames.append(sample['img_path'])
+            label = sample['word']
+            label_len = len(label)
+            
+            
+            tgt = np.concatenate((
+                label,
+                np.zeros(max_label_len - label_len, dtype=np.int32)))
+            tgt_input.append(tgt)
+
+            one_mask_len = label_len - 1
+
+            target_weights.append(np.concatenate((
+                np.ones(one_mask_len, dtype=np.float32),
+                np.zeros(max_label_len - one_mask_len,dtype=np.float32))))
+            
+        img = np.array(img, dtype=np.float32)
+
+
+        tgt_input = np.array(tgt_input, dtype=np.int64).T
+        tgt_output = np.roll(tgt_input, -1, 0).T
+        tgt_output[:, -1]=0
+        
+        # random mask token
+        if self.masked_language_model:
+            mask = np.random.random(size=tgt_input.shape) < 0.05
+            mask = mask & (tgt_input != 0) & (tgt_input != 1) & (tgt_input != 2)
+            tgt_input[mask] = 3
+
+        tgt_padding_mask = np.array(target_weights)==0
+
+        rs = {
+            'img': torch.FloatTensor(img),
+            'tgt_input': torch.LongTensor(tgt_input),
+            'tgt_output': torch.LongTensor(tgt_output),
+            'tgt_padding_mask': torch.BoolTensor(tgt_padding_mask),
+            'filenames': filenames
+        }   
+        
+        return rs

vietocr/vietocr/loader/dataloader_v1.py

@@ -0,0 +1,155 @@
+import torch
+import numpy as np
+from PIL import Image
+import random
+from vietocr.vietocr.model.vocab import Vocab
+from vietocr.vietocr.tool.translate import process_image
+import os
+from collections import defaultdict
+import math
+from prefetch_generator import background
+
+class BucketData(object):
+    def __init__(self, device):
+        self.max_label_len = 0
+        self.data_list = []
+        self.label_list = []
+        self.file_list = []
+        self.device = device
+
+    def append(self, datum, label, filename):
+        self.data_list.append(datum)
+        self.label_list.append(label)
+        self.file_list.append(filename)
+        
+        self.max_label_len = max(len(label), self.max_label_len)
+
+        return len(self.data_list)
+
+    def flush_out(self):                           
+        """
+        Shape:
+            - img: (N, C, H, W) 
+            - tgt_input: (T, N) 
+            - tgt_output: (N, T) 
+            - tgt_padding_mask: (N, T) 
+        """
+        # encoder part
+        img = np.array(self.data_list, dtype=np.float32)
+        
+        # decoder part
+        target_weights = []
+        tgt_input = []
+        for label in self.label_list:
+            label_len = len(label)
+            
+            tgt = np.concatenate((
+                label,
+                np.zeros(self.max_label_len - label_len, dtype=np.int32)))
+            tgt_input.append(tgt)
+            
+            one_mask_len = label_len - 1
+            
+            target_weights.append(np.concatenate((
+                np.ones(one_mask_len, dtype=np.float32),
+                np.zeros(self.max_label_len - one_mask_len,dtype=np.float32))))
+
+        # reshape to fit input shape
+        tgt_input = np.array(tgt_input, dtype=np.int64).T
+        tgt_output = np.roll(tgt_input, -1, 0).T
+        tgt_output[:, -1]=0
+        
+        tgt_padding_mask = np.array(target_weights)==0
+
+        filenames = self.file_list
+
+        self.data_list, self.label_list, self.file_list = [], [], []
+        self.max_label_len = 0
+        
+        rs = {
+            'img': torch.FloatTensor(img).to(self.device),
+            'tgt_input': torch.LongTensor(tgt_input).to(self.device),
+            'tgt_output': torch.LongTensor(tgt_output).to(self.device),
+            'tgt_padding_mask':torch.BoolTensor(tgt_padding_mask).to(self.device),
+            'filenames': filenames
+        }
+        
+        return rs
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __iadd__(self, other):
+        self.data_list += other.data_list
+        self.label_list += other.label_list
+        self.max_label_len = max(self.max_label_len, other.max_label_len)
+        self.max_width = max(self.max_width, other.max_width)
+
+    def __add__(self, other):
+        res = BucketData()
+        res.data_list = self.data_list + other.data_list
+        res.label_list = self.label_list + other.label_list
+        res.max_width = max(self.max_width, other.max_width)
+        res.max_label_len = max((self.max_label_len, other.max_label_len))
+        return res
+
+class DataGen(object):
+
+    def __init__(self,data_root, annotation_fn, vocab, device, image_height=32, image_min_width=32, image_max_width=512):
+        
+        self.image_height = image_height
+        self.image_min_width = image_min_width
+        self.image_max_width = image_max_width
+
+        self.data_root = data_root
+        self.annotation_path = os.path.join(data_root, annotation_fn)
+        
+        self.vocab = vocab
+        self.device = device
+        
+        self.clear()
+
+    def clear(self):
+        self.bucket_data = defaultdict(lambda: BucketData(self.device))
+
+    @background(max_prefetch=1) 
+    def gen(self, batch_size, last_batch=True):
+        with open(self.annotation_path, 'r') as ann_file:
+            lines = ann_file.readlines()
+            np.random.shuffle(lines)
+            for l in lines:     
+                
+                img_path, lex = l.strip().split('\t')
+                
+                img_path = os.path.join(self.data_root, img_path)
+                
+                try:
+                    img_bw, word = self.read_data(img_path, lex)
+                except IOError:
+                    print('ioread image:{}'.format(img_path))
+                    
+                width = img_bw.shape[-1]
+
+                bs = self.bucket_data[width].append(img_bw, word, img_path)
+                if bs >= batch_size:
+                    b = self.bucket_data[width].flush_out()
+                    yield b
+
+        if last_batch: 
+            for bucket in self.bucket_data.values():
+                if len(bucket) > 0:
+                    b = bucket.flush_out()
+                    yield b
+
+        self.clear()
+
+    def read_data(self, img_path, lex):        
+        
+        with open(img_path, 'rb') as img_file:
+            img = Image.open(img_file).convert('RGB')
+            img_bw = process_image(img, self.image_height, self.image_min_width, self.image_max_width)
+
+        word = self.vocab.encode(lex)
+
+        return img_bw, word
+

vietocr/vietocr/model/backbone/cnn.py

@@ -0,0 +1,28 @@
+import torch
+from torch import nn
+
+import vietocr.vietocr.model.backbone.vgg as vgg
+from vietocr.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/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/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(pretrained=pretrained)
+        elif name == 'vgg19_bn':
+            cnn = models.vgg19_bn(pretrained=pretrained)
+
+        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/vietocr/model/beam.py

@@ -0,0 +1,104 @@
+import torch
+
+class Beam:
+
+    def __init__(self, beam_size=8, min_length=0, n_top=1, ranker=None,
+                 start_token_id=1, end_token_id=2):
+        self.beam_size = beam_size
+        self.min_length = min_length
+        self.ranker = ranker
+
+        self.end_token_id = end_token_id
+        self.top_sentence_ended = False
+
+        self.prev_ks = []
+        self.next_ys = [torch.LongTensor(beam_size).fill_(start_token_id)] # remove padding
+
+        self.current_scores = torch.FloatTensor(beam_size).zero_()
+        self.all_scores = []
+
+        # Time and k pair for finished.
+        self.finished = []
+        self.n_top = n_top
+
+        self.ranker = ranker
+
+    def advance(self, next_log_probs):
+        # next_probs : beam_size X vocab_size
+
+        vocabulary_size = next_log_probs.size(1)
+        # current_beam_size = next_log_probs.size(0)
+
+        current_length = len(self.next_ys)
+        if current_length < self.min_length:
+            for beam_index in range(len(next_log_probs)):
+                next_log_probs[beam_index][self.end_token_id] = -1e10
+
+        if len(self.prev_ks) > 0:
+            beam_scores = next_log_probs + self.current_scores.unsqueeze(1).expand_as(next_log_probs)
+            # Don't let EOS have children.
+            last_y = self.next_ys[-1]
+            for beam_index in range(last_y.size(0)):
+                if last_y[beam_index] == self.end_token_id:
+                    beam_scores[beam_index] = -1e10 # -1e20 raises error when executing
+        else:
+            beam_scores = next_log_probs[0]
+            
+        flat_beam_scores = beam_scores.view(-1)
+        top_scores, top_score_ids = flat_beam_scores.topk(k=self.beam_size, dim=0, largest=True, sorted=True)
+
+        self.current_scores = top_scores
+        self.all_scores.append(self.current_scores)
+        
+        prev_k = top_score_ids // vocabulary_size  # (beam_size, )
+        next_y = top_score_ids - prev_k * vocabulary_size  # (beam_size, )
+                
+        
+        self.prev_ks.append(prev_k)
+        self.next_ys.append(next_y)
+
+        for beam_index, last_token_id in enumerate(next_y):
+            
+            if last_token_id == self.end_token_id:
+                
+                # skip scoring
+                self.finished.append((self.current_scores[beam_index], len(self.next_ys) - 1, beam_index))
+
+        if next_y[0] == self.end_token_id:
+            self.top_sentence_ended = True
+
+    def get_current_state(self):
+        "Get the outputs for the current timestep."
+        return torch.stack(self.next_ys, dim=1)
+
+    def get_current_origin(self):
+        "Get the backpointers for the current timestep."
+        return self.prev_ks[-1]
+
+    def done(self):
+        return self.top_sentence_ended and len(self.finished) >= self.n_top
+
+    def get_hypothesis(self, timestep, k):
+        hypothesis = []
+        for j in range(len(self.prev_ks[:timestep]) - 1, -1, -1):
+            hypothesis.append(self.next_ys[j + 1][k])
+            # for RNN, [:, k, :], and for trnasformer, [k, :, :]
+            k = self.prev_ks[j][k]
+
+        return hypothesis[::-1]
+
+    def sort_finished(self, minimum=None):
+        if minimum is not None:
+            i = 0
+            # Add from beam until we have minimum outputs.
+            while len(self.finished) < minimum:
+                # global_scores = self.global_scorer.score(self, self.scores)
+                # s = global_scores[i]
+                s = self.current_scores[i]
+                self.finished.append((s, len(self.next_ys) - 1, i))
+                i += 1
+
+        self.finished = sorted(self.finished, key=lambda a: a[0], reverse=True)
+        scores = [sc for sc, _, _ in self.finished]
+        ks = [(t, k) for _, t, k in self.finished]
+        return scores, ks

vietocr/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/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/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/vietocr/model/trainer.py

@@ -0,0 +1,363 @@
+from vietocr.vietocr.optim.optim import ScheduledOptim
+from vietocr.vietocr.optim.labelsmoothingloss import LabelSmoothingLoss
+from torch.optim import Adam, SGD, AdamW
+from torch import nn
+from vietocr.vietocr.tool.translate import build_model
+from vietocr.vietocr.tool.translate import translate, batch_translate_beam_search
+from vietocr.vietocr.tool.utils import download_weights
+from vietocr.vietocr.tool.logger import Logger
+from vietocr.vietocr.loader.aug import ImgAugTransform
+
+import yaml
+import torch
+from vietocr.vietocr.loader.dataloader_v1 import DataGen
+from vietocr.vietocr.loader.dataloader import OCRDataset, ClusterRandomSampler, Collator
+from torch.utils.data import DataLoader
+from einops import rearrange
+from torch.optim.lr_scheduler import CosineAnnealingLR, CyclicLR, OneCycleLR
+
+import torchvision 
+
+from vietocr.vietocr.tool.utils import compute_accuracy
+from PIL import Image
+import numpy as np
+import os
+import matplotlib.pyplot as plt
+import time
+
+class Trainer():
+    def __init__(self, config, pretrained=True, augmentor=ImgAugTransform()):
+
+        self.config = config
+        self.model, self.vocab = build_model(config)
+        
+        self.device = config['device']
+        self.num_iters = config['trainer']['iters']
+        self.beamsearch = config['predictor']['beamsearch']
+
+        self.data_root = config['dataset']['data_root']
+        self.train_annotation = config['dataset']['train_annotation']
+        self.valid_annotation = config['dataset']['valid_annotation']
+        self.dataset_name = config['dataset']['name']
+
+        self.batch_size = config['trainer']['batch_size']
+        self.print_every = config['trainer']['print_every']
+        self.valid_every = config['trainer']['valid_every']
+        
+        self.image_aug = config['aug']['image_aug']
+        self.masked_language_model = config['aug']['masked_language_model']
+
+        self.checkpoint = config['trainer']['checkpoint']
+        self.export_weights = config['trainer']['export']
+        self.metrics = config['trainer']['metrics']
+        logger = config['trainer']['log']
+    
+        if logger:
+            self.logger = Logger(logger) 
+
+        if pretrained:
+            weight_file = download_weights(config['pretrain'], quiet=config['quiet'])
+            self.load_weights(weight_file)
+
+        self.iter = 0
+        
+        self.optimizer = AdamW(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09)
+        self.scheduler = OneCycleLR(self.optimizer, total_steps=self.num_iters, **config['optimizer'])
+#        self.optimizer = ScheduledOptim(
+#            Adam(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09),
+#            #config['transformer']['d_model'], 
+#            512,
+#            **config['optimizer'])
+
+        self.criterion = LabelSmoothingLoss(len(self.vocab), padding_idx=self.vocab.pad, smoothing=0.1)
+        
+        transforms = None
+        if self.image_aug:
+            transforms =  augmentor
+
+        self.train_gen = self.data_gen('train_{}'.format(self.dataset_name), 
+                self.data_root, self.train_annotation, self.masked_language_model, transform=transforms)
+        if self.valid_annotation:
+            self.valid_gen = self.data_gen('valid_{}'.format(self.dataset_name), 
+                    self.data_root, self.valid_annotation, masked_language_model=False)
+
+        self.train_losses = []
+        
+    def train(self):
+        total_loss = 0
+        
+        total_loader_time = 0
+        total_gpu_time = 0
+        best_acc = 0
+
+        data_iter = iter(self.train_gen)
+        for i in range(self.num_iters):
+            self.iter += 1
+
+            start = time.time()
+
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(self.train_gen)
+                batch = next(data_iter)
+
+            total_loader_time += time.time() - start
+
+            start = time.time()
+            loss = self.step(batch)
+            total_gpu_time += time.time() - start
+
+            total_loss += loss
+            self.train_losses.append((self.iter, loss))
+
+            if self.iter % self.print_every == 0:
+                info = 'iter: {:06d} - train loss: {:.3f} - lr: {:.2e} - load time: {:.2f} - gpu time: {:.2f}'.format(self.iter, 
+                        total_loss/self.print_every, self.optimizer.param_groups[0]['lr'], 
+                        total_loader_time, total_gpu_time)
+
+                total_loss = 0
+                total_loader_time = 0
+                total_gpu_time = 0
+                print(info) 
+                self.logger.log(info)
+
+            if self.valid_annotation and self.iter % self.valid_every == 0:
+                val_loss = self.validate()
+                acc_full_seq, acc_per_char = self.precision(self.metrics)
+
+                info = 'iter: {:06d} - valid loss: {:.3f} - acc full seq: {:.4f} - acc per char: {:.4f}'.format(self.iter, val_loss, acc_full_seq, acc_per_char)
+                print(info)
+                self.logger.log(info)
+
+                if acc_full_seq > best_acc:
+                    self.save_weights(self.export_weights)
+                    best_acc = acc_full_seq
+
+            
+    def validate(self):
+        self.model.eval()
+
+        total_loss = []
+        
+        with torch.no_grad():
+            for step, batch in enumerate(self.valid_gen):
+                batch = self.batch_to_device(batch)
+                img, tgt_input, tgt_output, tgt_padding_mask = batch['img'], batch['tgt_input'], batch['tgt_output'], batch['tgt_padding_mask']
+
+                outputs = self.model(img, tgt_input, tgt_padding_mask)
+#                loss = self.criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_output, 'b o -> (b o)'))
+               
+                outputs = outputs.flatten(0,1)
+                tgt_output = tgt_output.flatten()
+                loss = self.criterion(outputs, tgt_output)
+
+                total_loss.append(loss.item())
+                
+                del outputs
+                del loss
+
+        total_loss = np.mean(total_loss)
+        self.model.train()
+        
+        return total_loss
+    
+    def predict(self, sample=None):
+        pred_sents = []
+        actual_sents = []
+        img_files = []
+
+        for batch in  self.valid_gen:
+            batch = self.batch_to_device(batch)
+
+            if self.beamsearch:
+                translated_sentence = batch_translate_beam_search(batch['img'], self.model)
+                prob = None
+            else:
+                translated_sentence, prob = translate(batch['img'], self.model)
+
+            pred_sent = self.vocab.batch_decode(translated_sentence.tolist())
+            actual_sent = self.vocab.batch_decode(batch['tgt_output'].tolist())
+
+            img_files.extend(batch['filenames'])
+
+            pred_sents.extend(pred_sent)
+            actual_sents.extend(actual_sent)
+            
+            if sample != None and len(pred_sents) > sample:
+                break
+
+        return pred_sents, actual_sents, img_files, prob
+
+    def precision(self, sample=None):
+
+        pred_sents, actual_sents, _, _ = self.predict(sample=sample)
+
+        acc_full_seq = compute_accuracy(actual_sents, pred_sents, mode='full_sequence')
+        acc_per_char = compute_accuracy(actual_sents, pred_sents, mode='per_char')
+    
+        return acc_full_seq, acc_per_char
+    
+    def visualize_prediction(self, sample=16, errorcase=False, fontname='serif', fontsize=16):
+        
+        pred_sents, actual_sents, img_files, probs = self.predict(sample)
+
+        if errorcase:
+            wrongs = []
+            for i in range(len(img_files)):
+                if pred_sents[i]!= actual_sents[i]:
+                    wrongs.append(i)
+
+            pred_sents = [pred_sents[i] for i in wrongs]
+            actual_sents = [actual_sents[i] for i in wrongs]
+            img_files = [img_files[i] for i in wrongs]
+            probs = [probs[i] for i in wrongs]
+
+        img_files = img_files[:sample]
+
+        fontdict = {
+                'family':fontname,
+                'size':fontsize
+                } 
+
+        for vis_idx in range(0, len(img_files)):
+            img_path = img_files[vis_idx]
+            pred_sent = pred_sents[vis_idx]
+            actual_sent = actual_sents[vis_idx]
+            prob = probs[vis_idx]
+
+            img = Image.open(open(img_path, 'rb'))
+            plt.figure()
+            plt.imshow(img)
+            plt.title('prob: {:.3f} - pred: {} - actual: {}'.format(prob, pred_sent, actual_sent), loc='left', fontdict=fontdict)
+            plt.axis('off')
+
+        plt.show()
+    
+    def visualize_dataset(self, sample=16, fontname='serif'):
+        n = 0
+        for batch in self.train_gen:
+            for i in range(self.batch_size):
+                img = batch['img'][i].numpy().transpose(1,2,0)
+                sent = self.vocab.decode(batch['tgt_input'].T[i].tolist())
+                
+                plt.figure()
+                plt.title('sent: {}'.format(sent), loc='center', fontname=fontname)
+                plt.imshow(img)
+                plt.axis('off')
+                
+                n += 1
+                if n >= sample:
+                    plt.show()
+                    return
+
+
+    def load_checkpoint(self, filename):
+        checkpoint = torch.load(filename)
+        
+        optim = ScheduledOptim(
+	       Adam(self.model.parameters(), betas=(0.9, 0.98), eps=1e-09),
+            	self.config['transformer']['d_model'], **self.config['optimizer'])
+
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        self.model.load_state_dict(checkpoint['state_dict'])
+        self.iter = checkpoint['iter']
+
+        self.train_losses = checkpoint['train_losses']
+
+    def save_checkpoint(self, filename):
+        state = {'iter':self.iter, 'state_dict': self.model.state_dict(),
+                'optimizer': self.optimizer.state_dict(), 'train_losses': self.train_losses}
+        
+        path, _ = os.path.split(filename)
+        os.makedirs(path, exist_ok=True)
+
+        torch.save(state, filename)
+
+    def load_weights(self, filename):
+        state_dict = torch.load(filename, map_location=torch.device(self.device))
+        
+        for name, param in self.model.named_parameters():
+            if name not in state_dict:
+                print('{} not found'.format(name))
+            elif state_dict[name].shape != param.shape:
+                print('{} missmatching shape, required {} but found {}'.format(name, param.shape, state_dict[name].shape))
+                del state_dict[name]
+
+        self.model.load_state_dict(state_dict, strict=False)
+
+    def save_weights(self, filename):
+        path, _ = os.path.split(filename)
+        os.makedirs(path, exist_ok=True)
+       
+        torch.save(self.model.state_dict(), filename)
+
+    def batch_to_device(self, batch):
+        img = batch['img'].to(self.device, non_blocking=True)
+        tgt_input = batch['tgt_input'].to(self.device, non_blocking=True)
+        tgt_output = batch['tgt_output'].to(self.device, non_blocking=True)
+        tgt_padding_mask = batch['tgt_padding_mask'].to(self.device, non_blocking=True)
+
+        batch = {
+                'img': img, 'tgt_input':tgt_input, 
+                'tgt_output':tgt_output, 'tgt_padding_mask':tgt_padding_mask, 
+                'filenames': batch['filenames']
+                }
+
+        return batch
+
+    def data_gen(self, lmdb_path, data_root, annotation, masked_language_model=True, transform=None):
+        dataset = OCRDataset(lmdb_path=lmdb_path, 
+                root_dir=data_root, annotation_path=annotation, 
+                vocab=self.vocab, transform=transform, 
+                image_height=self.config['dataset']['image_height'], 
+                image_min_width=self.config['dataset']['image_min_width'], 
+                image_max_width=self.config['dataset']['image_max_width'])
+
+        sampler = ClusterRandomSampler(dataset, self.batch_size, True)
+        collate_fn = Collator(masked_language_model)
+
+        gen = DataLoader(
+                dataset,
+                batch_size=self.batch_size, 
+                sampler=sampler,
+                collate_fn = collate_fn,
+                shuffle=False,
+                drop_last=False,
+                **self.config['dataloader'])
+       
+        return gen
+
+    def data_gen_v1(self, lmdb_path, data_root, annotation):
+        data_gen = DataGen(data_root, annotation, self.vocab, 'cpu', 
+                image_height = self.config['dataset']['image_height'],        
+                image_min_width = self.config['dataset']['image_min_width'],
+                image_max_width = self.config['dataset']['image_max_width'])
+
+        return data_gen
+
+    def step(self, batch):
+        self.model.train()
+
+        batch = self.batch_to_device(batch)
+        img, tgt_input, tgt_output, tgt_padding_mask = batch['img'], batch['tgt_input'], batch['tgt_output'], batch['tgt_padding_mask']    
+        
+        outputs = self.model(img, tgt_input, tgt_key_padding_mask=tgt_padding_mask)
+#        loss = self.criterion(rearrange(outputs, 'b t v -> (b t) v'), rearrange(tgt_output, 'b o -> (b o)'))
+        outputs = outputs.view(-1, outputs.size(2))#flatten(0, 1)
+        tgt_output = tgt_output.view(-1)#flatten()
+        
+        loss = self.criterion(outputs, tgt_output)
+
+        self.optimizer.zero_grad()
+
+        loss.backward()
+        
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1) 
+
+        self.optimizer.step()
+        self.scheduler.step()
+
+        loss_item = loss.item()
+
+        return loss_item

vietocr/vietocr/model/transformerocr.py

@@ -0,0 +1,45 @@
+from vietocr.vietocr.model.backbone.cnn import CNN
+from vietocr.vietocr.model.seqmodel.transformer import LanguageTransformer
+from vietocr.vietocr.model.seqmodel.seq2seq import Seq2Seq
+from vietocr.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/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/vietocr/optim/labelsmoothingloss.py

@@ -0,0 +1,25 @@
+import torch
+from torch import nn
+
+class LabelSmoothingLoss(nn.Module):
+    def __init__(self, classes, padding_idx, smoothing=0.0, dim=-1):
+        super(LabelSmoothingLoss, self).__init__()
+        self.confidence = 1.0 - smoothing
+        self.smoothing = smoothing
+        self.cls = classes
+        self.dim = dim
+        self.padding_idx = padding_idx
+
+    def forward(self, pred, target):
+        pred = pred.log_softmax(dim=self.dim)
+        with torch.no_grad():
+            # true_dist = pred.data.clone()
+            true_dist = torch.zeros_like(pred)
+            true_dist.fill_(self.smoothing / (self.cls - 2))
+            true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
+            true_dist[:, self.padding_idx] = 0
+            mask = torch.nonzero(target.data == self.padding_idx, as_tuple=False)
+            if mask.dim() > 0:
+                true_dist.index_fill_(0, mask.squeeze(), 0.0)
+            
+        return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))

vietocr/vietocr/optim/optim.py

@@ -0,0 +1,58 @@
+class ScheduledOptim():
+    '''A simple wrapper class for learning rate scheduling'''
+
+    def __init__(self, optimizer, d_model, init_lr, n_warmup_steps):
+        assert n_warmup_steps > 0, 'must be greater than 0'
+
+        self._optimizer = optimizer
+        self.init_lr = init_lr
+        self.d_model = d_model
+        self.n_warmup_steps = n_warmup_steps
+        self.n_steps = 0
+
+
+    def step(self):
+        "Step with the inner optimizer"
+        self._update_learning_rate()
+        self._optimizer.step()
+
+
+    def zero_grad(self):
+        "Zero out the gradients with the inner optimizer"
+        self._optimizer.zero_grad()
+
+
+    def _get_lr_scale(self):
+        d_model = self.d_model
+        n_steps, n_warmup_steps = self.n_steps, self.n_warmup_steps
+        return (d_model ** -0.5) * min(n_steps ** (-0.5), n_steps * n_warmup_steps ** (-1.5))
+
+    def state_dict(self):
+        optimizer_state_dict = {
+            'init_lr':self.init_lr,
+            'd_model':self.d_model,
+            'n_warmup_steps':self.n_warmup_steps,
+            'n_steps':self.n_steps,
+            '_optimizer':self._optimizer.state_dict(),
+        }
+        
+        return optimizer_state_dict
+    
+    def load_state_dict(self, state_dict):
+        self.init_lr = state_dict['init_lr']
+        self.d_model = state_dict['d_model']
+        self.n_warmup_steps = state_dict['n_warmup_steps']
+        self.n_steps = state_dict['n_steps']
+        
+        self._optimizer.load_state_dict(state_dict['_optimizer'])
+        
+    def _update_learning_rate(self):
+        ''' Learning rate scheduling per step '''
+
+        self.n_steps += 1
+
+        for param_group in self._optimizer.param_groups:
+            lr = self.init_lr*self._get_lr_scale()
+            self.lr = lr
+
+            param_group['lr'] = lr

vietocr/vietocr/tool/config.py

@@ -0,0 +1,40 @@
+import yaml
+from vietocr.vietocr.tool.utils import download_config
+
+url_config = {
+        'vgg_transformer':'vgg-transformer.yml',
+        'resnet_transformer':'resnet_transformer.yml',
+        'resnet_fpn_transformer':'resnet_fpn_transformer.yml',
+        'vgg_seq2seq':'vgg-seq2seq.yml',
+        'vgg_convseq2seq':'vgg_convseq2seq.yml',
+        'vgg_decoderseq2seq':'vgg_decoderseq2seq.yml',
+        'base':'base.yml',
+        }
+
+class Cfg(dict):
+    def __init__(self, config_dict):
+        super(Cfg, self).__init__(**config_dict)
+        self.__dict__ = self
+
+    @staticmethod
+    def load_config_from_file(fname):
+        #base_config = download_config(url_config['base'])
+        base_config = {}
+        with open(fname, encoding='utf-8') as f:
+            config = yaml.safe_load(f)
+        base_config.update(config)
+
+        return Cfg(base_config)
+
+    @staticmethod
+    def load_config_from_name(name):
+        base_config = download_config(url_config['base'])
+        config = download_config(url_config[name])
+
+        base_config.update(config)
+        return Cfg(base_config)
+
+    def save(self, fname):
+        with open(fname, 'w') as outfile:
+            yaml.dump(dict(self), outfile, default_flow_style=False, allow_unicode=True)
+

vietocr/vietocr/tool/create_dataset.py

@@ -0,0 +1,105 @@
+import sys
+import os
+import lmdb # install lmdb by "pip install lmdb"
+import cv2
+import numpy as np
+from tqdm import tqdm
+
+def checkImageIsValid(imageBin):
+    isvalid = True
+    imgH = None
+    imgW = None
+
+    imageBuf = np.fromstring(imageBin, dtype=np.uint8)
+    try:
+        img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
+
+        imgH, imgW = img.shape[0], img.shape[1]
+        if imgH * imgW == 0:
+            isvalid = False
+    except Exception as e:
+        isvalid = False
+
+    return isvalid, imgH, imgW
+
+def writeCache(env, cache):
+    with env.begin(write=True) as txn:
+        for k, v in cache.items():
+            txn.put(k.encode(), v)
+
+def createDataset(outputPath, root_dir, annotation_path):
+    """
+    Create LMDB dataset for CRNN training.
+    ARGS:
+        outputPath    : LMDB output path
+        imagePathList : list of image path
+        labelList     : list of corresponding groundtruth texts
+        lexiconList   : (optional) list of lexicon lists
+        checkValid    : if true, check the validity of every image
+    """
+
+    annotation_path = os.path.join(root_dir, annotation_path)
+    annotations = []
+    with open(annotation_path, 'r') as ann_file:
+        lines = ann_file.readlines()
+        # for l in lines:
+        #     try:
+        #         annotations.append(l.strip().split('\t'))
+        #     except:
+        #         pass
+        annotations = [l.strip().split('\t') for l in lines]
+
+    nSamples = len(annotations)
+    env = lmdb.open(outputPath, map_size=1099511627776)
+    cache = {}
+    cnt = 0
+    error = 0
+    
+    pbar = tqdm(range(nSamples), ncols = 100, desc='Create {}'.format(outputPath)) 
+    for i in pbar:
+        if len(annotations[i]) >= 2:
+            imageFile, label = annotations[i]
+        else:
+            print("Error: Not enough values to unpack")
+            #sys.exit()
+
+        #imageFile, label = annotations[i]
+        imagePath = os.path.join(root_dir, imageFile)
+
+        if not os.path.exists(imagePath):
+            error += 1
+            continue
+        
+        with open(imagePath, 'rb') as f:
+            imageBin = f.read()
+        isvalid, imgH, imgW = checkImageIsValid(imageBin)
+
+        if not isvalid:
+            error += 1
+            continue
+
+        imageKey = 'image-%09d' % cnt
+        labelKey = 'label-%09d' % cnt
+        pathKey = 'path-%09d' % cnt
+        dimKey = 'dim-%09d' % cnt
+
+        cache[imageKey] = imageBin
+        cache[labelKey] = label.encode()
+        cache[pathKey] = imageFile.encode()
+        cache[dimKey] = np.array([imgH, imgW], dtype=np.int32).tobytes()
+
+        cnt += 1
+
+        if cnt % 1000 == 0:
+            writeCache(env, cache)
+            cache = {}
+
+    nSamples = cnt-1
+    cache['num-samples'] = str(nSamples).encode()
+    writeCache(env, cache)
+
+    if error > 0:
+        print('Remove {} invalid images'.format(error))
+    print('Created dataset with %d samples' % nSamples)
+    sys.stdout.flush()
+

vietocr/vietocr/tool/logger.py

@@ -0,0 +1,17 @@
+import os
+
+
+class Logger():
+    def __init__(self, fname):
+        path, _ = os.path.split(fname)
+        os.makedirs(path, exist_ok=True)
+
+        self.logger = open(fname, 'w')
+
+    def log(self, string):
+        self.logger.write(string+'\n')
+        self.logger.flush()
+
+    def close(self):
+        self.logger.close()
+

vietocr/vietocr/tool/predictor.py

@@ -0,0 +1,85 @@
+from vietocr.vietocr.tool.translate import build_model, translate, translate_beam_search, process_input, predict, batch_translate_beam_search
+from vietocr.vietocr.tool.utils import download_weights
+
+import torch
+from collections import defaultdict
+
+class Predictor():
+    def __init__(self, config):
+
+        device = config['device']
+        
+        model, vocab = build_model(config)
+        weights = '/tmp/weights.pth'
+
+        if config['weights'].startswith('http'):
+            weights = download_weights(config['weights'])
+        else:
+            weights = config['weights']
+
+        model.load_state_dict(torch.load(weights, map_location=torch.device(device)))
+
+        self.config = config
+        self.model = model
+        self.vocab = vocab
+        self.device = device
+
+    def predict(self, img, return_prob=False):
+        img = process_input(img, self.config['dataset']['image_height'], 
+                self.config['dataset']['image_min_width'], self.config['dataset']['image_max_width'])        
+        img = img.to(self.config['device'])
+
+        if self.config['predictor']['beamsearch']:
+            sent = translate_beam_search(img, self.model)
+            s = sent
+            prob = None
+        else:
+            s, prob = translate(img, self.model)
+            s = s[0].tolist()
+            prob = prob[0]
+
+        s = self.vocab.decode(s)
+        
+        if return_prob:
+            return s, prob
+        else:
+            return s
+
+    def predict_batch(self, imgs, return_prob=False):
+        bucket = defaultdict(list)
+        bucket_idx = defaultdict(list)
+        bucket_pred = {}
+        
+        sents, probs = [0]*len(imgs), [0]*len(imgs)
+
+        for i, img in enumerate(imgs):
+            img = process_input(img, self.config['dataset']['image_height'], 
+                self.config['dataset']['image_min_width'], self.config['dataset']['image_max_width'])        
+        
+            bucket[img.shape[-1]].append(img)
+            bucket_idx[img.shape[-1]].append(i)
+
+
+        for k, batch in bucket.items():
+            batch = torch.cat(batch, 0).to(self.device)
+            s, prob = translate(batch, self.model)
+            prob = prob.tolist()
+
+            s = s.tolist()
+            s = self.vocab.batch_decode(s)
+
+            bucket_pred[k] = (s, prob)
+
+
+        for k in bucket_pred:
+            idx = bucket_idx[k]
+            sent, prob = bucket_pred[k]
+            for i, j in enumerate(idx):
+                sents[j] = sent[i]
+                probs[j] = prob[i]
+   
+        if return_prob: 
+            return sents, probs
+        else: 
+            return sents
+

vietocr/vietocr/tool/translate.py

@@ -0,0 +1,172 @@
+import torch
+import numpy as np
+import math
+from PIL import Image
+from torch.nn.functional import log_softmax, softmax
+
+from vietocr.vietocr.model.transformerocr import VietOCR
+from vietocr.vietocr.model.vocab import Vocab
+from vietocr.vietocr.model.beam import Beam
+
+def batch_translate_beam_search(img, model, beam_size=4, candidates=1, max_seq_length=128, sos_token=1, eos_token=2):
+    # img: NxCxHxW
+    model.eval()
+    device = img.device
+    sents = []
+
+    with torch.no_grad():
+        src = model.cnn(img)
+        print(src.shap)
+        memories = model.transformer.forward_encoder(src)
+        for i in range(src.size(0)):
+#            memory = memories[:,i,:].repeat(1, beam_size, 1) # TxNxE
+            memory = model.transformer.get_memory(memories, i)
+            sent = beamsearch(memory, model, device, beam_size, candidates, max_seq_length, sos_token, eos_token)
+            sents.append(sent)
+
+    sents = np.asarray(sents)
+
+    return sents
+   
+def translate_beam_search(img, model, beam_size=4, candidates=1, max_seq_length=128, sos_token=1, eos_token=2):
+    # img: 1xCxHxW
+    model.eval()
+    device = img.device
+
+    with torch.no_grad():
+        src = model.cnn(img)
+        memory = model.transformer.forward_encoder(src) #TxNxE
+        sent = beamsearch(memory, model, device, beam_size, candidates, max_seq_length, sos_token, eos_token)
+
+    return sent
+        
+def beamsearch(memory, model, device, beam_size=4, candidates=1, max_seq_length=128, sos_token=1, eos_token=2):    
+    # memory: Tx1xE
+    model.eval()
+
+    beam = Beam(beam_size=beam_size, min_length=0, n_top=candidates, ranker=None, start_token_id=sos_token, end_token_id=eos_token)
+
+    with torch.no_grad():
+#        memory = memory.repeat(1, beam_size, 1) # TxNxE
+        memory = model.transformer.expand_memory(memory, beam_size)
+
+        for _ in range(max_seq_length):
+            
+            tgt_inp = beam.get_current_state().transpose(0,1).to(device)  # TxN
+            decoder_outputs, memory = model.transformer.forward_decoder(tgt_inp, memory)
+
+            log_prob = log_softmax(decoder_outputs[:,-1, :].squeeze(0), dim=-1)
+            beam.advance(log_prob.cpu())
+            
+            if beam.done():
+                break
+                
+        scores, ks = beam.sort_finished(minimum=1)
+
+        hypothesises = []
+        for i, (times, k) in enumerate(ks[:candidates]):
+            hypothesis = beam.get_hypothesis(times, k)
+            hypothesises.append(hypothesis)
+    
+    return [1] + [int(i) for i in hypothesises[0][:-1]]
+
+def translate(img, model, max_seq_length=128, sos_token=1, eos_token=2):
+    "data: BxCXHxW"
+    model.eval()
+    device = img.device
+
+    with torch.no_grad():
+        src = model.cnn(img)
+        memory = model.transformer.forward_encoder(src)
+
+        translated_sentence = [[sos_token]*len(img)]
+        char_probs = [[1]*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).to(device)
+            
+#            output = model(img, tgt_inp, tgt_key_padding_mask=None)
+#            output = model.transformer(src, tgt_inp, tgt_key_padding_mask=None)
+            output, memory = model.transformer.forward_decoder(tgt_inp, memory)
+            output = softmax(output, dim=-1)
+            output = output.to('cpu')
+
+            values, indices  = torch.topk(output, 5)
+            
+            indices = indices[:, -1, 0]
+            indices = indices.tolist()
+            
+            values = values[:, -1, 0]
+            values = values.tolist()
+            char_probs.append(values)
+
+            translated_sentence.append(indices)   
+            max_length += 1
+
+            del output
+
+        translated_sentence = np.asarray(translated_sentence).T
+        
+        char_probs = np.asarray(char_probs).T
+        char_probs = np.multiply(char_probs, translated_sentence>3)
+        char_probs = np.sum(char_probs, axis=-1)/(char_probs>0).sum(-1)
+    
+    return translated_sentence, char_probs
+
+
+def build_model(config):
+    vocab = Vocab(config['vocab'])
+    device = config['device']
+    
+    model = VietOCR(len(vocab),
+            config['backbone'],
+            config['cnn'], 
+            config['transformer'],
+            config['seq_modeling'])
+    
+    model = model.to(device)
+
+    return model, vocab
+
+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.ANTIALIAS)
+
+    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
+
+def predict(filename, config):
+    img = Image.open(filename)
+    img = process_input(img)
+
+    img = img.to(config['device'])
+
+    model, vocab = build_model(config)
+    s = translate(img, model)[0].tolist()
+    s = vocab.decode(s)
+    
+    return s
+

vietocr/vietocr/tool/utils.py

@@ -0,0 +1,93 @@
+import os
+import gdown
+import yaml
+import numpy as np
+import uuid
+import requests
+import tempfile
+from tqdm import tqdm
+
+def download_weights(uri, cached=None, md5=None, quiet=False):
+    if uri.startswith('http'):
+        return download(url=uri, quiet=quiet)
+    return uri
+
+def download(url, quiet=False):
+    tmp_dir = tempfile.gettempdir()
+    filename = url.split('/')[-1]
+    full_path = os.path.join(tmp_dir, filename)
+    
+    if os.path.exists(full_path):
+        print('Model weight {} exsits. Ignore download!'.format(full_path))
+        return full_path
+
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(full_path, 'wb') as f:
+            for chunk in tqdm(r.iter_content(chunk_size=8192)):
+                # If you have chunk encoded response uncomment if
+                # and set chunk_size parameter to None.
+                #if chunk:
+                f.write(chunk)
+    return full_path
+
+def download_config(id):
+    url = 'https://vocr.vn/data/vietocr/config/{}'.format(id)
+    r = requests.get(url)
+    config = yaml.safe_load(r.text)
+    return config
+
+def compute_accuracy(ground_truth, predictions, mode='full_sequence'):
+    """
+    Computes accuracy
+    :param ground_truth:
+    :param predictions:
+    :param display: Whether to print values to stdout
+    :param mode: if 'per_char' is selected then
+                 single_label_accuracy = correct_predicted_char_nums_of_single_sample / single_label_char_nums
+                 avg_label_accuracy = sum(single_label_accuracy) / label_nums
+                 if 'full_sequence' is selected then
+                 single_label_accuracy = 1 if the prediction result is exactly the same as label else 0
+                 avg_label_accuracy = sum(single_label_accuracy) / label_nums
+    :return: avg_label_accuracy
+    """
+    if mode == 'per_char':
+
+        accuracy = []
+
+        for index, label in enumerate(ground_truth):
+            prediction = predictions[index]
+            total_count = len(label)
+            correct_count = 0
+            try:
+                for i, tmp in enumerate(label):
+                    if tmp == prediction[i]:
+                        correct_count += 1
+            except IndexError:
+                continue
+            finally:
+                try:
+                    accuracy.append(correct_count / total_count)
+                except ZeroDivisionError:
+                    if len(prediction) == 0:
+                        accuracy.append(1)
+                    else:
+                        accuracy.append(0)
+        avg_accuracy = np.mean(np.array(accuracy).astype(np.float32), axis=0)
+    elif mode == 'full_sequence':
+        try:
+            correct_count = 0
+            for index, label in enumerate(ground_truth):
+                prediction = predictions[index]
+                if prediction == label:
+                    correct_count += 1
+            avg_accuracy = correct_count / len(ground_truth)
+        except ZeroDivisionError:
+            if not predictions:
+                avg_accuracy = 1
+            else:
+                avg_accuracy = 0
+    else:
+        raise NotImplementedError('Other accuracy compute mode has not been implemented')
+
+    return avg_accuracy