openocr-python 0.0.9__py3-none-any.whl → 0.1.0.dev0__py3-none-any.whl

openocr/openrec/postprocess/smtr_postprocess.py

@@ -33,7 +33,7 @@ class SMTRLabelDecode(BaseRecLabelDecode):
         if batch is None:
             return text
         label = batch[1]
-        label = self.decode(label[:, 1:].detach().cpu().numpy())
+        label = self.decode(label[:, 1:])
         return text, label
 
     def add_special_char(self, dict_character):

openocr/openrec/postprocess/srn_postprocess.py

@@ -71,7 +71,7 @@ class SRNLabelDecode(BaseRecLabelDecode):
         if batch is None:
             return text
 
-        label = batch[1].cpu().numpy()
+        label = batch[1]
         # print(f"label.shape:{label.shape}")
         label = self.decode(label, is_remove_duplicate=False)
         return text, label

openocr/openrec/postprocess/unirec_postprocess.py

@@ -0,0 +1,58 @@
+import re
+from .ctc_postprocess import BaseRecLabelDecode
+
+rules = [
+    (r'-<\|sn\|>', ''),
+    (r' <\|sn\|>', ' '),
+    (r'<\|sn\|>', ' '),
+    (r'<\|unk\|>', ''),
+    (r'<s>', ''),
+    (r'</s>', ''),
+    (r'\uffff', ''),
+    (r'_{4,}', '___'),
+    (r'\.{4,}', '...'),
+]
+
+
+def clean_special_tokens(text):
+    text = text.replace('Ġ',
+                        ' ').replace('Ċ', '\n').replace('<|bos|>', '').replace(
+                            '<|eos|>', '').replace('<|pad|>', '')
+    for rule in rules:
+        text = re.sub(rule[0], rule[1], text)
+    return text
+
+
+class UniRecLabelDecode(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,
+                 tokenizer_path='./configs/rec/unirec/unirec-0.1b',
+                 **kwargs):
+        super(UniRecLabelDecode, self).__init__(character_dict_path,
+                                                use_space_char)
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+    def __call__(self, preds, batch=None, *args, **kwargs):
+        result_list = []
+        pred_ids = preds
+        res = [
+            ''.join(self.tokenizer.convert_ids_to_tokens(seq.tolist()))
+            for seq in pred_ids
+        ]
+        for i in range(len(res)):
+            res[i] = clean_special_tokens(res[i])
+            result_list.append(
+                (res[i],
+                 0.0))  # Assuming confidence is not available, set to 0.0
+        return result_list
+
+    def add_special_char(self, dict_character):
+        dict_character = self.list_token + dict_character
+        return dict_character

openocr/openrec/postprocess/visionlan_postprocess.py

@@ -77,5 +77,5 @@ class VisionLANLabelDecode(BaseRecLabelDecode):
             text.append((preds_text, float(preds_prob)))
         if batch is None:
             return text
-        label = self.decode(label.detach().cpu().numpy())
+        label = self.decode(label)
         return text, label

openocr/openrec/preprocess/__init__.py

@@ -101,11 +101,15 @@ MODULE_MAPPING = {
     'SRNLabelEncode': '.srn_label_encode',
     'VisionLANLabelEncode': '.visionlan_label_encode',
     'CAMLabelEncode': '.cam_label_encode',
+    'MDiffLabelEncode': '.mdiff_label_encode',
+    'UniRecLabelEncode': '.unirec_label_encode',
+    'CMERProcessor': '.cmer_label_encode',
     'ABINetAug': '.rec_aug',
     'BDA': '.rec_aug',
     'PARSeqAug': '.rec_aug',
     'PARSeqAugPIL': '.rec_aug',
     'SVTRAug': '.rec_aug',
+    'DocAug': '.rec_aug',
     'ABINetResize': '.resize',
     'CDistNetResize': '.resize',
     'LongResize': '.resize',
@@ -117,6 +121,7 @@ MODULE_MAPPING = {
     'SVTRResize': '.resize',
     'VisionLANResize': '.resize',
     'RecDynamicResize': '.resize',
+    'NaSizeResize': '.resize',
 }
 
 

openocr/openrec/preprocess/ce_label_encode.py

@@ -50,7 +50,7 @@ class BaseRecLabelEncode(object):
         text_re = []
         c_current = ''
         for c in text:
-            if not bool(re.search('[a-zA-Z0-9 :*./%+-١٢٣٤٥٦٧٨٩٠]', c)):
+            if not bool(re.search('[a-zA-Z0-9 :*./%+١٢٣٤٥٦٧٨٩٠-]', c)):
                 if c_current != '':
                     text_re.append(c_current)
                 text_re.append(c)