openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/optimizer/__init__.py

@@ -0,0 +1,73 @@
+import copy
+
+import torch
+from torch import nn
+
+__all__ = ['build_optimizer']
+
+
+def param_groups_weight_decay(model: nn.Module,
+                              weight_decay=1e-5,
+                              no_weight_decay_list=()):
+    no_weight_decay_list = set(no_weight_decay_list)
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if param.ndim <= 1 or name.endswith(
+                '.bias') or any(nd in name for nd in no_weight_decay_list):
+            no_decay.append(param)
+        else:
+            decay.append(param)
+
+    return [
+        {
+            'params': no_decay,
+            'weight_decay': 0.0
+        },
+        {
+            'params': decay,
+            'weight_decay': weight_decay
+        },
+    ]
+
+
+def build_optimizer(optim_config, lr_scheduler_config, epochs, step_each_epoch,
+                    model):
+    from . import lr
+
+    config = copy.deepcopy(optim_config)
+
+    if isinstance(model, nn.Module):
+        # a model was passed in, extract parameters and add weight decays to appropriate layers
+        weight_decay = config.get('weight_decay', 0.0)
+        filter_bias_and_bn = (config.pop('filter_bias_and_bn')
+                              if 'filter_bias_and_bn' in config else False)
+        if weight_decay > 0.0 and filter_bias_and_bn:
+            no_weight_decay = {}
+            if hasattr(model, 'no_weight_decay'):
+                no_weight_decay = model.no_weight_decay()
+            parameters = param_groups_weight_decay(model, weight_decay,
+                                                   no_weight_decay)
+            config['weight_decay'] = 0.0
+            # print('debug adamw')
+        else:
+            parameters = model.parameters()
+    else:
+        # iterable of parameters or param groups passed in
+        parameters = model
+
+    optim = getattr(torch.optim, config.pop('name'))(params=parameters,
+                                                     **config)
+
+    lr_config = copy.deepcopy(lr_scheduler_config)
+    lr_config.update({
+        'epochs': epochs,
+        'step_each_epoch': step_each_epoch,
+        'lr': config['lr']
+    })
+    lr_scheduler = getattr(lr,
+                           lr_config.pop('name'))(**lr_config)(optimizer=optim)
+    return optim, lr_scheduler

openocr/openrec/optimizer/lr.py

@@ -0,0 +1,227 @@
+import math
+from functools import partial
+
+import numpy as np
+from torch.optim import lr_scheduler
+
+
+class StepLR(object):
+
+    def __init__(self,
+                 step_each_epoch,
+                 step_size,
+                 warmup_epoch=0,
+                 gamma=0.1,
+                 last_epoch=-1,
+                 **kwargs):
+        super(StepLR, self).__init__()
+        self.step_size = step_each_epoch * step_size
+        self.gamma = gamma
+        self.last_epoch = last_epoch
+        self.warmup_epoch = warmup_epoch
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1, self.warmup_epoch))
+        return self.gamma**(current_step // self.step_size)
+
+
+class MultiStepLR(object):
+
+    def __init__(self,
+                 step_each_epoch,
+                 milestones,
+                 warmup_epoch=0,
+                 gamma=0.1,
+                 last_epoch=-1,
+                 **kwargs):
+        super(MultiStepLR, self).__init__()
+        self.milestones = [step_each_epoch * e for e in milestones]
+        self.gamma = gamma
+        self.last_epoch = last_epoch
+        self.warmup_epoch = warmup_epoch
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1, self.warmup_epoch))
+        return self.gamma**len(
+            [m for m in self.milestones if m <= current_step])
+
+
+class ConstLR(object):
+
+    def __init__(self,
+                 step_each_epoch,
+                 warmup_epoch=0,
+                 last_epoch=-1,
+                 **kwargs):
+        super(ConstLR, self).__init__()
+        self.last_epoch = last_epoch
+        self.warmup_epoch = warmup_epoch * step_each_epoch
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1.0, self.warmup_epoch))
+        return 1.0
+
+
+class LinearLR(object):
+
+    def __init__(self,
+                 epochs,
+                 step_each_epoch,
+                 warmup_epoch=0,
+                 last_epoch=-1,
+                 **kwargs):
+        super(LinearLR, self).__init__()
+        self.epochs = epochs * step_each_epoch
+        self.last_epoch = last_epoch
+        self.warmup_epoch = warmup_epoch * step_each_epoch
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1, self.warmup_epoch))
+        return max(
+            0.0,
+            float(self.epochs - current_step) /
+            float(max(1, self.epochs - self.warmup_epoch)),
+        )
+
+
+class CosineAnnealingLR(object):
+
+    def __init__(self,
+                 epochs,
+                 step_each_epoch,
+                 warmup_epoch=0,
+                 last_epoch=-1,
+                 **kwargs):
+        super(CosineAnnealingLR, self).__init__()
+        self.epochs = epochs * step_each_epoch
+        self.last_epoch = last_epoch
+        self.warmup_epoch = warmup_epoch * step_each_epoch
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step, num_cycles=0.5):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1, self.warmup_epoch))
+        progress = float(current_step - self.warmup_epoch) / float(
+            max(1, self.epochs - self.warmup_epoch))
+        return max(
+            0.0, 0.5 *
+            (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+
+class OneCycleLR(object):
+
+    def __init__(self,
+                 epochs,
+                 step_each_epoch,
+                 last_epoch=-1,
+                 lr=0.00148,
+                 warmup_epoch=1.0,
+                 cycle_momentum=True,
+                 **kwargs):
+        super(OneCycleLR, self).__init__()
+        self.epochs = epochs
+        self.last_epoch = last_epoch
+        self.step_each_epoch = step_each_epoch
+        self.lr = lr
+        self.pct_start = warmup_epoch / epochs
+        self.cycle_momentum = cycle_momentum
+
+    def __call__(self, optimizer):
+        return lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=self.lr,
+            total_steps=self.epochs * self.step_each_epoch,
+            pct_start=self.pct_start,
+            cycle_momentum=self.cycle_momentum,
+        )
+
+
+class PolynomialLR(object):
+
+    def __init__(self,
+                 step_each_epoch,
+                 epochs,
+                 lr_end=1e-7,
+                 power=1.0,
+                 warmup_epoch=0,
+                 last_epoch=-1,
+                 **kwargs):
+        super(PolynomialLR, self).__init__()
+        self.lr_end = lr_end
+        self.power = power
+        self.epochs = epochs * step_each_epoch
+        self.warmup_epoch = warmup_epoch * step_each_epoch
+        self.last_epoch = last_epoch
+
+    def __call__(self, optimizer):
+        lr_lambda = partial(
+            self.lambda_func,
+            lr_init=optimizer.defaults['lr'],
+        )
+        return lr_scheduler.LambdaLR(optimizer, lr_lambda, self.last_epoch)
+
+    def lambda_func(self, current_step, lr_init):
+        if current_step < self.warmup_epoch:
+            return float(current_step) / float(max(1, self.warmup_epoch))
+        elif current_step > self.epochs:
+            return self.lr_end / lr_init  # as LambdaLR multiplies by lr_init
+        else:
+            lr_range = lr_init - self.lr_end
+            decay_steps = self.epochs - self.warmup_epoch
+            pct_remaining = 1 - (current_step -
+                                 self.warmup_epoch) / decay_steps
+            decay = lr_range * pct_remaining**self.power + self.lr_end
+            return decay / lr_init  # as LambdaLR multiplies by lr_init
+
+
+class CdistNetLR(object):
+
+    def __init__(self,
+                 step_each_epoch,
+                 lr=0.0442,
+                 n_warmup_steps=10000,
+                 step2_epoch=7,
+                 last_epoch=-1,
+                 **kwargs):
+        super(CdistNetLR, self).__init__()
+        self.last_epoch = last_epoch
+        self.step2_epoch = step2_epoch * step_each_epoch
+        self.n_current_steps = 0
+        self.n_warmup_steps = n_warmup_steps
+        self.init_lr = lr
+        self.step2_lr = 0.00001
+
+    def __call__(self, optimizer):
+        return lr_scheduler.LambdaLR(optimizer, self.lambda_func,
+                                     self.last_epoch)
+
+    def lambda_func(self, current_step):
+        if current_step < self.step2_epoch:
+            return np.min([
+                np.power(current_step, -0.5),
+                np.power(self.n_warmup_steps, -1.5) * current_step,
+            ])
+        return self.step2_lr / self.init_lr

openocr/openrec/postprocess/__init__.py

@@ -0,0 +1,72 @@
+import copy
+
+__all__ = ['build_post_process']
+
+from .abinet_postprocess import ABINetLabelDecode
+from .ar_postprocess import ARLabelDecode
+from .ce_postprocess import CELabelDecode
+from .char_postprocess import CharLabelDecode
+from .cppd_postprocess import CPPDLabelDecode
+from .ctc_postprocess import CTCLabelDecode
+from .igtr_postprocess import IGTRLabelDecode
+from .lister_postprocess import LISTERLabelDecode
+from .mgp_postprocess import MPGLabelDecode
+from .nrtr_postprocess import NRTRLabelDecode
+from .smtr_postprocess import SMTRLabelDecode
+from .srn_postprocess import SRNLabelDecode
+from .visionlan_postprocess import VisionLANLabelDecode
+
+support_dict = [
+    'CTCLabelDecode', 'CharLabelDecode', 'CELabelDecode', 'CPPDLabelDecode',
+    'NRTRLabelDecode', 'ABINetLabelDecode', 'ARLabelDecode', 'IGTRLabelDecode',
+    'VisionLANLabelDecode', 'SMTRLabelDecode', 'SRNLabelDecode',
+    'LISTERLabelDecode', 'GTCLabelDecode', 'MPGLabelDecode'
+]
+
+
+def build_post_process(config, global_config=None):
+    config = copy.deepcopy(config)
+    module_name = config.pop('name')
+    if global_config is not None:
+        config.update(global_config)
+    assert module_name in support_dict, Exception(
+        'post process only support {}'.format(support_dict))
+    module_class = eval(module_name)(**config)
+    return module_class
+
+
+class GTCLabelDecode(object):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 gtc_label_decode=None,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 only_gtc=False,
+                 with_ratio=False,
+                 **kwargs):
+        gtc_label_decode['character_dict_path'] = character_dict_path
+        gtc_label_decode['use_space_char'] = use_space_char
+        self.gtc_label_decode = build_post_process(gtc_label_decode)
+        self.ctc_label_decode = CTCLabelDecode(
+            character_dict_path=character_dict_path,
+            use_space_char=use_space_char)
+        self.gtc_character = self.gtc_label_decode.character
+        self.ctc_character = self.ctc_label_decode.character
+        self.only_gtc = only_gtc
+        self.with_ratio = with_ratio
+
+    def get_character_num(self):
+        return [len(self.gtc_character), len(self.ctc_character)]
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        if self.with_ratio:
+            batch = batch[:-1]
+        gtc = self.gtc_label_decode(preds['gtc_pred'],
+                                    batch[:-2] if batch is not None else None)
+        if self.only_gtc:
+            return gtc
+        ctc = self.ctc_label_decode(preds['ctc_pred'], [None] +
+                                    batch[-2:] if batch is not None else None)
+
+        return [gtc, ctc]

openocr/openrec/postprocess/abinet_postprocess.py

@@ -0,0 +1,37 @@
+import torch
+
+from .nrtr_postprocess import NRTRLabelDecode
+
+
+class ABINetLabelDecode(NRTRLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(ABINetLabelDecode, self).__init__(character_dict_path,
+                                                use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        if isinstance(preds, dict):
+            if len(preds['align']) > 0:
+                preds = preds['align'][-1].detach().cpu().numpy()
+            else:
+                preds = preds['vision'].detach().cpu().numpy()
+        elif isinstance(preds, torch.Tensor):
+            preds = preds.detach().cpu().numpy()
+        else:
+            preds = preds
+
+        preds_idx = preds.argmax(axis=2)
+        preds_prob = preds.max(axis=2)
+        text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
+        if batch is None:
+            return text
+        label = self.decode(batch[1].cpu().numpy())
+        return text, label
+
+    def add_special_char(self, dict_character):
+        dict_character = ['</s>'] + dict_character
+        return dict_character

openocr/openrec/postprocess/ar_postprocess.py

@@ -0,0 +1,63 @@
+import numpy as np
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class ARLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    BOS = '<s>'
+    EOS = '</s>'
+    PAD = '<pad>'
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 **kwargs):
+        super(ARLabelDecode, self).__init__(character_dict_path,
+                                            use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        if isinstance(preds, list):
+            preds = preds[-1]
+        if isinstance(preds, torch.Tensor):
+            preds = preds.detach().cpu().numpy()
+        preds_idx = preds.argmax(axis=2)
+        preds_prob = preds.max(axis=2)
+        text = self.decode(preds_idx, preds_prob)
+        if batch is None:
+            return text
+        label = batch[1]
+        label = self.decode(label[:, 1:].detach().cpu().numpy())
+        return text, label
+
+    def add_special_char(self, dict_character):
+        dict_character = [self.EOS] + dict_character + [self.BOS, self.PAD]
+        return dict_character
+
+    def decode(self, text_index, text_prob=None):
+        """convert text-index into text-label."""
+        result_list = []
+        batch_size = len(text_index)
+        for batch_idx in range(batch_size):
+            char_list = []
+            conf_list = []
+            for idx in range(len(text_index[batch_idx])):
+                try:
+                    char_idx = self.character[int(text_index[batch_idx][idx])]
+                except:
+                    continue
+                if char_idx == self.EOS:  # end
+                    break
+                if char_idx == self.BOS or char_idx == self.PAD:
+                    continue
+                char_list.append(char_idx)
+                if text_prob is not None:
+                    conf_list.append(text_prob[batch_idx][idx])
+                else:
+                    conf_list.append(1)
+            text = ''.join(char_list)
+            result_list.append((text, np.mean(conf_list).tolist()))
+        return result_list

openocr/openrec/postprocess/ce_postprocess.py

@@ -0,0 +1,43 @@
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class CELabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(CELabelDecode, self).__init__(character_dict_path,
+                                            use_space_char)
+
+    def __call__(self, preds, label=None, *args, **kwargs):
+        if isinstance(preds, tuple) or isinstance(preds, list):
+            preds = preds[-1]
+        if isinstance(preds, torch.Tensor):
+            preds = preds.numpy()
+        preds_idx = preds.argmax(axis=-1)
+        preds_prob = preds.max(axis=-1)
+        text = self.decode(preds_idx, preds_prob)
+        if label is None:
+            return text
+        label = self.decode(label.flatten())
+        return text, label
+
+    def decode(self, text_index, text_prob=None):
+        """convert text-index into text-label."""
+        result_list = []
+        batch_size = len(text_index)
+        for batch_idx in range(batch_size):
+            text = self.character[text_index[batch_idx]]
+            if text_prob is not None:
+                conf_list = text_prob[batch_idx]
+            else:
+                conf_list = 1.0
+            result_list.append((text, conf_list))
+        return result_list
+
+    def add_special_char(self, dict_character):
+        return dict_character

openocr/openrec/postprocess/char_postprocess.py

@@ -0,0 +1,108 @@
+import numpy as np
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class CharLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 **kwargs):
+        super(CharLabelDecode, self).__init__(character_dict_path,
+                                              use_space_char)
+
+    def __call__(self, preds, label=None, *args, **kwargs):
+        if len(preds) >= 4:
+            preds_id = preds[0]
+            preds_prob = preds[1]
+            char_preds = preds[2]
+            if isinstance(preds_id, torch.Tensor):
+                preds_id = preds_id.numpy()
+            if isinstance(preds_prob, torch.Tensor):
+                preds_prob = preds_prob.numpy()
+            if preds_id[0][0] == 2:
+                preds_idx = preds_id[:, 1:]
+                preds_prob = preds_prob[:, 1:]
+                # char_preds = char_preds[:, 1:]
+            else:
+                preds_idx = preds_id
+            char_preds = char_preds.numpy()
+            char_preds_idx = char_preds.argmax(-1) + 4
+            char_preds_prob = char_preds.max(-1)
+            text, text_box = self.decode(preds_idx, preds_prob, char_preds_idx,
+                                         char_preds_prob)
+        else:
+            preds_logit = preds[0].numpy()
+            char_preds = preds[1].numpy()
+            # if isinstance(preds, torch.Tensor):
+            #     preds = preds.numpy()
+            preds_idx = preds_logit.argmax(axis=2)
+            preds_prob = preds_logit.max(axis=2)
+            char_preds_idx = char_preds.argmax(-1) + 4
+            char_preds_prob = char_preds.max(-1)
+            text, text_box = self.decode(preds_idx, preds_prob, char_preds_idx,
+                                         char_preds_prob)
+
+        if label is None:
+            return text, text_box
+        label = self.decode(label[:, 1:])
+        return text, text_box, label
+
+    def add_special_char(self, dict_character):
+        dict_character = ['blank', '<unk>', '<s>', '</s>'] + dict_character
+        return dict_character
+
+    def decode(
+        self,
+        text_index,
+        text_prob=None,
+        char_text_index=None,
+        char_text_prob=None,
+        is_remove_duplicate=False,
+    ):
+        """convert text-index into text-label."""
+        result_list = []
+        box_result_list = []
+        batch_size = len(text_index)
+        for batch_idx in range(batch_size):
+            char_list = []
+            conf_list = []
+            char_box_list = []
+            conf_box_list = []
+            for idx in range(len(text_index[batch_idx])):
+                try:
+                    char_idx = self.character[int(text_index[batch_idx][idx])]
+                    if char_text_index is not None:
+                        char_box_idx = self.character[int(
+                            char_text_index[batch_idx][idx])]
+                except:
+                    continue
+                if char_idx == '</s>':  # end
+                    break
+                char_list.append(char_idx)
+
+                if char_text_index is not None:
+                    char_box_list.append(char_box_idx)
+
+                if text_prob is not None:
+                    conf_list.append(text_prob[batch_idx][idx])
+                else:
+                    conf_list.append(1)
+
+                if char_text_prob is not None:
+                    conf_box_list.append(char_text_prob[batch_idx][idx])
+                else:
+                    conf_box_list.append(1)
+            text = ''.join(char_list)
+            result_list.append((text, np.mean(conf_list).tolist()))
+
+            if char_text_index is not None:
+                text_box = ''.join(char_box_list)
+                box_result_list.append(
+                    (text_box, np.mean(conf_box_list).tolist()))
+        if char_text_index is not None:
+            return result_list, box_result_list
+        return result_list

openocr/openrec/postprocess/cppd_postprocess.py

@@ -0,0 +1,42 @@
+import torch
+
+from .nrtr_postprocess import NRTRLabelDecode
+
+
+class CPPDLabelDecode(NRTRLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(CPPDLabelDecode, self).__init__(character_dict_path,
+                                              use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        if isinstance(preds, tuple):
+            if isinstance(preds[-1], dict):
+                preds = preds[-1]['align'][-1].detach().cpu().numpy()
+            else:
+                preds = preds[-1].detach().cpu().numpy()
+        if isinstance(preds, list):
+            preds = preds[-1].detach().cpu().numpy()
+        if isinstance(preds, torch.Tensor):
+            preds = preds.detach().cpu().numpy()
+        elif isinstance(preds, dict):
+            preds = preds['align'][-1].detach().cpu().numpy()
+        else:
+            preds = preds
+        preds_idx = preds.argmax(axis=2)
+        preds_prob = preds.max(axis=2)
+        text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
+        if batch is None:
+            return text
+        label = batch[1]
+        label = self.decode(label.detach().cpu().numpy())
+        return text, label
+
+    def add_special_char(self, dict_character):
+        dict_character = ['</s>'] + dict_character
+        return dict_character