gfgf

FAST-ABINet-OCR/.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

FAST-ABINet-OCR/README.md

@@ -0,0 +1,12 @@
+---
+title: FAST ABINet OCR
+emoji: 🌍
+colorFrom: yellow
+colorTo: pink
+sdk: gradio
+sdk_version: 3.2
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

LICENSE

@@ -0,0 +1,25 @@
+ABINet for non-commercial purposes
+
+Copyright (c) 2021, USTC
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,12 +1,13 @@
 ---
-title: FAST ABINet OCR
-emoji: 🌍
-colorFrom: yellow
-colorTo: pink
+title: ABINet OCR
+emoji: 🏃
+colorFrom: indigo
+colorTo: red
 sdk: gradio
-sdk_version: 3.2
+sdk_version: 2.8.12
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,36 @@
+
+import os
+os.system('pip install --upgrade gdown')
+import gdown
+gdown.download(id='1mYM_26qHUom_5NU7iutHneB_KHlLjL5y', output='workdir.zip')
+os.system('unzip workdir.zip')
+
+import glob
+import gradio as gr
+from demo import get_model, preprocess, postprocess, load
+from utils import Config, Logger, CharsetMapper
+
+config = Config('configs/train_abinet.yaml')
+config.model_vision_checkpoint = None
+model = get_model(config)
+model = load(model, 'workdir/train-abinet/best-train-abinet.pth')
+charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
+
+def process_image(image):
+  img = image.convert('RGB')
+  img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
+  res = model(img)
+  return postprocess(res, charset, 'alignment')[0][0]
+
+title = "Interactive demo: ABINet"
+description = "Demo for ABINet, ABINet uses a vision model and an explicit language model to recognize text in the wild, which are trained in end-to-end way. The language model (BCN) achieves bidirectional language representation in simulating cloze test, additionally utilizing iterative correction strategy. To use it, simply upload a (single-text line) image or use one of the example images below and click 'submit'. Results will show up in a few seconds."
+article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2103.06495.pdf'>Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition</a> | <a href='https://github.com/FangShancheng/ABINet'>Github Repo</a></p>"
+
+iface = gr.Interface(fn=process_image, 
+                     inputs=gr.inputs.Image(type="pil"), 
+                     outputs=gr.outputs.Textbox(),
+                     title=title,
+                     description=description,
+                     article=article,
+                     examples=glob.glob('figs/test/*.png'))
+iface.launch(debug=True, share=True)

callbacks.py

@@ -0,0 +1,360 @@
+import logging
+import shutil
+import time
+
+import editdistance as ed
+import torchvision.utils as vutils
+from fastai.callbacks.tensorboard import (LearnerTensorboardWriter,
+                                          SummaryWriter, TBWriteRequest,
+                                          asyncTBWriter)
+from fastai.vision import *
+from torch.nn.parallel import DistributedDataParallel
+from torchvision import transforms
+
+import dataset
+from utils import CharsetMapper, Timer, blend_mask
+
+
+class IterationCallback(LearnerTensorboardWriter):
+    "A `TrackerCallback` that monitor in each iteration."
+    def __init__(self, learn:Learner, name:str='model', checpoint_keep_num=5,
+                 show_iters:int=50, eval_iters:int=1000, save_iters:int=20000,
+                 start_iters:int=0, stats_iters=20000):
+        #if self.learn.rank is not None: time.sleep(self.learn.rank)  # keep all event files
+        super().__init__(learn, base_dir='.', name=learn.path, loss_iters=show_iters, 
+                        stats_iters=stats_iters, hist_iters=stats_iters)
+        self.name, self.bestname = Path(name).name, f'best-{Path(name).name}'
+        self.show_iters = show_iters
+        self.eval_iters = eval_iters
+        self.save_iters = save_iters
+        self.start_iters = start_iters
+        self.checpoint_keep_num = checpoint_keep_num
+        self.metrics_root = 'metrics/'  # rewrite
+        self.timer = Timer()
+        self.host = self.learn.rank is None or self.learn.rank == 0
+
+    def _write_metrics(self, iteration:int, names:List[str], last_metrics:MetricsList)->None:
+        "Writes training metrics to Tensorboard."
+        for i, name in enumerate(names):
+            if last_metrics is None or len(last_metrics) < i+1: return
+            scalar_value = last_metrics[i]
+            self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)
+
+    def _write_sub_loss(self, iteration:int, last_losses:dict)->None:
+        "Writes sub loss to Tensorboard."
+        for name, loss in last_losses.items():
+            scalar_value = to_np(loss)
+            tag = self.metrics_root + name
+            self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    def _save(self, name):
+        if isinstance(self.learn.model, DistributedDataParallel):
+            tmp = self.learn.model
+            self.learn.model = self.learn.model.module
+            self.learn.save(name)
+            self.learn.model = tmp
+        else: self.learn.save(name)
+
+    def _validate(self, dl=None, callbacks=None, metrics=None, keeped_items=False):
+        "Validate on `dl` with potential `callbacks` and `metrics`."
+        dl = ifnone(dl, self.learn.data.valid_dl)
+        metrics = ifnone(metrics, self.learn.metrics)
+        cb_handler = CallbackHandler(ifnone(callbacks, []), metrics)
+        cb_handler.on_train_begin(1, None, metrics); cb_handler.on_epoch_begin()
+        if keeped_items: cb_handler.state_dict.update(dict(keeped_items=[]))
+        val_metrics = validate(self.learn.model, dl, self.loss_func, cb_handler)
+        cb_handler.on_epoch_end(val_metrics)
+        if keeped_items: return cb_handler.state_dict['keeped_items']
+        else: return cb_handler.state_dict['last_metrics']
+
+    def jump_to_epoch_iter(self, epoch:int, iteration:int)->None:
+        try:
+            self.learn.load(f'{self.name}_{epoch}_{iteration}', purge=False)
+            logging.info(f'Loaded {self.name}_{epoch}_{iteration}')
+        except: logging.info(f'Model {self.name}_{epoch}_{iteration} not found.')
+
+    def on_train_begin(self, n_epochs, **kwargs):
+        # TODO: can not write graph here
+        # super().on_train_begin(**kwargs)
+        self.best = -float('inf')
+        self.timer.tic()
+        if self.host:
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if checkpoint_path.exists():
+                os.remove(checkpoint_path)
+            open(checkpoint_path, 'w').close()
+        return {'skip_validate': True, 'iteration':self.start_iters}  # disable default validate
+
+    def on_batch_begin(self, **kwargs:Any)->None:
+        self.timer.toc_data()
+        super().on_batch_begin(**kwargs)
+
+    def on_batch_end(self, iteration, epoch, last_loss, smooth_loss, train, **kwargs):
+        super().on_batch_end(last_loss, iteration, train, **kwargs)
+        if iteration == 0: return
+
+        if iteration % self.loss_iters == 0:
+            last_losses = self.learn.loss_func.last_losses
+            self._write_sub_loss(iteration=iteration, last_losses=last_losses)
+            self.tbwriter.add_scalar(tag=self.metrics_root + 'lr',
+                scalar_value=self.opt.lr, global_step=iteration)
+
+        if iteration % self.show_iters == 0:
+            log_str = f'epoch {epoch} iter {iteration}: loss = {last_loss:6.4f},  ' \
+                      f'smooth loss = {smooth_loss:6.4f}'
+            logging.info(log_str)
+            # log_str = f'data time = {self.timer.data_diff:.4f}s, runing time = {self.timer.running_diff:.4f}s'
+            # logging.info(log_str)
+
+        if iteration % self.eval_iters == 0:
+            # TODO: or remove time to on_epoch_end
+            # 1. Record time 
+            log_str = f'average data time = {self.timer.average_data_time():.4f}s, ' \
+                      f'average running time = {self.timer.average_running_time():.4f}s'
+            logging.info(log_str)
+
+            # 2. Call validate
+            last_metrics = self._validate()
+            self.learn.model.train()
+            log_str = f'epoch {epoch} iter {iteration}: eval loss = {last_metrics[0]:6.4f},  ' \
+                      f'ccr = {last_metrics[1]:6.4f},  cwr = {last_metrics[2]:6.4f},  ' \
+                      f'ted = {last_metrics[3]:6.4f},  ned = {last_metrics[4]:6.4f},  ' \
+                      f'ted/w = {last_metrics[5]:6.4f}, '
+            logging.info(log_str)
+            names = ['eval_loss', 'ccr', 'cwr', 'ted', 'ned', 'ted/w']
+            self._write_metrics(iteration, names, last_metrics)
+
+            # 3. Save best model
+            current = last_metrics[2]
+            if current is not None and current > self.best:
+                logging.info(f'Better model found at epoch {epoch}, '\
+                             f'iter {iteration} with accuracy value: {current:6.4f}.')
+                self.best = current
+                self._save(f'{self.bestname}')
+
+        if iteration % self.save_iters == 0 and self.host:
+            logging.info(f'Save model {self.name}_{epoch}_{iteration}')
+            filename = f'{self.name}_{epoch}_{iteration}'
+            self._save(filename)
+
+            checkpoint_path = self.learn.path/'checkpoint.yaml'
+            if not checkpoint_path.exists():
+                open(checkpoint_path, 'w').close()
+            with open(checkpoint_path, 'r') as file:
+                checkpoints = yaml.load(file, Loader=yaml.FullLoader) or dict()
+            checkpoints['all_checkpoints'] = (
+                checkpoints.get('all_checkpoints') or list())
+            checkpoints['all_checkpoints'].insert(0, filename)
+            if len(checkpoints['all_checkpoints']) > self.checpoint_keep_num:
+                removed_checkpoint = checkpoints['all_checkpoints'].pop()
+                removed_checkpoint =  self.learn.path/self.learn.model_dir/f'{removed_checkpoint}.pth'
+                os.remove(removed_checkpoint)
+            checkpoints['current_checkpoint'] = filename
+            with open(checkpoint_path, 'w') as file:
+                yaml.dump(checkpoints, file)
+
+
+        self.timer.toc_running()
+
+    def on_train_end(self, **kwargs):
+        #self.learn.load(f'{self.bestname}', purge=False)
+        pass
+    
+    def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:
+        self._write_embedding(iteration=iteration)
+
+
+class TextAccuracy(Callback):
+    _names = ['ccr', 'cwr', 'ted', 'ned', 'ted/w']
+    def __init__(self, charset_path, max_length, case_sensitive, model_eval):
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+
+        self.model_eval = model_eval or 'alignment'
+        assert self.model_eval in ['vision', 'language', 'alignment']
+ 
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+        self.total_ed = 0.
+        self.total_ned = 0.
+
+    def _get_output(self, last_output):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == self.model_eval: output = res
+        else: output = last_output
+        return output
+    
+    def _update_output(self, last_output, items):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == self.model_eval: res.update(items)
+        else: last_output.update(items)
+        return last_output
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        output = self._get_output(last_output)
+        logits, pt_lengths = output['logits'], output['pt_lengths']
+        pt_text, pt_scores, pt_lengths_ = self.decode(logits)
+        assert (pt_lengths == pt_lengths_).all(), f'{pt_lengths} != {pt_lengths_} for {pt_text}'
+        last_output = self._update_output(last_output, {'pt_text':pt_text, 'pt_scores':pt_scores})
+
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        label = last_target[0]
+        if label.dim() == 3: label = label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in label]
+        
+        for i in range(len(gt_text)):
+            if not self.case_sensitive:
+                gt_text[i], pt_text[i] = gt_text[i].lower(), pt_text[i].lower()
+            distance = ed.eval(gt_text[i], pt_text[i])
+            self.total_ed += distance
+            self.total_ned += float(distance) / max(len(gt_text[i]), 1)
+
+            if gt_text[i] == pt_text[i]:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+            
+            for j in range(min(len(gt_text[i]), len(pt_text[i]))):
+                if gt_text[i][j] == pt_text[i][j]:
+                    self.correct_num_char += 1
+            self.total_num_char += len(gt_text[i])
+
+        return {'last_output': last_output}
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                self.total_ed,
+                self.total_ned,
+                self.total_ed / self.total_num_word]
+        return add_metrics(last_metrics, mets)
+
+    def decode(self, logit):
+        """ Greed decode """
+        # TODO: test running time and decode on GPU
+        out = F.softmax(logit, dim=2)
+        pt_text, pt_scores, pt_lengths = [], [], []
+        for o in out:
+            text = self.charset.get_text(o.argmax(dim=1), padding=False, trim=False)
+            text = text.split(self.charset.null_char)[0]  # end at end-token
+            pt_text.append(text)
+            pt_scores.append(o.max(dim=1)[0])
+            pt_lengths.append(min(len(text) + 1, self.max_length))  # one for end-token
+        pt_scores = torch.stack(pt_scores)
+        pt_lengths = pt_scores.new_tensor(pt_lengths, dtype=torch.long)
+        return pt_text, pt_scores, pt_lengths
+
+
+class TopKTextAccuracy(TextAccuracy):
+    _names = ['ccr', 'cwr']
+    def __init__(self, k, charset_path, max_length, case_sensitive, model_eval):
+        self.k = k
+        self.charset_path = charset_path
+        self.max_length = max_length
+        self.case_sensitive = case_sensitive
+        self.charset = CharsetMapper(charset_path, self.max_length)
+        self.names = self._names
+ 
+    def on_epoch_begin(self, **kwargs):
+        self.total_num_char = 0.
+        self.total_num_word = 0.
+        self.correct_num_char = 0.
+        self.correct_num_word = 0.
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        logits, pt_lengths = last_output['logits'], last_output['pt_lengths']
+        gt_labels, gt_lengths = last_target[:]
+
+        for logit, pt_length, label, length in zip(logits, pt_lengths, gt_labels, gt_lengths):
+            word_flag = True
+            for i in range(length):
+                char_logit = logit[i].topk(self.k)[1]
+                char_label = label[i].argmax(-1)
+                if char_label in char_logit: self.correct_num_char += 1
+                else: word_flag = False
+                self.total_num_char += 1
+            if pt_length == length and word_flag:
+                self.correct_num_word += 1
+            self.total_num_word += 1
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        mets = [self.correct_num_char / self.total_num_char,
+                self.correct_num_word / self.total_num_word,
+                0., 0., 0.]
+        return add_metrics(last_metrics, mets)
+
+
+class DumpPrediction(LearnerCallback):
+
+    def __init__(self, learn, dataset, charset_path, model_eval, image_only=False, debug=False):
+        super().__init__(learn=learn)
+        self.debug = debug
+        self.model_eval = model_eval or 'alignment'
+        self.image_only = image_only
+        assert self.model_eval in ['vision', 'language', 'alignment']
+
+        self.dataset, self.root = dataset, Path(self.learn.path)/f'{dataset}-{self.model_eval}'
+        self.attn_root = self.root/'attn'
+        self.charset = CharsetMapper(charset_path)
+        if self.root.exists(): shutil.rmtree(self.root)
+        self.root.mkdir(), self.attn_root.mkdir()
+
+        self.pil = transforms.ToPILImage()
+        self.tensor = transforms.ToTensor()
+        size = self.learn.data.img_h, self.learn.data.img_w
+        self.resize = transforms.Resize(size=size, interpolation=0)
+        self.c = 0
+
+    def on_batch_end(self, last_input, last_output, last_target, **kwargs):
+        if isinstance(last_output, (tuple, list)):
+            for res in last_output:
+                if res['name'] == self.model_eval: pt_text = res['pt_text']
+                if res['name'] == 'vision': attn_scores = res['attn_scores'].detach().cpu()
+                if res['name'] == self.model_eval: logits = res['logits']
+        else:
+            pt_text = last_output['pt_text']
+            attn_scores = last_output['attn_scores'].detach().cpu()
+            logits = last_output['logits']
+
+        images = last_input[0] if isinstance(last_input, (tuple, list)) else last_input
+        images = images.detach().cpu()
+        pt_text = [self.charset.trim(t) for t in pt_text]
+        gt_label = last_target[0]
+        if gt_label.dim() == 3: gt_label = gt_label.argmax(dim=-1)  # one-hot label
+        gt_text = [self.charset.get_text(l, trim=True) for l in gt_label]
+
+        prediction, false_prediction = [], []
+        for gt, pt, image, attn, logit in zip(gt_text, pt_text, images, attn_scores, logits):
+            prediction.append(f'{gt}\t{pt}\n')
+            if gt != pt:
+                if self.debug:
+                    scores = torch.softmax(logit, dim=-1)[:max(len(pt), len(gt)) + 1]
+                    logging.info(f'{self.c} gt {gt}, pt {pt}, logit {logit.shape}, scores {scores.topk(5, dim=-1)}')
+                false_prediction.append(f'{gt}\t{pt}\n')
+
+            image = self.learn.data.denorm(image)
+            if not self.image_only:
+                image_np = np.array(self.pil(image))
+                attn_pil = [self.pil(a) for a in attn[:, None, :, :]]
+                attn = [self.tensor(self.resize(a)).repeat(3, 1, 1) for a in attn_pil]
+                attn_sum = np.array([np.array(a) for a in attn_pil[:len(pt)]]).sum(axis=0)
+                blended_sum = self.tensor(blend_mask(image_np, attn_sum))
+                blended = [self.tensor(blend_mask(image_np, np.array(a))) for a in attn_pil]
+                save_image = torch.stack([image] + attn + [blended_sum] + blended)
+                save_image = save_image.view(2, -1, *save_image.shape[1:])
+                save_image = save_image.permute(1, 0, 2, 3, 4).flatten(0, 1)
+                vutils.save_image(save_image, self.attn_root/f'{self.c}_{gt}_{pt}.jpg', 
+                                nrow=2, normalize=True, scale_each=True)
+            else:
+                self.pil(image).save(self.attn_root/f'{self.c}_{gt}_{pt}.jpg')
+            self.c += 1
+
+        with open(self.root/f'{self.model_eval}.txt', 'a') as f: f.writelines(prediction)
+        with open(self.root/f'{self.model_eval}-false.txt', 'a') as f: f.writelines(false_prediction)  

configs/pretrain_language_model.yaml

@@ -0,0 +1,45 @@
+global:
+  name: pretrain-language-model
+  phase: train
+  stage: pretrain-language
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/WikiText-103.csv'],
+    batch_size: 4096
+  }
+  test: {
+    roots: ['data/WikiText-103_eval_d1.csv'],
+    batch_size: 4096
+  }
+
+training:
+  epochs: 80
+  show_iters: 50
+  eval_iters: 6000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [70, 10],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_language.BCNLanguage'
+  language: {
+    num_layers: 4,
+    loss_weight: 1.,
+    use_self_attn: False
+  }

configs/pretrain_vision_model.yaml

@@ -0,0 +1,58 @@
+global:
+  name: pretrain-vision-model
+  phase: train
+  stage: pretrain-vision
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 384
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 384
+  }
+  data_aug: True
+  multiscales: False
+  num_workers: 14
+
+training:
+  epochs: 8
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [6, 2],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_vision.BaseVision'
+  checkpoint: ~
+  vision: {
+    loss_weight: 1.,
+    attention: 'position',
+    backbone: 'transformer',
+    backbone_ln: 3,
+  }

configs/pretrain_vision_model_sv.yaml

@@ -0,0 +1,58 @@
+global:
+  name: pretrain-vision-model-sv
+  phase: train
+  stage: pretrain-vision
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 384
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 384
+  }
+  data_aug: True
+  multiscales: False
+  num_workers: 14
+
+training:
+  epochs: 8
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [6, 2],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_vision.BaseVision'
+  checkpoint: ~
+  vision: {
+    loss_weight: 1.,
+    attention: 'attention',
+    backbone: 'transformer',
+    backbone_ln: 2,
+  }

configs/template.yaml

@@ -0,0 +1,67 @@
+global:
+  name: exp
+  phase: train
+  stage: pretrain-vision
+  workdir: /tmp/workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 128
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 128
+  }
+  charset_path: data/charset_36.txt
+  num_workers: 4
+  max_length: 25  # 30
+  image_height: 32
+  image_width: 128
+  case_sensitive: False
+  eval_case_sensitive: False
+  data_aug: True
+  multiscales: False
+  pin_memory: True
+  smooth_label: False
+  smooth_factor: 0.1
+  one_hot_y: True
+  use_sm: False
+
+training:
+  epochs: 6
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 20000
+  start_iters: 0
+  stats_iters: 100000
+
+optimizer:
+  type: Adadelta # Adadelta, Adam
+  true_wd: False
+  wd: 0. # 0.001
+  bn_wd: False
+  args: {
+    # betas: !!python/tuple [0.9, 0.99], # betas=(0.9,0.99) for AdamW
+    # betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  clip_grad: 20
+  lr: [1.0, 1.0, 1.0]  # lr: [0.005, 0.005, 0.005]   
+  scheduler: {
+    periods: [3, 2, 1],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_abinet.ABINetModel'
+  checkpoint: ~
+  strict: True

configs/train_abinet.yaml

@@ -0,0 +1,71 @@
+global:
+  name: train-abinet
+  phase: train
+  stage: train-super
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 384
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 384
+  }
+  data_aug: True
+  multiscales: False
+  num_workers: 14
+
+training:
+  epochs: 10
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [6, 4],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_abinet_iter.ABINetIterModel'
+  iter_size: 3
+  ensemble: ''
+  use_vision: False
+  vision: {
+    checkpoint: workdir/pretrain-vision-model/best-pretrain-vision-model.pth,
+    loss_weight: 1.,
+    attention: 'position',
+    backbone: 'transformer',
+    backbone_ln: 3,
+  }
+  language: {
+    checkpoint:  workdir/pretrain-language-model/pretrain-language-model.pth,
+    num_layers: 4,
+    loss_weight: 1.,
+    detach: True,
+    use_self_attn: False
+  }
+  alignment: {
+    loss_weight: 1.,
+  }

configs/train_abinet_sv.yaml

@@ -0,0 +1,71 @@
+global:
+  name: train-abinet-sv
+  phase: train
+  stage: train-super
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 384
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 384
+  }
+  data_aug: True
+  multiscales: False
+  num_workers: 14
+
+training:
+  epochs: 10
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [6, 4],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_abinet_iter.ABINetIterModel'
+  iter_size: 3
+  ensemble: ''
+  use_vision: False
+  vision: {
+    checkpoint: workdir/pretrain-vision-model-sv/best-pretrain-vision-model-sv.pth,
+    loss_weight: 1.,
+    attention: 'attention',
+    backbone: 'transformer',
+    backbone_ln: 2,
+  }
+  language: {
+    checkpoint:  workdir/pretrain-language-model/pretrain-language-model.pth,
+    num_layers: 4,
+    loss_weight: 1.,
+    detach: True,
+    use_self_attn: False
+  }
+  alignment: {
+    loss_weight: 1.,
+  }

configs/train_abinet_wo_iter.yaml

@@ -0,0 +1,68 @@
+global:
+  name: train-abinet-wo-iter
+  phase: train
+  stage: train-super
+  workdir: workdir
+  seed: ~
+ 
+dataset:
+  train: {
+    roots: ['data/training/MJ/MJ_train/', 
+            'data/training/MJ/MJ_test/', 
+            'data/training/MJ/MJ_valid/', 
+            'data/training/ST'],
+    batch_size: 384
+  }
+  test: {
+    roots: ['data/evaluation/IIIT5k_3000', 
+            'data/evaluation/SVT', 
+            'data/evaluation/SVTP',
+            'data/evaluation/IC13_857',
+            'data/evaluation/IC15_1811',
+            'data/evaluation/CUTE80'],
+    batch_size: 384
+  }
+  data_aug: True
+  multiscales: False
+  num_workers: 14
+
+training:
+  epochs: 10
+  show_iters: 50
+  eval_iters: 3000
+  save_iters: 3000
+
+optimizer:
+  type: Adam
+  true_wd: False
+  wd: 0.0
+  bn_wd: False
+  clip_grad: 20
+  lr: 0.0001
+  args: {
+    betas: !!python/tuple [0.9, 0.999], # for default Adam 
+  }
+  scheduler: {
+    periods: [6, 4],
+    gamma: 0.1,
+  }
+
+model:
+  name: 'modules.model_abinet.ABINetModel'
+  vision: {
+    checkpoint: workdir/pretrain-vision-model/best-pretrain-vision-model.pth,
+    loss_weight: 1.,
+    attention: 'position',
+    backbone: 'transformer',
+    backbone_ln: 3,
+  }
+  language: {
+    checkpoint:  workdir/pretrain-language-model/pretrain-language-model.pth,
+    num_layers: 4,
+    loss_weight: 1.,
+    detach: True,
+    use_self_attn: False
+  }
+  alignment: {
+    loss_weight: 1.,
+  }

data/charset_36.txt

@@ -0,0 +1,36 @@
+0	a
+1	b
+2	c
+3	d
+4	e
+5	f
+6	g
+7	h
+8	i
+9	j
+10	k
+11	l
+12	m
+13	n
+14	o
+15	p
+16	q
+17	r
+18	s
+19	t
+20	u
+21	v
+22	w
+23	x
+24	y
+25	z
+26	1
+27	2
+28	3
+29	4
+30	5
+31	6
+32	7
+33	8
+34	9
+35	0

data/charset_62.txt

@@ -0,0 +1,62 @@
+0	0
+1	1
+2	2
+3	3
+4	4
+5	5
+6	6
+7	7
+8	8
+9	9
+10	A
+11	B
+12	C
+13	D
+14	E
+15	F
+16	G
+17	H
+18	I
+19	J
+20	K
+21	L
+22	M
+23	N
+24	O
+25	P
+26	Q
+27	R
+28	S
+29	T
+30	U
+31	V
+32	W
+33	X
+34	Y
+35	Z
+36	a
+37	b
+38	c
+39	d
+40	e
+41	f
+42	g
+43	h
+44	i
+45	j
+46	k
+47	l
+48	m
+49	n
+50	o
+51	p
+52	q
+53	r
+54	s
+55	t
+56	u
+57	v
+58	w
+59	x
+60	y
+61	z

dataset.py

@@ -0,0 +1,278 @@
+import logging
+import re
+
+import cv2
+import lmdb
+import six
+from fastai.vision import *
+from torchvision import transforms
+
+from transforms import CVColorJitter, CVDeterioration, CVGeometry
+from utils import CharsetMapper, onehot
+
+
+class ImageDataset(Dataset):
+    "`ImageDataset` read data from LMDB database."
+
+    def __init__(self,
+                 path:PathOrStr,
+                 is_training:bool=True,
+                 img_h:int=32,
+                 img_w:int=100,
+                 max_length:int=25,
+                 check_length:bool=True,
+                 case_sensitive:bool=False,
+                 charset_path:str='data/charset_36.txt',
+                 convert_mode:str='RGB',
+                 data_aug:bool=True,
+                 deteriorate_ratio:float=0.,
+                 multiscales:bool=True,
+                 one_hot_y:bool=True,
+                 return_idx:bool=False,
+                 return_raw:bool=False,
+                 **kwargs):
+        self.path, self.name = Path(path), Path(path).name
+        assert self.path.is_dir() and self.path.exists(), f"{path} is not a valid directory."
+        self.convert_mode, self.check_length = convert_mode, check_length
+        self.img_h, self.img_w = img_h, img_w
+        self.max_length, self.one_hot_y = max_length, one_hot_y
+        self.return_idx, self.return_raw = return_idx, return_raw
+        self.case_sensitive, self.is_training = case_sensitive, is_training
+        self.data_aug, self.multiscales = data_aug, multiscales
+        self.charset = CharsetMapper(charset_path, max_length=max_length+1)
+        self.c = self.charset.num_classes
+
+        self.env = lmdb.open(str(path), readonly=True, lock=False, readahead=False, meminit=False)
+        assert self.env, f'Cannot open LMDB dataset from {path}.'
+        with self.env.begin(write=False) as txn:
+            self.length = int(txn.get('num-samples'.encode()))
+
+        if self.is_training and self.data_aug:
+            self.augment_tfs = transforms.Compose([
+                CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
+                CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
+                CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25)
+            ])
+        self.totensor = transforms.ToTensor()
+    
+    def __len__(self): return self.length
+
+    def _next_image(self, index):
+        next_index = random.randint(0, len(self) - 1)
+        return self.get(next_index)
+
+    def _check_image(self, x, pixels=6):
+        if x.size[0] <= pixels or x.size[1] <= pixels: return False
+        else: return True
+
+    def resize_multiscales(self, img, borderType=cv2.BORDER_CONSTANT): 
+        def _resize_ratio(img, ratio, fix_h=True):
+            if ratio * self.img_w < self.img_h:
+                if fix_h: trg_h = self.img_h
+                else: trg_h = int(ratio * self.img_w)
+                trg_w = self.img_w
+            else: trg_h, trg_w = self.img_h, int(self.img_h / ratio)
+            img = cv2.resize(img, (trg_w, trg_h))
+            pad_h, pad_w = (self.img_h - trg_h) / 2, (self.img_w - trg_w) / 2
+            top, bottom = math.ceil(pad_h), math.floor(pad_h)
+            left, right = math.ceil(pad_w), math.floor(pad_w)
+            img = cv2.copyMakeBorder(img, top, bottom, left, right, borderType)
+            return img
+        
+        if self.is_training: 
+            if random.random() < 0.5:
+                base, maxh, maxw = self.img_h, self.img_h, self.img_w
+                h, w = random.randint(base, maxh), random.randint(base, maxw)
+                return _resize_ratio(img, h/w)
+            else: return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio
+        else:  return _resize_ratio(img, img.shape[0] / img.shape[1])  # keep aspect ratio
+
+    def resize(self, img):
+        if self.multiscales: return self.resize_multiscales(img, cv2.BORDER_REPLICATE)
+        else: return cv2.resize(img, (self.img_w, self.img_h))
+         
+    def get(self, idx):
+        with self.env.begin(write=False) as txn:
+            image_key, label_key = f'image-{idx+1:09d}', f'label-{idx+1:09d}'
+            try:
+                label = str(txn.get(label_key.encode()), 'utf-8')  # label
+                label = re.sub('[^0-9a-zA-Z]+', '', label)
+                if self.check_length and self.max_length > 0:
+                    if len(label) > self.max_length or len(label) <= 0:
+                        #logging.info(f'Long or short text image is found: {self.name}, {idx}, {label}, {len(label)}')
+                        return self._next_image(idx)
+                label = label[:self.max_length]
+
+                imgbuf = txn.get(image_key.encode())  # image
+                buf = six.BytesIO()
+                buf.write(imgbuf)
+                buf.seek(0)
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", UserWarning) # EXIF warning from TiffPlugin
+                    image = PIL.Image.open(buf).convert(self.convert_mode)
+                if self.is_training and not self._check_image(image):
+                    #logging.info(f'Invalid image is found: {self.name}, {idx}, {label}, {len(label)}')
+                    return self._next_image(idx)
+            except:
+                import traceback
+                traceback.print_exc()
+                logging.info(f'Corrupted image is found: {self.name}, {idx}, {label}, {len(label)}')
+                return self._next_image(idx)
+            return image, label, idx
+
+    def _process_training(self, image):
+        if self.data_aug: image = self.augment_tfs(image)
+        image = self.resize(np.array(image))
+        return image
+
+    def _process_test(self, image):
+        return self.resize(np.array(image)) # TODO:move is_training to here
+
+    def __getitem__(self, idx):
+        image, text, idx_new = self.get(idx)
+        if not self.is_training: assert idx == idx_new, f'idx {idx} != idx_new {idx_new} during testing.'
+
+        if self.is_training: image = self._process_training(image)
+        else: image = self._process_test(image)
+        if self.return_raw: return image, text
+        image = self.totensor(image)
+
+        length = tensor(len(text) + 1).to(dtype=torch.long)  # one for end token
+        label = self.charset.get_labels(text, case_sensitive=self.case_sensitive)
+        label = tensor(label).to(dtype=torch.long)
+        if self.one_hot_y: label = onehot(label, self.charset.num_classes)
+
+        if self.return_idx: y = [label, length, idx_new]
+        else: y = [label, length]
+        return image, y
+
+
+class TextDataset(Dataset):
+    def __init__(self,
+                 path:PathOrStr, 
+                 delimiter:str='\t',
+                 max_length:int=25, 
+                 charset_path:str='data/charset_36.txt', 
+                 case_sensitive=False, 
+                 one_hot_x=True,
+                 one_hot_y=True,
+                 is_training=True,
+                 smooth_label=False,
+                 smooth_factor=0.2,
+                 use_sm=False,
+                 **kwargs):
+        self.path = Path(path)
+        self.case_sensitive, self.use_sm = case_sensitive, use_sm
+        self.smooth_factor, self.smooth_label = smooth_factor, smooth_label
+        self.charset = CharsetMapper(charset_path, max_length=max_length+1)
+        self.one_hot_x, self.one_hot_y, self.is_training = one_hot_x, one_hot_y, is_training
+        if self.is_training and self.use_sm: self.sm = SpellingMutation(charset=self.charset)
+
+        dtype = {'inp': str, 'gt': str}
+        self.df = pd.read_csv(self.path, dtype=dtype, delimiter=delimiter, na_filter=False)
+        self.inp_col, self.gt_col = 0, 1
+
+    def __len__(self): return len(self.df)
+
+    def __getitem__(self, idx):
+        text_x = self.df.iloc[idx, self.inp_col]
+        text_x = re.sub('[^0-9a-zA-Z]+', '', text_x)
+        if not self.case_sensitive: text_x = text_x.lower()
+        if self.is_training and self.use_sm: text_x = self.sm(text_x)
+
+        length_x = tensor(len(text_x) + 1).to(dtype=torch.long)  # one for end token
+        label_x = self.charset.get_labels(text_x, case_sensitive=self.case_sensitive)
+        label_x = tensor(label_x)
+        if self.one_hot_x:
+            label_x = onehot(label_x, self.charset.num_classes)
+            if self.is_training and self.smooth_label: 
+                label_x = torch.stack([self.prob_smooth_label(l) for l in label_x])
+        x =  [label_x, length_x]
+    
+        text_y = self.df.iloc[idx, self.gt_col]
+        text_y = re.sub('[^0-9a-zA-Z]+', '', text_y)
+        if not self.case_sensitive: text_y = text_y.lower()
+        length_y = tensor(len(text_y) + 1).to(dtype=torch.long)  # one for end token
+        label_y = self.charset.get_labels(text_y, case_sensitive=self.case_sensitive)
+        label_y = tensor(label_y)
+        if self.one_hot_y: label_y = onehot(label_y, self.charset.num_classes)
+        y = [label_y, length_y]
+
+        return x, y
+
+    def prob_smooth_label(self, one_hot):
+        one_hot = one_hot.float()
+        delta = torch.rand([]) * self.smooth_factor
+        num_classes = len(one_hot)
+        noise = torch.rand(num_classes)
+        noise = noise / noise.sum() * delta
+        one_hot = one_hot * (1 - delta) + noise
+        return one_hot
+
+
+class SpellingMutation(object):
+    def __init__(self, pn0=0.7, pn1=0.85, pn2=0.95, pt0=0.7, pt1=0.85, charset=None):
+        """ 
+        Args:
+            pn0: the prob of not modifying characters is (pn0)
+            pn1: the prob of modifying one characters is (pn1 - pn0)
+            pn2: the prob of modifying two characters is (pn2 - pn1), 
+                 and three (1 - pn2)
+            pt0: the prob of replacing operation is pt0.
+            pt1: the prob of inserting operation is (pt1 - pt0),
+                 and deleting operation is (1 - pt1)
+        """
+        super().__init__()
+        self.pn0, self.pn1, self.pn2 = pn0, pn1, pn2
+        self.pt0, self.pt1 = pt0, pt1
+        self.charset = charset
+        logging.info(f'the probs: pn0={self.pn0}, pn1={self.pn1} ' + 
+                     f'pn2={self.pn2}, pt0={self.pt0}, pt1={self.pt1}')
+
+    def is_digit(self, text, ratio=0.5):
+        length = max(len(text), 1)
+        digit_num = sum([t in self.charset.digits for t in text])
+        if digit_num / length < ratio: return False
+        return True
+
+    def is_unk_char(self, char):
+        # return char == self.charset.unk_char
+        return (char not in self.charset.digits) and (char not in self.charset.alphabets)
+
+    def get_num_to_modify(self, length):
+        prob = random.random()
+        if prob < self.pn0: num_to_modify = 0
+        elif prob < self.pn1: num_to_modify = 1
+        elif prob < self.pn2: num_to_modify = 2
+        else: num_to_modify = 3
+        
+        if length <= 1: num_to_modify = 0
+        elif length >= 2 and length <= 4: num_to_modify = min(num_to_modify, 1)
+        else: num_to_modify = min(num_to_modify, length // 2)  # smaller than length // 2
+        return num_to_modify
+
+    def __call__(self, text, debug=False):
+        if self.is_digit(text): return text
+        length = len(text)
+        num_to_modify = self.get_num_to_modify(length)
+        if num_to_modify <= 0: return text
+
+        chars = []
+        index = np.arange(0, length)
+        random.shuffle(index)
+        index = index[: num_to_modify]
+        if debug: self.index = index
+        for i, t in enumerate(text):
+            if i not in index: chars.append(t)
+            elif self.is_unk_char(t): chars.append(t)
+            else:
+                prob = random.random()
+                if prob < self.pt0: # replace
+                    chars.append(random.choice(self.charset.alphabets))
+                elif prob < self.pt1: # insert
+                    chars.append(random.choice(self.charset.alphabets))
+                    chars.append(t)
+                else: # delete
+                    continue
+        new_text = ''.join(chars[: self.charset.max_length-1])
+        return new_text if len(new_text) >= 1 else text

demo.py

@@ -0,0 +1,109 @@
+import argparse
+import logging
+import os
+import glob
+import tqdm
+import torch
+import PIL
+import cv2
+import numpy as np
+import torch.nn.functional as F
+from torchvision import transforms
+from utils import Config, Logger, CharsetMapper
+
+def get_model(config):
+    import importlib
+    names = config.model_name.split('.')
+    module_name, class_name = '.'.join(names[:-1]), names[-1]
+    cls = getattr(importlib.import_module(module_name), class_name)
+    model = cls(config)
+    logging.info(model)
+    model = model.eval()
+    return model
+
+def preprocess(img, width, height):
+    img = cv2.resize(np.array(img), (width, height))
+    img = transforms.ToTensor()(img).unsqueeze(0)
+    mean = torch.tensor([0.485, 0.456, 0.406])
+    std  = torch.tensor([0.229, 0.224, 0.225])
+    return (img-mean[...,None,None]) / std[...,None,None]
+
+def postprocess(output, charset, model_eval):
+    def _get_output(last_output, model_eval):
+        if isinstance(last_output, (tuple, list)): 
+            for res in last_output:
+                if res['name'] == model_eval: output = res
+        else: output = last_output
+        return output
+
+    def _decode(logit):
+        """ Greed decode """
+        out = F.softmax(logit, dim=2)
+        pt_text, pt_scores, pt_lengths = [], [], []
+        for o in out:
+            text = charset.get_text(o.argmax(dim=1), padding=False, trim=False)
+            text = text.split(charset.null_char)[0]  # end at end-token
+            pt_text.append(text)
+            pt_scores.append(o.max(dim=1)[0])
+            pt_lengths.append(min(len(text) + 1, charset.max_length))  # one for end-token
+        return pt_text, pt_scores, pt_lengths
+
+    output = _get_output(output, model_eval)
+    logits, pt_lengths = output['logits'], output['pt_lengths']
+    pt_text, pt_scores, pt_lengths_ = _decode(logits)
+    
+    return pt_text, pt_scores, pt_lengths_
+
+def load(model, file, device=None, strict=True):
+    if device is None: device = 'cpu'
+    elif isinstance(device, int): device = torch.device('cuda', device)
+    assert os.path.isfile(file)
+    state = torch.load(file, map_location=device)
+    if set(state.keys()) == {'model', 'opt'}:
+        state = state['model']
+    model.load_state_dict(state, strict=strict)
+    return model
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, default='configs/train_abinet.yaml',
+                        help='path to config file')
+    parser.add_argument('--input', type=str, default='figs/test')
+    parser.add_argument('--cuda', type=int, default=-1)
+    parser.add_argument('--checkpoint', type=str, default='workdir/train-abinet/best-train-abinet.pth')
+    parser.add_argument('--model_eval', type=str, default='alignment', 
+                        choices=['alignment', 'vision', 'language'])
+    args = parser.parse_args()
+    config = Config(args.config)
+    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
+    if args.model_eval is not None: config.model_eval = args.model_eval
+    config.global_phase = 'test'
+    config.model_vision_checkpoint, config.model_language_checkpoint = None, None
+    device = 'cpu' if args.cuda < 0 else f'cuda:{args.cuda}'
+
+    Logger.init(config.global_workdir, config.global_name, config.global_phase)
+    Logger.enable_file()
+    logging.info(config)
+
+    logging.info('Construct model.')
+    model = get_model(config).to(device)
+    model = load(model, config.model_checkpoint, device=device)
+    charset = CharsetMapper(filename=config.dataset_charset_path,
+                            max_length=config.dataset_max_length + 1)
+
+    if os.path.isdir(args.input):
+        paths = [os.path.join(args.input, fname) for fname in os.listdir(args.input)]
+    else:
+        paths = glob.glob(os.path.expanduser(args.input))
+        assert paths, "The input path(s) was not found"
+    paths = sorted(paths)
+    for path in tqdm.tqdm(paths):
+        img = PIL.Image.open(path).convert('RGB')
+        img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
+        img = img.to(device)
+        res = model(img)
+        pt_text, _, __ = postprocess(res, charset, config.model_eval)
+        logging.info(f'{path}: {pt_text[0]}')
+
+if __name__ == '__main__':
+    main()

docker/Dockerfile

@@ -0,0 +1,25 @@
+FROM anibali/pytorch:cuda-9.0
+MAINTAINER fangshancheng <fangsc@ustc.edu.cn>
+RUN sudo rm -rf /etc/apt/sources.list.d && \
+    sudo apt update && \
+    sudo apt install -y build-essential vim && \
+    conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/ && \
+    conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/ && \
+    conda config --set show_channel_urls yes && \
+    pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/ && \
+    pip install torch==1.1.0 torchvision==0.3.0 && \
+    pip install fastai==1.0.60 && \
+    pip install ipdb jupyter ipython lmdb editdistance tensorboardX natsort nltk && \
+    conda uninstall -y --force pillow pil jpeg libtiff libjpeg-turbo && \
+    pip uninstall -y pillow pil jpeg libtiff libjpeg-turbo && \
+    conda install -yc conda-forge libjpeg-turbo && \
+    CFLAGS="${CFLAGS} -mavx2" pip install --no-cache-dir --force-reinstall --no-binary :all: --compile pillow-simd==6.2.2.post1 && \
+    conda install -y jpeg libtiff opencv && \
+    sudo rm -rf /var/lib/apt/lists/* && \
+    sudo rm -rf /tmp/* && \
+    sudo rm -rf ~/.cache && \
+    sudo apt clean all && \
+    conda clean -y -a
+EXPOSE 8888
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8

losses.py

@@ -0,0 +1,72 @@
+from fastai.vision import *
+
+from modules.model import Model
+
+
+class MultiLosses(nn.Module):
+    def __init__(self, one_hot=True):
+        super().__init__()
+        self.ce = SoftCrossEntropyLoss() if one_hot else torch.nn.CrossEntropyLoss()
+        self.bce = torch.nn.BCELoss()
+
+    @property
+    def last_losses(self):
+        return self.losses
+
+    def _flatten(self, sources, lengths):
+        return torch.cat([t[:l] for t, l in zip(sources, lengths)])
+
+    def _merge_list(self, all_res):
+        if not isinstance(all_res, (list, tuple)):
+            return all_res
+        def merge(items):
+            if isinstance(items[0], torch.Tensor): return torch.cat(items, dim=0)
+            else: return items[0]
+        res = dict()
+        for key in all_res[0].keys():
+            items = [r[key] for r in all_res]
+            res[key] = merge(items)
+        return res
+
+    def _ce_loss(self, output, gt_labels, gt_lengths, idx=None, record=True):
+        loss_name = output.get('name')
+        pt_logits, weight = output['logits'], output['loss_weight']
+
+        assert pt_logits.shape[0] % gt_labels.shape[0] == 0
+        iter_size = pt_logits.shape[0] // gt_labels.shape[0]
+        if iter_size > 1:
+            gt_labels = gt_labels.repeat(3, 1, 1)
+            gt_lengths = gt_lengths.repeat(3)
+        flat_gt_labels = self._flatten(gt_labels, gt_lengths)
+        flat_pt_logits = self._flatten(pt_logits, gt_lengths)
+
+        nll = output.get('nll')
+        if nll is not None:
+            loss = self.ce(flat_pt_logits, flat_gt_labels, softmax=False) * weight
+        else:
+            loss = self.ce(flat_pt_logits, flat_gt_labels) * weight
+        if record and loss_name is not None: self.losses[f'{loss_name}_loss'] = loss
+
+        return loss
+
+    def forward(self, outputs, *args):
+        self.losses = {}
+        if isinstance(outputs, (tuple, list)):
+            outputs = [self._merge_list(o) for o in outputs]
+            return sum([self._ce_loss(o, *args) for o in outputs if o['loss_weight'] > 0.])
+        else:
+            return self._ce_loss(outputs, *args, record=False)
+
+
+class SoftCrossEntropyLoss(nn.Module):
+    def __init__(self, reduction="mean"):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, input, target, softmax=True):
+        if softmax: log_prob = F.log_softmax(input, dim=-1)
+        else: log_prob = torch.log(input)
+        loss = -(target * log_prob).sum(dim=-1)
+        if self.reduction == "mean": return loss.mean()
+        elif self.reduction == "sum": return loss.sum()
+        else: return loss

main.py

@@ -0,0 +1,246 @@
+import argparse
+import logging
+import os
+import random
+
+import torch
+from fastai.callbacks.general_sched import GeneralScheduler, TrainingPhase
+from fastai.distributed import *
+from fastai.vision import *
+from torch.backends import cudnn
+
+from callbacks import DumpPrediction, IterationCallback, TextAccuracy, TopKTextAccuracy
+from dataset import ImageDataset, TextDataset
+from losses import MultiLosses
+from utils import Config, Logger, MyDataParallel, MyConcatDataset
+
+
+def _set_random_seed(seed):
+    if seed is not None:
+        random.seed(seed)
+        torch.manual_seed(seed)
+        cudnn.deterministic = True
+        logging.warning('You have chosen to seed training. '
+                        'This will slow down your training!')
+
+def _get_training_phases(config, n):
+    lr = np.array(config.optimizer_lr)
+    periods = config.optimizer_scheduler_periods
+    sigma = [config.optimizer_scheduler_gamma ** i for i in range(len(periods))]
+    phases = [TrainingPhase(n * periods[i]).schedule_hp('lr', lr * sigma[i])
+                for i in range(len(periods))]
+    return phases
+
+def _get_dataset(ds_type, paths, is_training, config, **kwargs):
+    kwargs.update({
+        'img_h': config.dataset_image_height,
+        'img_w': config.dataset_image_width,
+        'max_length': config.dataset_max_length,
+        'case_sensitive': config.dataset_case_sensitive,
+        'charset_path': config.dataset_charset_path,
+        'data_aug': config.dataset_data_aug,
+        'deteriorate_ratio': config.dataset_deteriorate_ratio,
+        'is_training': is_training,
+        'multiscales': config.dataset_multiscales,
+        'one_hot_y': config.dataset_one_hot_y,
+    })
+    datasets = [ds_type(p, **kwargs) for p in paths]
+    if len(datasets) > 1: return MyConcatDataset(datasets)
+    else: return datasets[0]
+
+
+def _get_language_databaunch(config):
+    kwargs = {
+        'max_length': config.dataset_max_length,
+        'case_sensitive': config.dataset_case_sensitive,
+        'charset_path': config.dataset_charset_path,
+        'smooth_label': config.dataset_smooth_label,
+        'smooth_factor': config.dataset_smooth_factor,
+        'one_hot_y': config.dataset_one_hot_y,
+        'use_sm': config.dataset_use_sm,
+    }
+    train_ds = TextDataset(config.dataset_train_roots[0], is_training=True, **kwargs)
+    valid_ds = TextDataset(config.dataset_test_roots[0], is_training=False, **kwargs)
+    data = DataBunch.create(
+        path=train_ds.path,
+        train_ds=train_ds,
+        valid_ds=valid_ds,
+        bs=config.dataset_train_batch_size,
+        val_bs=config.dataset_test_batch_size,
+        num_workers=config.dataset_num_workers,
+        pin_memory=config.dataset_pin_memory)
+    logging.info(f'{len(data.train_ds)} training items found.')
+    if not data.empty_val:
+        logging.info(f'{len(data.valid_ds)} valid items found.')
+    return data
+
+def _get_databaunch(config):
+    # An awkward way to reduce loadding data time during test
+    if config.global_phase == 'test': config.dataset_train_roots = config.dataset_test_roots
+    train_ds = _get_dataset(ImageDataset, config.dataset_train_roots, True, config)
+    valid_ds = _get_dataset(ImageDataset, config.dataset_test_roots, False, config)
+    data = ImageDataBunch.create(
+        train_ds=train_ds,
+        valid_ds=valid_ds,
+        bs=config.dataset_train_batch_size,
+        val_bs=config.dataset_test_batch_size,
+        num_workers=config.dataset_num_workers,
+        pin_memory=config.dataset_pin_memory).normalize(imagenet_stats)
+    ar_tfm = lambda x: ((x[0], x[1]), x[1])  # auto-regression only for dtd
+    data.add_tfm(ar_tfm)
+
+    logging.info(f'{len(data.train_ds)} training items found.')
+    if not data.empty_val:
+        logging.info(f'{len(data.valid_ds)} valid items found.')
+    
+    return data
+
+def _get_model(config):
+    import importlib
+    names = config.model_name.split('.')
+    module_name, class_name = '.'.join(names[:-1]), names[-1]
+    cls = getattr(importlib.import_module(module_name), class_name)
+    model = cls(config)
+    logging.info(model)
+    return model
+
+
+def _get_learner(config, data, model, local_rank=None):
+    strict = ifnone(config.model_strict, True)
+    if config.global_stage == 'pretrain-language':
+        metrics = [TopKTextAccuracy(
+            k=ifnone(config.model_k, 5),
+            charset_path=config.dataset_charset_path,
+            max_length=config.dataset_max_length + 1,
+            case_sensitive=config.dataset_eval_case_sensisitves,
+            model_eval=config.model_eval)] 
+    else:
+        metrics = [TextAccuracy(
+            charset_path=config.dataset_charset_path,
+            max_length=config.dataset_max_length + 1,
+            case_sensitive=config.dataset_eval_case_sensisitves,
+            model_eval=config.model_eval)]
+    opt_type = getattr(torch.optim, config.optimizer_type)
+    learner = Learner(data, model, silent=True, model_dir='.',
+        true_wd=config.optimizer_true_wd, 
+        wd=config.optimizer_wd,
+        bn_wd=config.optimizer_bn_wd,
+        path=config.global_workdir,
+        metrics=metrics,
+        opt_func=partial(opt_type, **config.optimizer_args or dict()), 
+        loss_func=MultiLosses(one_hot=config.dataset_one_hot_y))
+    learner.split(lambda m: children(m))
+
+    if config.global_phase == 'train':
+        num_replicas = 1 if local_rank is None else torch.distributed.get_world_size()
+        phases = _get_training_phases(config, len(learner.data.train_dl)//num_replicas)
+        learner.callback_fns += [
+            partial(GeneralScheduler, phases=phases),
+            partial(GradientClipping, clip=config.optimizer_clip_grad),
+            partial(IterationCallback, name=config.global_name,
+                    show_iters=config.training_show_iters,
+                    eval_iters=config.training_eval_iters,
+                    save_iters=config.training_save_iters,
+                    start_iters=config.training_start_iters,
+                    stats_iters=config.training_stats_iters)]
+    else:
+        learner.callbacks += [
+            DumpPrediction(learn=learner,
+                    dataset='-'.join([Path(p).name for p in config.dataset_test_roots]),charset_path=config.dataset_charset_path,
+                    model_eval=config.model_eval,
+                    debug=config.global_debug,
+                    image_only=config.global_image_only)]
+
+    learner.rank = local_rank
+    if local_rank is not None:
+        logging.info(f'Set model to distributed with rank {local_rank}.')
+        learner.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(learner.model)
+        learner.model.to(local_rank)
+        learner = learner.to_distributed(local_rank)
+
+    if torch.cuda.device_count() > 1 and local_rank is None:
+        logging.info(f'Use {torch.cuda.device_count()} GPUs.')
+        learner.model = MyDataParallel(learner.model)
+
+    if config.model_checkpoint:
+        if Path(config.model_checkpoint).exists():
+            with open(config.model_checkpoint, 'rb') as f:
+                buffer = io.BytesIO(f.read())
+            learner.load(buffer, strict=strict)
+        else:
+            from distutils.dir_util import copy_tree
+            src = Path('/data/fangsc/model')/config.global_name
+            trg = Path('/output')/config.global_name
+            if src.exists(): copy_tree(str(src), str(trg))
+            learner.load(config.model_checkpoint, strict=strict)
+        logging.info(f'Read model from {config.model_checkpoint}')
+    elif config.global_phase == 'test':
+        learner.load(f'best-{config.global_name}', strict=strict)
+        logging.info(f'Read model from best-{config.global_name}')
+
+    if learner.opt_func.func.__name__ == 'Adadelta':    # fastai bug, fix after 1.0.60
+        learner.fit(epochs=0, lr=config.optimizer_lr)
+        learner.opt.mom = 0.
+
+    return learner
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True,
+                        help='path to config file')
+    parser.add_argument('--phase', type=str, default=None, choices=['train', 'test'])
+    parser.add_argument('--name', type=str, default=None)
+    parser.add_argument('--checkpoint', type=str, default=None)
+    parser.add_argument('--test_root', type=str, default=None)
+    parser.add_argument("--local_rank", type=int, default=None)
+    parser.add_argument('--debug', action='store_true', default=None)
+    parser.add_argument('--image_only', action='store_true', default=None)
+    parser.add_argument('--model_strict', action='store_false', default=None)
+    parser.add_argument('--model_eval', type=str, default=None, 
+                        choices=['alignment', 'vision', 'language'])
+    args = parser.parse_args()
+    config = Config(args.config)
+    if args.name is not None: config.global_name = args.name
+    if args.phase is not None: config.global_phase = args.phase
+    if args.test_root is not None: config.dataset_test_roots = [args.test_root]
+    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
+    if args.debug is not None: config.global_debug = args.debug
+    if args.image_only is not None: config.global_image_only = args.image_only
+    if args.model_eval is not None: config.model_eval = args.model_eval
+    if args.model_strict is not None: config.model_strict = args.model_strict
+
+    Logger.init(config.global_workdir, config.global_name, config.global_phase)
+    Logger.enable_file()
+    _set_random_seed(config.global_seed)
+    logging.info(config)
+
+    if args.local_rank is not None:
+        logging.info(f'Init distribution training at device {args.local_rank}.')
+        torch.cuda.set_device(args.local_rank)
+        torch.distributed.init_process_group(backend='nccl', init_method='env://')
+
+    logging.info('Construct dataset.')
+    if config.global_stage == 'pretrain-language': data = _get_language_databaunch(config)
+    else: data = _get_databaunch(config)
+
+    logging.info('Construct model.')
+    model = _get_model(config)
+
+    logging.info('Construct learner.')
+    learner = _get_learner(config, data, model, args.local_rank)
+
+    if config.global_phase == 'train':
+        logging.info('Start training.')
+        learner.fit(epochs=config.training_epochs,
+                    lr=config.optimizer_lr)
+    else:
+        logging.info('Start validate')
+        last_metrics = learner.validate()
+        log_str = f'eval loss = {last_metrics[0]:6.3f},  ' \
+                  f'ccr = {last_metrics[1]:6.3f},  cwr = {last_metrics[2]:6.3f},  ' \
+                  f'ted = {last_metrics[3]:6.3f},  ned = {last_metrics[4]:6.0f},  ' \
+                  f'ted/w = {last_metrics[5]:6.3f}, '
+        logging.info(log_str)
+
+if __name__ == '__main__':
+    main()

modules/attention.py

@@ -0,0 +1,97 @@
+import torch
+import torch.nn as nn
+from .transformer import PositionalEncoding
+
+class Attention(nn.Module):
+    def __init__(self, in_channels=512, max_length=25, n_feature=256):
+        super().__init__()
+        self.max_length = max_length
+
+        self.f0_embedding = nn.Embedding(max_length, in_channels)
+        self.w0 = nn.Linear(max_length, n_feature)
+        self.wv = nn.Linear(in_channels, in_channels)
+        self.we = nn.Linear(in_channels, max_length)
+
+        self.active = nn.Tanh()
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, enc_output):
+        enc_output = enc_output.permute(0, 2, 3, 1).flatten(1, 2)
+        reading_order = torch.arange(self.max_length, dtype=torch.long, device=enc_output.device)
+        reading_order = reading_order.unsqueeze(0).expand(enc_output.size(0), -1)  # (S,) -> (B, S)
+        reading_order_embed = self.f0_embedding(reading_order)  # b,25,512
+
+        t = self.w0(reading_order_embed.permute(0, 2, 1))  # b,512,256
+        t = self.active(t.permute(0, 2, 1) + self.wv(enc_output))  # b,256,512
+
+        attn = self.we(t)  # b,256,25
+        attn = self.softmax(attn.permute(0, 2, 1))  # b,25,256
+        g_output = torch.bmm(attn, enc_output)  # b,25,512
+        return g_output, attn.view(*attn.shape[:2], 8, 32)
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
+    align_corners = None if mode=='nearest' else True
+    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor, 
+                                     mode=mode, align_corners=align_corners),
+                         nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+
+class PositionAttention(nn.Module):
+    def __init__(self, max_length, in_channels=512, num_channels=64, 
+                 h=8, w=32, mode='nearest', **kwargs):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2))
+        )
+        self.k_decoder = nn.Sequential(
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
+        )
+
+        self.pos_encoder = PositionalEncoding(in_channels, dropout=0, max_len=max_length)
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x):
+        N, E, H, W = x.size()
+        k, v = x, x  # (N, E, H, W)
+
+        # calculate key vector
+        features = []
+        for i in range(0, len(self.k_encoder)):
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k)
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k)
+
+        # calculate query vector
+        # TODO q=f(q,k)
+        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)
+        q = self.pos_encoder(zeros)  # (T, N, E)
+        q = q.permute(1, 0, 2)  # (N, T, E)
+        q = self.project(q)  # (N, T, E)
+        
+        # calculate attention
+        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
+        attn_scores = attn_scores / (E ** 0.5)
+        attn_scores = torch.softmax(attn_scores, dim=-1)
+
+        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
+        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)
+
+        return attn_vecs, attn_scores.view(N, -1, H, W)

modules/backbone.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+from fastai.vision import *
+
+from modules.model import _default_tfmer_cfg
+from modules.resnet import resnet45
+from modules.transformer import (PositionalEncoding,
+                                 TransformerEncoder,
+                                 TransformerEncoderLayer)
+
+
+class ResTranformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.resnet = resnet45()
+
+        self.d_model = ifnone(config.model_vision_d_model, _default_tfmer_cfg['d_model'])
+        nhead = ifnone(config.model_vision_nhead, _default_tfmer_cfg['nhead'])
+        d_inner = ifnone(config.model_vision_d_inner, _default_tfmer_cfg['d_inner'])
+        dropout = ifnone(config.model_vision_dropout, _default_tfmer_cfg['dropout'])
+        activation = ifnone(config.model_vision_activation, _default_tfmer_cfg['activation'])
+        num_layers = ifnone(config.model_vision_backbone_ln, 2)
+
+        self.pos_encoder = PositionalEncoding(self.d_model, max_len=8*32)
+        encoder_layer = TransformerEncoderLayer(d_model=self.d_model, nhead=nhead, 
+                dim_feedforward=d_inner, dropout=dropout, activation=activation)
+        self.transformer = TransformerEncoder(encoder_layer, num_layers)
+
+    def forward(self, images):
+        feature = self.resnet(images)
+        n, c, h, w = feature.shape
+        feature = feature.view(n, c, -1).permute(2, 0, 1)
+        feature = self.pos_encoder(feature)
+        feature = self.transformer(feature)
+        feature = feature.permute(1, 2, 0).view(n, c, h, w)
+        return feature