openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/postprocess/ctc_postprocess.py

@@ -0,0 +1,119 @@
+import re
+
+import numpy as np
+import torch
+
+
+class BaseRecLabelDecode(object):
+    """Convert between text-label and text-index."""
+
+    def __init__(self, character_dict_path=None, use_space_char=False):
+        self.beg_str = 'sos'
+        self.end_str = 'eos'
+        self.reverse = False
+        self.character_str = []
+
+        if character_dict_path is None:
+            self.character_str = '0123456789abcdefghijklmnopqrstuvwxyz'
+            dict_character = list(self.character_str)
+        else:
+            with open(character_dict_path, 'rb') as fin:
+                lines = fin.readlines()
+                for line in lines:
+                    line = line.decode('utf-8').strip('\n').strip('\r\n')
+                    self.character_str.append(line)
+            if use_space_char:
+                self.character_str.append(' ')
+            dict_character = list(self.character_str)
+            if 'arabic' in character_dict_path:
+                self.reverse = True
+
+        dict_character = self.add_special_char(dict_character)
+        self.dict = {}
+        for i, char in enumerate(dict_character):
+            self.dict[char] = i
+        self.character = dict_character
+
+    def pred_reverse(self, pred):
+        pred_re = []
+        c_current = ''
+        for c in pred:
+            if not bool(re.search('[a-zA-Z0-9 :*./%+-]', c)):
+                if c_current != '':
+                    pred_re.append(c_current)
+                pred_re.append(c)
+                c_current = ''
+            else:
+                c_current += c
+        if c_current != '':
+            pred_re.append(c_current)
+
+        return ''.join(pred_re[::-1])
+
+    def add_special_char(self, dict_character):
+        return dict_character
+
+    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
+        """convert text-index into text-label."""
+        result_list = []
+        ignored_tokens = self.get_ignored_tokens()
+        batch_size = len(text_index)
+        for batch_idx in range(batch_size):
+            selection = np.ones(len(text_index[batch_idx]), dtype=bool)
+            if is_remove_duplicate:
+                selection[1:] = text_index[batch_idx][1:] != text_index[
+                    batch_idx][:-1]
+            for ignored_token in ignored_tokens:
+                selection &= text_index[batch_idx] != ignored_token
+
+            char_list = [
+                self.character[text_id]
+                for text_id in text_index[batch_idx][selection]
+            ]
+            if text_prob is not None:
+                conf_list = text_prob[batch_idx][selection]
+            else:
+                conf_list = [1] * len(selection)
+            if len(conf_list) == 0:
+                conf_list = [0]
+
+            text = ''.join(char_list)
+
+            if self.reverse:  # for arabic rec
+                text = self.pred_reverse(text)
+
+            result_list.append((text, np.mean(conf_list).tolist()))
+        return result_list
+
+    def get_ignored_tokens(self):
+        return [0]  # for ctc blank
+
+    def get_character_num(self):
+        return len(self.character)
+
+
+class CTCLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(CTCLabelDecode, self).__init__(character_dict_path,
+                                             use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        # preds = preds['res']
+        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, is_remove_duplicate=True)
+        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 = ['blank'] + dict_character
+        return dict_character

openocr/openrec/postprocess/igtr_postprocess.py

@@ -0,0 +1,100 @@
+import numpy as np
+import torch
+
+from .nrtr_postprocess import NRTRLabelDecode
+
+
+class IGTRLabelDecode(NRTRLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(IGTRLabelDecode, self).__init__(character_dict_path,
+                                              use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        if isinstance(preds, list):
+            if isinstance(preds[0], dict):
+                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)
+            else:
+                preds_idx = preds[0].detach().cpu().numpy()
+                preds_prob = preds[1].detach().cpu().numpy()
+                text = self.decode(preds_idx,
+                                   preds_prob,
+                                   is_remove_duplicate=False)
+        else:
+            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_idx_top5 = preds.argsort(axis=2)[:, :, -5:]
+            preds_prob = preds.max(axis=2)
+            text = self.decode(preds_idx,
+                               preds_prob,
+                               is_remove_duplicate=False,
+                               idx_top5=preds_idx_top5)
+        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 + ['<s>', '<pad>']
+        return dict_character
+
+    def decode(self,
+               text_index,
+               text_prob=None,
+               is_remove_duplicate=False,
+               idx_top5=None):
+        """convert text-index into text-label."""
+        result_list = []
+        batch_size = len(text_index)
+        for batch_idx in range(batch_size):
+            char_list = []
+            char_list_top5 = []
+            conf_list = []
+            for idx in range(len(text_index[batch_idx])):
+                char_idx_top5 = []
+                try:
+                    char_idx = self.character[int(text_index[batch_idx][idx])]
+                    if idx_top5 is not None:
+                        for top5_i in idx_top5[batch_idx][idx]:
+                            char_idx_top5.append(self.character[top5_i])
+                except:
+                    continue
+                if char_idx == '</s>':  # end
+                    break
+                if char_idx == '<s>' or char_idx == '<pad>':
+                    continue
+                char_list.append(char_idx)
+                char_list_top5.append(char_idx_top5)
+                if text_prob is not None:
+                    conf_list.append(text_prob[batch_idx][idx])
+                else:
+                    conf_list.append(1)
+            text = ''.join(char_list)
+            if idx_top5 is not None:
+                result_list.append(
+                    (text, [np.mean(conf_list).tolist(), char_list_top5]))
+            else:
+                result_list.append((text, np.mean(conf_list).tolist()))
+        return result_list

openocr/openrec/postprocess/lister_postprocess.py

@@ -0,0 +1,59 @@
+import numpy as np
+import torch
+
+from openrec.postprocess.ctc_postprocess import BaseRecLabelDecode
+
+
+class LISTERLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 **kwargs):
+        super(LISTERLabelDecode, self).__init__(character_dict_path,
+                                                use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        preds = preds[1]['logits']
+        if isinstance(preds, torch.Tensor):
+            preds = preds.detach().cpu().numpy()
+        preds_idx = preds.argmax(axis=2)
+        # preds_idx_top5 = preds.argsort(axis=2)[:, :, -5:]
+        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 + ['<pad>']
+        return dict_character
+
+    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
+        """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 == '</s>':  # end
+                    break
+                if char_idx == '<s>' or char_idx == '<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/mgp_postprocess.py

@@ -0,0 +1,143 @@
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class MPGLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+    SPACE = '[s]'
+    GO = '[GO]'
+    list_token = [GO, SPACE]
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 only_char=False,
+                 **kwargs):
+        super(MPGLabelDecode, self).__init__(character_dict_path,
+                                             use_space_char)
+        self.only_char = only_char
+        self.EOS = '[s]'
+        self.PAD = '[GO]'
+        if not only_char:
+            # transformers==4.2.1
+            from transformers import BertTokenizer, GPT2Tokenizer
+            self.bpe_tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+            self.wp_tokenizer = BertTokenizer.from_pretrained(
+                'bert-base-uncased')
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        if isinstance(preds, list):
+            char_preds = preds[0].detach().cpu().numpy()
+        else:
+            char_preds = preds.detach().cpu().numpy()
+
+        preds_idx = char_preds.argmax(axis=2)
+        preds_prob = char_preds.max(axis=2)
+        char_text = self.char_decode(preds_idx[:, 1:], preds_prob[:, 1:])
+        if batch is None:
+            return char_text
+        label = batch[1]
+        label = self.char_decode(label[:, 1:].detach().cpu().numpy())
+        if self.only_char:
+            return char_text, label
+        else:
+            bpe_preds = preds[1].detach().cpu().numpy()
+            wp_preds = preds[2]
+
+            bpe_preds_idx = bpe_preds.argmax(axis=2)
+            bpe_preds_prob = bpe_preds.max(axis=2)
+            bpe_text = self.bpe_decode(bpe_preds_idx[:, 1:],
+                                       bpe_preds_prob[:, 1:])
+
+            wp_preds = wp_preds.detach()  #.cpu().numpy()
+            wp_preds_prob, wp_preds_idx = wp_preds.max(-1)
+            wp_text = self.wp_decode(wp_preds_idx[:, 1:], wp_preds_prob[:, 1:])
+
+            final_text = self.final_decode(char_text, bpe_text, wp_text)
+            return char_text, bpe_text, wp_text, final_text, label
+
+    def add_special_char(self, dict_character):
+        dict_character = self.list_token + dict_character
+        return dict_character
+
+    def final_decode(self, char_text, bpe_text, wp_text):
+        result_list = []
+        for (char_pred,
+             char_pred_conf), (bpe_pred,
+                               bpe_pred_conf), (wp_pred, wp_pred_conf) in zip(
+                                   char_text, bpe_text, wp_text):
+            final_text = char_pred
+            final_prob = char_pred_conf
+            if bpe_pred_conf > final_prob:
+                final_text = bpe_pred
+                final_prob = bpe_pred_conf
+            if wp_pred_conf > final_prob:
+                final_text = wp_pred
+                final_prob = wp_pred_conf
+            result_list.append((final_text, final_prob))
+        return result_list
+
+    def char_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 = 1.0
+            for idx in range(len(text_index[batch_idx])):
+                try:
+                    char_idx = self.character[int(text_index[batch_idx][idx])]
+                except:
+                    continue
+                if text_prob is not None:
+                    conf_list *= text_prob[batch_idx][idx]
+
+                if char_idx == self.EOS:  # end
+                    break
+                if char_idx == self.PAD:
+                    continue
+                char_list.append(char_idx)
+
+            text = ''.join(char_list)
+            result_list.append((text, conf_list))
+        return result_list
+
+    def bpe_decode(self, text_index, text_prob):
+        """ convert text-index into text-label. """
+        result_list = []
+        for text, probs in zip(text_index, text_prob):
+            text_decoded = []
+            conf_list = 1.0
+            for bpeindx, prob in zip(text, probs):
+                tokenstr = self.bpe_tokenizer.decode([bpeindx])
+                if tokenstr == '#':
+                    break
+                text_decoded.append(tokenstr)
+                conf_list *= prob
+            text = ''.join(text_decoded)
+            result_list.append((text, conf_list))
+        return result_list
+
+    def wp_decode(self, text_index, text_prob=None):
+        """ convert text-index into text-label. """
+        result_list = []
+        for batch_idx, text in enumerate(text_index):
+            wp_pred = self.wp_tokenizer.decode(text)
+            wp_pred_EOS = wp_pred.find('[SEP]')
+            wp_pred = wp_pred[:wp_pred_EOS]
+            if text_prob is not None:
+                try:
+                    # print(text.cpu().tolist())
+                    wp_pred_EOS_index = text.cpu().tolist().index(102) + 1
+                except:
+                    wp_pred_EOS_index = -1
+                wp_pred_max_prob = text_prob[batch_idx][:wp_pred_EOS_index]
+                try:
+                    wp_confidence_score = wp_pred_max_prob.cumprod(
+                        dim=0)[-1].cpu().numpy().sum()
+                except:
+                    wp_confidence_score = 0.0
+            else:
+                wp_confidence_score = 1.0
+            result_list.append((wp_pred, wp_confidence_score))
+        return result_list

openocr/openrec/postprocess/nrtr_postprocess.py

@@ -0,0 +1,75 @@
+import numpy as np
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class NRTRLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 **kwargs):
+        super(NRTRLabelDecode, self).__init__(character_dict_path,
+                                              use_space_char)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        preds = preds['res']
+        if len(preds) == 2:
+            preds_id = preds[0]
+            preds_prob = preds[1]
+            if isinstance(preds_id, torch.Tensor):
+                preds_id = preds_id.detach().cpu().numpy()
+            if isinstance(preds_prob, torch.Tensor):
+                preds_prob = preds_prob.detach().cpu().numpy()
+            if preds_id[0][0] == 2:
+                preds_idx = preds_id[:, 1:]
+                preds_prob = preds_prob[:, 1:]
+            else:
+                preds_idx = preds_id
+            text = self.decode(preds_idx,
+                               preds_prob,
+                               is_remove_duplicate=False)
+            if batch is None:
+                return text
+            label = self.decode(batch[1][:, 1:].cpu().numpy())
+        else:
+            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,
+                               is_remove_duplicate=False)
+            if batch is None:
+                return text
+            label = self.decode(batch[1][:, 1:].cpu().numpy())
+        return text, 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, is_remove_duplicate=False):
+        """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 == '</s>':  # end
+                    break
+                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/smtr_postprocess.py

@@ -0,0 +1,73 @@
+import numpy as np
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class SMTRLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    BOS = '<s>'
+    EOS = '</s>'
+    IN_F = '<INF>'  # ignore
+    IN_B = '<INB>'  # ignore
+    PAD = '<pad>'
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=True,
+                 next_mode=True,
+                 **kwargs):
+        super(SMTRLabelDecode, self).__init__(character_dict_path,
+                                              use_space_char)
+        self.next_mode = next_mode
+
+    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, is_remove_duplicate=False)
+        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.IN_F, self.IN_B, self.PAD
+        ]
+        self.num_character = len(dict_character)
+        return dict_character
+
+    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
+        """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 == '</s>':  # end
+                    break
+                if char_idx == '<s>' or char_idx == '<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)
+            if self.next_mode or text_prob is None:
+                text = ''.join(char_list)
+            else:
+                text = ''.join(char_list[::-1])
+            result_list.append((text, np.mean(conf_list).tolist()))
+        return result_list

openocr/openrec/postprocess/srn_postprocess.py

@@ -0,0 +1,80 @@
+import numpy as np
+import torch
+
+from .ctc_postprocess import BaseRecLabelDecode
+
+
+class SRNLabelDecode(BaseRecLabelDecode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 **kwargs):
+        super(SRNLabelDecode, self).__init__(character_dict_path,
+                                             use_space_char)
+        self.max_len = 25
+
+    def add_special_char(self, dict_character):
+        dict_character = dict_character + ['<BOS>', '<EOS>']
+        self.start_idx = len(dict_character) - 2
+        self.end_idx = len(dict_character) - 1
+        return dict_character
+
+    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
+        """convert text-index into text-label."""
+        result_list = []
+        ignored_tokens = self.get_ignored_tokens()
+        # [B,25]
+        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])):
+                # print(f"text_index[{batch_idx}][{idx}]:{text_index[batch_idx][idx]}")
+                if text_index[batch_idx][idx] in ignored_tokens:
+                    continue
+                if int(text_index[batch_idx][idx]) == int(self.end_idx):
+                    if text_prob is None and idx == 0:
+                        continue
+                    else:
+                        break
+                if is_remove_duplicate:
+                    # only for predict
+                    if idx > 0 and text_index[batch_idx][
+                            idx - 1] == text_index[batch_idx][idx]:
+                        continue
+                char_list.append(self.character[int(
+                    text_index[batch_idx][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
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+
+        if isinstance(preds, torch.Tensor):
+            preds = preds.reshape([-1, self.max_len, preds.shape[-1]])
+            preds = preds.detach().cpu().numpy()
+        else:
+            preds = preds[-1]
+            preds = preds.reshape([-1, self.max_len,
+                                   preds.shape[-1]]).detach().cpu().numpy()
+
+        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].cpu().numpy()
+        # print(f"label.shape:{label.shape}")
+        label = self.decode(label, is_remove_duplicate=False)
+        return text, label
+
+    def get_ignored_tokens(self):
+        return [self.start_idx, self.end_idx]