openocr-python 0.0.2__py3-none-any.whl

openocr/openrec/preprocess/igtr_label_encode.py

@@ -0,0 +1,360 @@
+import copy
+import random
+
+import numpy as np
+
+from openrec.preprocess.ctc_label_encode import BaseRecLabelEncode
+
+
+class IGTRLabelEncode(BaseRecLabelEncode):
+    """Convert between text-label and text-index."""
+
+    def __init__(self,
+                 max_text_length,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 k=1,
+                 ch=False,
+                 prompt_error=False,
+                 **kwargs):
+        super(IGTRLabelEncode,
+              self).__init__(max_text_length, character_dict_path,
+                             use_space_char)
+        self.ignore_index = self.dict['<pad>']
+        self.k = k
+        self.prompt_error = prompt_error
+        self.ch = ch
+        rare_file = kwargs.get('rare_file', None)
+        siml_file = kwargs.get('siml_file', None)
+        siml_char_dict = {}
+        siml_char_list = [0 for _ in range(self.num_character)]
+        if siml_file is not None:
+            with open(siml_file, 'r') as f:
+                for lin in f.readlines():
+                    lin_s = lin.strip().split('\t')
+                    char_siml = lin_s[0]
+                    if char_siml in self.dict:
+                        siml_list = []
+                        siml_prob = []
+                        for i in range(1, len(lin_s), 2):
+                            c = lin_s[i]
+                            prob = int(lin_s[i + 1])
+                            if c in self.dict and prob >= 1:
+                                siml_list.append(self.dict[c])
+                                siml_prob.append(prob)
+                        siml_prob = np.array(siml_prob,
+                                             dtype=np.float32) / sum(siml_prob)
+                        siml_char_dict[self.dict[char_siml]] = [
+                            siml_list, siml_prob.tolist()
+                        ]
+                        siml_char_list[self.dict[char_siml]] = 1
+        self.siml_char_dict = siml_char_dict
+        self.siml_char_list = siml_char_list
+
+        rare_char_list = [0 for _ in range(self.num_character)]
+        if rare_file is not None:
+            with open(rare_file, 'r') as f:
+                for lin in f.readlines():
+                    lin_s = lin.strip().split('\t')
+                    # print(lin_s)
+                    char_rare = lin_s[0]
+                    num_appear = int(lin_s[1])
+                    if char_rare in self.dict and num_appear < 1000:
+                        rare_char_list[self.dict[char_rare]] = 1
+
+        self.rare_char_list = rare_char_list  # [self.dict[char] for char in rare_char_list]
+
+    def __call__(self, data):
+        text = data['label']  # coffee
+
+        encoder_result = self.encode(text)
+        if encoder_result is None:
+            return None
+
+        text, text_char_num, ques_list_s, prompt_list_s = encoder_result
+
+        if len(text) > self.max_text_len:
+            return None
+        data['length'] = np.array(len(text))
+
+        text = [self.dict['<s>']] + text + [self.dict['</s>']]
+        text = text + [self.dict['<pad>']
+                       ] * (self.max_text_len + 2 - len(text))
+        data['label'] = np.array(text)  # 6
+
+        ques_len_list = []
+        ques2_len_list = []
+        prompt_len_list = []
+
+        prompt_pos_idx_list = []
+        prompt_char_idx_list = []
+        ques_pos_idx_list = []
+        ques1_answer_list = []
+        ques2_char_idx_list = []
+        ques2_answer_list = []
+        ques4_char_num_list = []
+        train_step = 0
+        for prompt_list, ques_list in zip(prompt_list_s, ques_list_s):
+
+            prompt_len = len(prompt_list) + 1
+            prompt_len_list.append(prompt_len)
+            prompt_list = np.array(
+                [[0, self.dict['<s>'], 0]] + prompt_list +
+                [[self.max_text_len + 2, self.dict['<pad>'], 0]] *
+                (self.max_text_len - len(prompt_list)))
+            prompt_pos_idx_list.append(prompt_list[:, 0])
+            prompt_char_idx_list.append(prompt_list[:, 1])
+
+            ques_len = len(ques_list)
+            ques_len_list.append(ques_len)
+
+            ques_list = np.array(
+                ques_list + [[self.max_text_len + 2, self.dict['<pad>'], 0]] *
+                (self.max_text_len + 1 - ques_len))
+            ques_pos_idx_list.append(ques_list[:, 0])
+            # what is the first and third char?
+            # Is the first character 't'? and  Is the third character 'f'?
+            # How many 'c', 's' and 'f' are there in the text image?
+            ques1_answer_list.append(ques_list[:, 1])
+            ques2_char_idx = copy.deepcopy(ques_list[:ques_len, :2])
+            new_ques2_char_idx = []
+            ques2_answer = []
+            for q_2, ques2_idx in enumerate(ques2_char_idx.tolist()):
+
+                if (train_step == 2 or train_step == 3) and q_2 == ques_len - 1:
+                    new_ques2_char_idx.append(ques2_idx)
+                    ques2_answer.append(1)
+                    continue
+                if ques2_idx[1] != self.dict['<pad>'] and random.random() > 0.5:
+                    select_idx = random.randint(0, self.num_character - 3)
+                    new_ques2_char_idx.append([ques2_idx[0], select_idx])
+                    if select_idx == ques2_idx[1]:
+                        ques2_answer.append(1)
+                    else:
+                        ques2_answer.append(0)
+
+                    if self.siml_char_list[
+                            ques2_idx[1]] == 1 and random.random() > 0.5:
+                        select_idx_sim_list = random.sample(
+                            self.siml_char_dict[ques2_idx[1]][0],
+                            min(3, len(self.siml_char_dict[ques2_idx[1]][0])),
+                        )
+                        for select_idx in select_idx_sim_list:
+                            new_ques2_char_idx.append(
+                                [ques2_idx[0], select_idx])
+                            if select_idx == ques2_idx[1]:
+                                ques2_answer.append(1)
+                            else:
+                                ques2_answer.append(0)
+                else:
+                    new_ques2_char_idx.append(ques2_idx)
+                    ques2_answer.append(1)
+            ques2_len_list.append(len(new_ques2_char_idx))
+            ques2_char_idx_new = np.array(
+                new_ques2_char_idx +
+                [[self.max_text_len + 2, self.dict['<pad>']]] *
+                (self.max_text_len * 4 + 1 - len(new_ques2_char_idx)))
+            ques2_answer = np.array(
+                ques2_answer + [0] *
+                (self.max_text_len * 4 + 1 - len(ques2_answer)))
+            ques2_char_idx_list.append(ques2_char_idx_new)
+            ques2_answer_list.append(ques2_answer)
+
+            ques4_char_num_list.append(ques_list[:, 2])
+            train_step += 1
+
+        data['ques_len_list'] = np.array(ques_len_list, dtype=np.int64)
+        data['ques2_len_list'] = np.array(ques2_len_list, dtype=np.int64)
+        data['prompt_len_list'] = np.array(prompt_len_list, dtype=np.int64)
+
+        data['prompt_pos_idx_list'] = np.array(prompt_pos_idx_list,
+                                               dtype=np.int64)
+        data['prompt_char_idx_list'] = np.array(prompt_char_idx_list,
+                                                dtype=np.int64)
+        data['ques_pos_idx_list'] = np.array(ques_pos_idx_list, dtype=np.int64)
+        data['ques1_answer_list'] = np.array(ques1_answer_list, dtype=np.int64)
+        data['ques2_char_idx_list'] = np.array(ques2_char_idx_list,
+                                               dtype=np.int64)
+        data['ques2_answer_list'] = np.array(ques2_answer_list,
+                                             dtype=np.float32)
+
+        data['ques3_answer'] = np.array(
+            text_char_num,
+            dtype=np.int64)  # np.array([1, 0, 2]) # answer 1, 0, 2
+        data['ques4_char_num_list'] = np.array(ques4_char_num_list)
+
+        return data
+
+    def add_special_char(self, dict_character):
+        dict_character = ['</s>'] + dict_character + ['<s>'] + ['<pad>']
+        self.num_character = len(dict_character)
+
+        return dict_character
+
+    def encode(self, text):
+        """convert text-label into text-index.
+        input:
+            text: text labels of each image. [batch_size]
+
+        output:
+            text: concatenated text index for CTCLoss.
+                    [sum(text_lengths)] = [text_index_0 + text_index_1 + ... + text_index_(n - 1)]
+            length: length of each text. [batch_size]
+        """
+        if len(text) == 0:
+            return None
+        if self.lower:
+            text = text.lower()
+        char_num = [0 for _ in range(self.num_character - 2)]
+        char_num[0] = 1
+        text_list = []
+        qa_text = []
+        pos_i = 0
+        rare_char_qa = []
+        unrare_char_qa = []
+        for char in text:
+            if char not in self.dict:
+                continue
+
+            char_id = self.dict[char]
+            text_list.append(char_id)
+            qa_text.append([pos_i + 1, char_id, char_num[char_id]])
+            if self.rare_char_list[char_id] == 1:
+                rare_char_qa.append([pos_i + 1, char_id, char_num[char_id]])
+            else:
+                unrare_char_qa.append([pos_i + 1, char_id, char_num[char_id]])
+            char_num[char_id] += 1
+            pos_i += 1
+
+        if self.ch:
+            char_num_ch = []
+            char_num_ch_none = []
+            rare_char_num_ch_none = []
+            for i, num in enumerate(char_num):
+                if self.rare_char_list[i] == 1:
+                    rare_char_num_ch_none.append([i, num])
+                if num > 0:
+                    char_num_ch.append([i, num])
+                else:
+                    char_num_ch_none.append([i, 0])
+            none_char_index = random.sample(
+                char_num_ch_none,
+                min(37 - len(char_num_ch), len(char_num_ch_none)))
+            if len(rare_char_num_ch_none) > 0:
+                none_rare_char_index = random.sample(
+                    rare_char_num_ch_none,
+                    min(40 - len(char_num_ch) - len(none_char_index),
+                        len(rare_char_num_ch_none)),
+                )
+                char_num_ch = char_num_ch + none_char_index + none_rare_char_index
+            else:
+                char_num_ch = char_num_ch + none_char_index
+            char_num_ch.sort(key=lambda x: x[0])
+            char_num = char_num_ch
+
+        len_ = len(text_list)
+        if len_ == 0:
+            return None
+        ques_list = [
+            qa_text + [[pos_i + 1, self.dict['</s>'], 0]],
+            [[pos_i + 1, self.dict['</s>'], 0]],
+        ]
+        prompt_list = [qa_text[len_:], qa_text]
+        if len_ == 1:
+            ques_list.append([[self.max_text_len + 1, self.dict['</s>'], 0]])
+            prompt_list.append(
+                [[self.max_text_len + 2, self.dict['<pad>'], 0]] * 4 + qa_text)
+            for _ in range(1, self.k):
+                ques_list.append(
+                    [[self.max_text_len + 2, self.dict['<pad>'], 0]])
+                prompt_list.append(qa_text[1:])
+        else:
+
+            next_id = random.sample(range(1, len_ + 1), 2)
+            for slice_id in next_id:
+                b_i = slice_id - 5 if slice_id - 5 > 0 else 0
+                if slice_id == len_:
+                    ques_list.append(
+                        [[self.max_text_len + 1, self.dict['</s>'], 0]])
+                else:
+                    ques_list.append(
+                        qa_text[slice_id:] +
+                        [[self.max_text_len + 1, qa_text[slice_id][1], 0]])
+                prompt_list.append(
+                    [[self.max_text_len + 2, self.dict['<pad>'], 0]] *
+                    (5 - slice_id + b_i) + qa_text[b_i:slice_id])
+
+            shuffle_id1 = random.sample(range(1, len_),
+                                        2) if len_ > 2 else [1, 0]
+            for slice_id in shuffle_id1:
+                if slice_id == 0:
+                    ques_list.append(
+                        [[self.max_text_len + 2, self.dict['<pad>'], 0]])
+                    prompt_list.append(qa_text[:0])
+                else:
+                    ques_list.append(qa_text[slice_id:] +
+                                     [[pos_i + 1, self.dict['</s>'], 0]])
+                    prompt_list.append(qa_text[:slice_id])
+
+            if len_ > 2:
+                shuffle_id2 = random.sample(
+                    range(1, len_),
+                    self.k - 4 if len_ - 1 > self.k - 4 else len_ - 1)
+                if self.k - 4 != len(shuffle_id2):
+                    shuffle_id2 += random.sample(range(1, len_),
+                                                 self.k - 4 - len(shuffle_id2))
+                rare_slice_id = len(rare_char_qa)
+                unrare_slice_id = len(unrare_char_qa)
+                for slice_id in shuffle_id2:
+                    random.shuffle(qa_text)
+                    if len(rare_char_qa) > 0 and random.random() < 0.5:
+                        ques_list.append(rare_char_qa[:rare_slice_id] +
+                                         unrare_char_qa[unrare_slice_id:] +
+                                         [[pos_i + 1, self.dict['</s>'], 0]])
+                        if len(unrare_char_qa[:unrare_slice_id]) > 0:
+                            prompt_list1 = random.sample(
+                                unrare_char_qa[:unrare_slice_id],
+                                random.randint(
+                                    1, len(unrare_char_qa[:unrare_slice_id]))
+                                if len(unrare_char_qa[:unrare_slice_id]) > 1
+                                else 1,
+                            )
+                        else:
+                            prompt_list1 = []
+                        if len(rare_char_qa[rare_slice_id:]) > 0:
+                            prompt_list2 = random.sample(
+                                rare_char_qa[rare_slice_id:],
+                                random.randint(
+                                    1,
+                                    len(rare_char_qa[rare_slice_id:])
+                                    if len(rare_char_qa[rare_slice_id:]) > 1
+                                    else 1,
+                                ),
+                            )
+                        else:
+                            prompt_list2 = []
+                        prompt_list.append(prompt_list1 + prompt_list2)
+                        random.shuffle(rare_char_qa)
+                        random.shuffle(unrare_char_qa)
+                        rare_slice_id = random.randint(
+                            1,
+                            len(rare_char_qa)) if len(rare_char_qa) > 1 else 1
+                        unrare_slice_id = random.randint(
+                            1, len(unrare_char_qa)
+                        ) if len(unrare_char_qa) > 1 else 1
+                    else:
+                        ques_list.append(qa_text[slice_id:] +
+                                         [[pos_i + 1, self.dict['</s>'], 0]])
+                        prompt_list.append(qa_text[:slice_id])
+            else:
+                ques_list.append(qa_text[1:] +
+                                 [[pos_i + 1, self.dict['</s>'], 0]])
+                prompt_list.append(qa_text[:1])
+                ques_list.append(qa_text[:1] +
+                                 [[pos_i + 1, self.dict['</s>'], 0]])
+                prompt_list.append(qa_text[1:])
+                ques_list += [[[self.max_text_len + 2, self.dict['<pad>'], 0]]
+                              ] * (self.k - 6)
+                prompt_list += [qa_text[:0]] * (self.k - 6)
+
+        return text_list, char_num, ques_list, prompt_list

openocr/openrec/preprocess/mgp_label_encode.py

@@ -0,0 +1,95 @@
+'''
+This code is refer from:
+https://github.com/AlibabaResearch/AdvancedLiterateMachinery/blob/main/OCR/MGP-STR
+'''
+import numpy as np
+
+from openrec.preprocess.ctc_label_encode import BaseRecLabelEncode
+
+
+class MGPLabelEncode(BaseRecLabelEncode):
+    """ Convert between text-label and text-index """
+    SPACE = '[s]'
+    GO = '[GO]'
+    list_token = [GO, SPACE]
+
+    def __init__(self,
+                 max_text_length,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 only_char=False,
+                 **kwargs):
+        super(MGPLabelEncode,
+              self).__init__(max_text_length, character_dict_path,
+                             use_space_char)
+        # character (str): set of the possible characters.
+        # [GO] for the start token of the attention decoder. [s] for end-of-sentence token.
+
+        self.batch_max_length = max_text_length + len(self.list_token)
+        self.only_char = only_char
+        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, data):
+        text = data['label']
+        char_text, char_len = self.encode(text)
+        if char_text is None:
+            return None
+        data['length'] = np.array(char_len)
+        data['char_label'] = np.array(char_text)
+        if self.only_char:
+            return data
+        bpe_text = self.bpe_encode(text)
+        if bpe_text is None:
+            return None
+        wp_text = self.wp_encode(text)
+        data['bpe_label'] = np.array(bpe_text)
+        data['wp_label'] = wp_text
+        return data
+
+    def add_special_char(self, dict_character):
+        dict_character = self.list_token + dict_character
+        return dict_character
+
+    def encode(self, text):
+        """ convert text-label into text-index.
+        """
+        if len(text) == 0:
+            return None, None
+        if self.lower:
+            text = text.lower()
+        length = len(text)
+        text = [self.GO] + list(text) + [self.SPACE]
+        text_list = []
+        for char in text:
+            if char not in self.dict:
+                continue
+            text_list.append(self.dict[char])
+        if len(text_list) == 0 or len(text_list) > self.batch_max_length:
+            return None, None
+        text_list = text_list + [self.dict[self.GO]
+                                 ] * (self.batch_max_length - len(text_list))
+        return text_list, length
+
+    def bpe_encode(self, text):
+        if len(text) == 0:
+            return None
+        token = self.bpe_tokenizer(text)['input_ids']
+        text_list = [1] + token + [2]
+        if len(text_list) == 0 or len(text_list) > self.batch_max_length:
+            return None
+        text_list = text_list + [self.dict[self.GO]
+                                 ] * (self.batch_max_length - len(text_list))
+        return text_list
+
+    def wp_encode(self, text):
+        wp_target = self.wp_tokenizer([text],
+                                      padding='max_length',
+                                      max_length=self.batch_max_length,
+                                      truncation=True,
+                                      return_tensors='np')
+        return wp_target['input_ids'][0]

openocr/openrec/preprocess/parseq_aug.py

@@ -0,0 +1,150 @@
+# Scene Text Recognition Model Hub
+# Copyright 2022 Darwin Bautista
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import partial
+
+import imgaug.augmenters as iaa
+import numpy as np
+from PIL import Image, ImageFilter
+
+from openrec.preprocess import auto_augment
+from openrec.preprocess.auto_augment import _LEVEL_DENOM, LEVEL_TO_ARG, NAME_TO_OP, _randomly_negate, rotate
+
+
+def rotate_expand(img, degrees, **kwargs):
+    """Rotate operation with expand=True to avoid cutting off the
+    characters."""
+    kwargs['expand'] = True
+    return rotate(img, degrees, **kwargs)
+
+
+def _level_to_arg(level, hparams, key, default):
+    magnitude = hparams.get(key, default)
+    level = (level / _LEVEL_DENOM) * magnitude
+    level = _randomly_negate(level)
+    return level,
+
+
+def apply():
+    # Overrides
+    NAME_TO_OP.update({'Rotate': rotate_expand})
+    LEVEL_TO_ARG.update({
+        'Rotate':
+        partial(_level_to_arg, key='rotate_deg', default=30.),
+        'ShearX':
+        partial(_level_to_arg, key='shear_x_pct', default=0.3),
+        'ShearY':
+        partial(_level_to_arg, key='shear_y_pct', default=0.3),
+        'TranslateXRel':
+        partial(_level_to_arg, key='translate_x_pct', default=0.45),
+        'TranslateYRel':
+        partial(_level_to_arg, key='translate_y_pct', default=0.45),
+    })
+
+
+apply()
+
+_OP_CACHE = {}
+
+
+def _get_op(key, factory):
+    try:
+        op = _OP_CACHE[key]
+    except KeyError:
+        op = factory()
+        _OP_CACHE[key] = op
+    return op
+
+
+def _get_param(level, img, max_dim_factor, min_level=1):
+    max_level = max(min_level, max_dim_factor * max(img.size))
+    return round(min(level, max_level))
+
+
+def gaussian_blur(img, radius, **__):
+    radius = _get_param(radius, img, 0.02)
+    key = 'gaussian_blur_' + str(radius)
+    op = _get_op(key, lambda: ImageFilter.GaussianBlur(radius))
+    return img.filter(op)
+
+
+def motion_blur(img, k, **__):
+    k = _get_param(k, img, 0.08, 3) | 1  # bin to odd values
+    key = 'motion_blur_' + str(k)
+    op = _get_op(key, lambda: iaa.MotionBlur(k))
+    return Image.fromarray(op(image=np.asarray(img)))
+
+
+def gaussian_noise(img, scale, **_):
+    scale = _get_param(scale, img, 0.25) | 1  # bin to odd values
+    key = 'gaussian_noise_' + str(scale)
+    op = _get_op(key, lambda: iaa.AdditiveGaussianNoise(scale=scale))
+    return Image.fromarray(op(image=np.asarray(img)))
+
+
+def poisson_noise(img, lam, **_):
+    lam = _get_param(lam, img, 0.2) | 1  # bin to odd values
+    key = 'poisson_noise_' + str(lam)
+    op = _get_op(key, lambda: iaa.AdditivePoissonNoise(lam))
+    return Image.fromarray(op(image=np.asarray(img)))
+
+
+def _level_to_arg(level, _hparams, max):
+    level = max * level / auto_augment._LEVEL_DENOM
+    return level,
+
+
+_RAND_TRANSFORMS = auto_augment._RAND_INCREASING_TRANSFORMS.copy()
+_RAND_TRANSFORMS.remove(
+    'SharpnessIncreasing')  # remove, interferes with *blur ops
+_RAND_TRANSFORMS.extend([
+    'GaussianBlur',
+    # 'MotionBlur',
+    # 'GaussianNoise',
+    'PoissonNoise'
+])
+auto_augment.LEVEL_TO_ARG.update({
+    'GaussianBlur':
+    partial(_level_to_arg, max=4),
+    'MotionBlur':
+    partial(_level_to_arg, max=20),
+    'GaussianNoise':
+    partial(_level_to_arg, max=0.1 * 255),
+    'PoissonNoise':
+    partial(_level_to_arg, max=40)
+})
+auto_augment.NAME_TO_OP.update({
+    'GaussianBlur': gaussian_blur,
+    'MotionBlur': motion_blur,
+    'GaussianNoise': gaussian_noise,
+    'PoissonNoise': poisson_noise
+})
+
+
+def rand_augment_transform(magnitude=5, num_layers=3):
+    # These are tuned for magnitude=5, which means that effective magnitudes are half of these values.
+    hparams = {
+        'rotate_deg': 30,
+        'shear_x_pct': 0.9,
+        'shear_y_pct': 0.2,
+        'translate_x_pct': 0.10,
+        'translate_y_pct': 0.30
+    }
+    ra_ops = auto_augment.rand_augment_ops(magnitude,
+                                           hparams=hparams,
+                                           transforms=_RAND_TRANSFORMS)
+    # Supply weights to disable replacement in random selection (i.e. avoid applying the same op twice)
+    choice_weights = [1. / len(ra_ops) for _ in range(len(ra_ops))]
+    return auto_augment.RandAugment(ra_ops, num_layers, choice_weights)