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

openocr/openrec/preprocess/mdiff_label_encode.py

@@ -0,0 +1,312 @@
+import random
+import numpy as np
+
+from openrec.preprocess.ctc_label_encode import BaseRecLabelEncode
+
+
+class MDiffLabelEncode(BaseRecLabelEncode):
+    """Convert between text-label and text-index."""
+
+    MASK = '<mask>'
+    EOS = '</s>'
+    PAD = '<pad>'
+
+    def __init__(self,
+                 max_text_length,
+                 character_dict_path=None,
+                 use_space_char=False,
+                 semi_ar=False,
+                 mask_tpye=[0, 1, 2, 3, 4, 5],
+                 train_all_layer=False,
+                 sample_num=1,
+                 **kwargs):
+        super(MDiffLabelEncode,
+              self).__init__(max_text_length, character_dict_path,
+                             use_space_char)
+
+        self.semi_ar = semi_ar
+        self.mask_tpye = mask_tpye  #[0, 1, 2, 3, 4, 5]
+        # 5种mask：
+        # 全部mask，纯并行
+        # 正向自回归
+        # 反向自回归
+        # block mask（1, 2，3，4，5，6步）
+        # 随机mask
+        self.train_all_layer = train_all_layer
+        if not self.train_all_layer:
+            self.sample_num = 1
+        else:
+            self.sample_num = sample_num
+
+    def random_mask(self, text):
+        l = len(text)
+        p_mask = random.random()
+
+        noisy_batch = text[:]
+        masked_indices = [False] * l
+        none_pad_indices = []
+        for i in range(l):
+            if random.random() < p_mask and text[i] != self.dict[self.PAD]:
+                noisy_batch[i] = self.dict[self.MASK]
+                masked_indices[i] = True
+            if text[i] != self.dict[self.PAD]:
+                none_pad_indices.append(i)
+            if noisy_batch[i] == self.dict[self.PAD]:
+                noisy_batch[i] = self.dict[self.MASK]
+
+        if not any(masked_indices) and len(none_pad_indices) > 0:
+            idx = random.choice(none_pad_indices)
+            noisy_batch[idx] = self.dict[self.MASK]
+            masked_indices[idx] = True
+        return noisy_batch, masked_indices
+
+    def full_mask(self, text):
+        noisy_batch = [self.dict[self.MASK]] * (self.max_text_len + 1)
+        masked_indices = [True] * len(text) + [False] * (self.max_text_len +
+                                                         1 - len(text))
+        return noisy_batch, masked_indices
+
+    def left_to_right_mask(self, text):
+        rand_split = random.randint(1, len(text) - 1)
+        noisy_batch = text[:rand_split] + [self.dict[self.MASK]] * (
+            self.max_text_len + 1 - rand_split)
+        masked_indices = [False] * rand_split + [True] * (len(
+            text) - rand_split) + [False] * (self.max_text_len + 1 - len(text))
+        return noisy_batch, masked_indices
+
+    def semi_left_to_right_mask(self, text, block_size=5, step=3, i=0):
+
+        text = text + [self.dict[self.PAD]
+                       ] * (self.max_text_len + 1 - len(text))
+        noisy_batch_step_i = []
+        masked_indices_step_i = []
+
+        if i == 0:
+            for t in text[:block_size]:
+                if t == self.dict[self.PAD]:
+                    masked_indices_step_i.append(False)
+                else:
+                    masked_indices_step_i.append(True)
+
+            noisy_batch_step_i = [self.dict[self.MASK]] * block_size + [
+                self.dict[self.PAD]
+            ] * (self.max_text_len + 1)
+            masked_indices_step_i = masked_indices_step_i + [False] * (
+                self.max_text_len + 1 - block_size)
+        elif i == 1:
+
+            noisy_batch_r, masked_indices_r = self.random_mask(
+                text[:block_size])
+            noisy_batch_step_i = noisy_batch_r
+            masked_indices_step_i = masked_indices_r
+            if i == step - 1:
+                for t in text[i * block_size:]:
+                    if t == self.dict[self.PAD]:
+                        masked_indices_step_i.append(False)
+                    else:
+                        masked_indices_step_i.append(True)
+                    noisy_batch_step_i.append(self.dict[self.MASK])
+            else:
+                for t in text[i * block_size:(i + 1) * block_size]:
+                    if t == self.dict[self.PAD]:
+                        masked_indices_step_i.append(False)
+                    else:
+                        masked_indices_step_i.append(True)
+                    noisy_batch_step_i.append(self.dict[self.MASK])
+
+            noisy_batch_step_i = noisy_batch_step_i + [self.dict[self.PAD]] * (
+                self.max_text_len + 1 - len(noisy_batch_step_i))
+            masked_indices_step_i = masked_indices_step_i + [False] * (
+                self.max_text_len + 1 - len(noisy_batch_step_i))
+        elif i >= 2:
+
+            for t in text[:(i - 1) * block_size]:
+                if t == self.dict[self.PAD]:
+                    noisy_batch_step_i.append(self.dict[self.MASK])
+                else:
+                    noisy_batch_step_i.append(text[t])
+            masked_indices_step_i = [False] * ((i - 1) * block_size)
+
+            noisy_batch_r, masked_indices_r = self.random_mask(
+                text[(i - 1) * block_size:i * block_size])
+
+            noisy_batch_step_i = noisy_batch_step_i + noisy_batch_r
+            masked_indices_step_i = masked_indices_step_i + masked_indices_r
+            if i == step - 1:
+                for t in text[i * block_size:]:
+                    if t == self.dict[self.PAD]:
+                        masked_indices_step_i.append(False)
+                    else:
+                        masked_indices_step_i.append(True)
+                    noisy_batch_step_i.append(self.dict[self.MASK])
+            else:
+                for t in text[i * block_size:(i + 1) * block_size]:
+                    if t == self.dict[self.PAD]:
+                        masked_indices_step_i.append(False)
+                    else:
+                        masked_indices_step_i.append(True)
+                    noisy_batch_step_i.append(self.dict[self.MASK])
+
+            noisy_batch_step_i = noisy_batch_step_i + [self.dict[self.PAD]] * (
+                self.max_text_len + 1 - len(noisy_batch_step_i))
+            masked_indices_step_i = masked_indices_step_i + [False] * (
+                self.max_text_len + 1 - len(masked_indices_step_i))
+        return noisy_batch_step_i, masked_indices_step_i
+
+    def forward_process_semi_ar(self, text):
+
+        step = 5
+        block_size = (self.max_text_len + 1) // step
+        noisy_batch_semi_ar = []
+        masked_indices_semi_ar = []
+        for i in range(step):
+            noisy_batch_step_i, masked_indices_step_i = self.semi_left_to_right_mask(
+                text, block_size=block_size, step=step, i=i)
+            noisy_batch_semi_ar.append(noisy_batch_step_i)
+            masked_indices_semi_ar.append(masked_indices_step_i)
+
+        rd_step = random.choice([2, 3, 4, 6, 7, 8])
+        block_size = (self.max_text_len + 1) // rd_step
+        rd_step_i = random.randint(0, rd_step - 1)
+        noisy_batch, masked_indices = self.semi_left_to_right_mask(
+            text, block_size=block_size, step=rd_step, i=rd_step_i)
+        noisy_batch_semi_ar.append(noisy_batch)
+        masked_indices_semi_ar.append(masked_indices)
+
+        # 随机将text中的部分token mask掉
+        noisy_batch, masked_indices = self.random_mask(text)
+        noisy_batch = noisy_batch + [self.dict[self.MASK]
+                                     ] * (self.max_text_len + 1 - len(text))
+        masked_indices = masked_indices + [False] * (self.max_text_len + 1 -
+                                                     len(text))
+
+        noisy_batch_semi_ar.append(noisy_batch)
+        masked_indices_semi_ar.append(masked_indices)
+        return noisy_batch_semi_ar, noisy_batch_semi_ar
+
+    def right_to_left_mask(self, text):
+        rand_split = random.randint(1, len(text) - 1)
+        noisy_batch = [self.dict[self.MASK]
+                       ] * rand_split + text[rand_split:] + [
+                           self.dict[self.MASK]
+                       ] * (self.max_text_len + 1 - len(text))
+        masked_indices = [True] * rand_split + [False] * (len(
+            text) - rand_split) + [False] * (self.max_text_len + 1 - len(text))
+        return noisy_batch, masked_indices
+
+    def forward_process(self, text):
+
+        rand_choice = random.choice(self.mask_tpye)
+        if rand_choice == 0:  # 并行mask full mask
+            return self.full_mask(text)
+        elif rand_choice == 1 and len(text) > 2:  # 正向自回归 right mask
+            return self.left_to_right_mask(text)
+        elif rand_choice == 2 and len(text) > 2:  # 反向自回归 left mask
+            return self.right_to_left_mask(text)
+        elif rand_choice == 3 and len(text) > 2:  # block mask
+            rand_step = min(random.randint(2, 6), len(text))
+            if rand_step <= 1:  # len(text) <= 1
+                return self.full_mask(text)
+            block_size = len(text) // rand_step
+            if block_size == 1:
+                return self.left_to_right_mask(text) if random.random(
+                ) < 0.5 else self.right_to_left_mask(text)
+            # 余数处理
+            if len(text) % rand_step != 0:
+                rand_step += 1
+            # 选择一个随机的block_size
+            rand_step_from_mask = random.randint(2, rand_step)
+            if rand_step == 2:
+                rand_step_from_mask = 1
+            else:
+                rand_step_from_mask = random.randint(2, rand_step)
+
+            noisy_batch = text[:block_size * (rand_step_from_mask - 1)]
+            masked_indices = [False] * (block_size * (rand_step_from_mask - 1))
+
+            noisy_batch = noisy_batch + [self.dict[self.MASK]] * (
+                self.max_text_len + 1 - len(noisy_batch))
+            masked_indices = masked_indices + [True] * (
+                len(text) - block_size *
+                (rand_step_from_mask - 1)) + [False] * (self.max_text_len + 1 -
+                                                        len(text))
+            return noisy_batch, masked_indices
+        elif rand_choice == 4 and len(text) > 2:  # cloze mask
+            noisy_batch = text[:]
+            masked_indices = [False] * len(text)
+            rand_index = random.randint(0, len(text) - 1)
+            noisy_batch[rand_index] = self.dict[self.MASK]
+            masked_indices[rand_index] = True
+            noisy_batch = noisy_batch + [self.dict[self.MASK]] * (
+                self.max_text_len + 1 - len(text))
+            masked_indices = masked_indices + [False] * (self.max_text_len +
+                                                         1 - len(text))
+            return noisy_batch, masked_indices
+        else:  # random mask
+            # 随机将text中的部分token mask掉
+            noisy_batch, masked_indices = self.random_mask(text)
+            noisy_batch = noisy_batch + [self.dict[self.MASK]] * (
+                self.max_text_len + 1 - len(text))
+            masked_indices = masked_indices + [False] * (self.max_text_len +
+                                                         1 - len(text))
+            return noisy_batch, masked_indices
+
+    def reflect_random_idices(self, text, eps=1e-3):
+        l = len(text)
+        t = random.random()
+        p_mask = (1 - eps) * t + eps
+        reflect_ids = text[:]
+        for i in range(l):
+            if random.random() < p_mask:
+                reflect_ids[i] = random.randint(0, len(self.dict) - 1)
+        return reflect_ids
+
+    def __call__(self, data):
+        text = data['label']
+        text = self.encode(text)
+        if text is None:
+            return None
+        data['length'] = np.array(len(text) + 1)
+        text = text + [self.dict[self.EOS]]
+        p_mask_list = []
+        noisy_batch_list = []
+        masked_indices_list = []
+        reflect_ids_list = []
+        for i in range(self.sample_num):
+            reflect_ids = self.reflect_random_idices(text)
+            reflect_ids = reflect_ids + [self.dict[self.MASK]] * (
+                self.max_text_len + 1 - len(reflect_ids))
+            if self.semi_ar:
+                noisy_batch, masked_indices = self.forward_process_semi_ar(
+                    text)
+            else:
+                noisy_batch, masked_indices = self.forward_process(text)
+                p_mask = float(sum(masked_indices)) / float(len(text))
+                p_mask_list.append(np.array(p_mask))
+            noisy_batch_list.append(np.array(noisy_batch))
+            masked_indices_list.append(np.array(masked_indices))
+            reflect_ids_list.append(np.array(reflect_ids))
+
+        if not self.semi_ar:
+            data['p_mask'] = np.array(
+                p_mask_list) if self.train_all_layer else np.array(
+                    p_mask_list[0])
+        data['noisy_batch'] = np.array(
+            noisy_batch_list) if self.train_all_layer else np.array(
+                noisy_batch_list[0])
+        data['masked_indices'] = np.array(
+            masked_indices_list) if self.train_all_layer else np.array(
+                masked_indices_list[0])
+        data['reflect_ids'] = np.array(
+            reflect_ids_list) if self.train_all_layer else np.array(
+                reflect_ids_list[0])
+
+        text = text + [self.dict[self.PAD]
+                       ] * (self.max_text_len + 1 - len(text))
+        data['label'] = np.array(text)
+
+        return data
+
+    def add_special_char(self, dict_character):
+        dict_character = [self.EOS] + dict_character + [self.MASK, self.PAD]
+        return dict_character

openocr/openrec/preprocess/rec_aug.py

@@ -2,13 +2,13 @@ import random
 
 import cv2
 import numpy as np
-from PIL import Image
-from .parseq_aug import rand_augment_transform
+from PIL import Image, ImageOps
 
 
 class PARSeqAugPIL(object):
 
     def __init__(self, **kwargs):
+        from .parseq_aug import rand_augment_transform
         self.transforms = rand_augment_transform()
 
     def __call__(self, data):
@@ -21,6 +21,7 @@ class PARSeqAugPIL(object):
 class PARSeqAug(object):
 
     def __init__(self, **kwargs):
+        from .parseq_aug import rand_augment_transform
         self.transforms = rand_augment_transform()
 
     def __call__(self, data):
@@ -149,6 +150,131 @@ class BaseDataAugmentation(object):
         return data
 
 
+class DocAug(object):
+
+    def __init__(self, **kwargs):
+        import albumentations as A
+        self.aug = A.Compose(
+            [
+                A.RandomShadow(
+                    shadow_roi=(0, 0.7, 1, 1), num_shadows_upper=2, p=0.5),
+                # 透视变换，模拟拍照角度，这个计算量稍大，但效果关键
+                A.OneOf(
+                    [
+                        A.Perspective(
+                            scale=(0.05, 0.1),  # 控制形变程度
+                            keep_size=False,
+                            fit_output=True,  # 保持输出大小一致
+                            pad_mode=cv2.BORDER_CONSTANT,
+                            pad_val=[255, 255, 255],  # 白色填充
+                            p=0.3  # 50% 概率应用
+                        ),
+                        A.GridDistortion(distort_limit=0.1,
+                                         border_mode=0,
+                                         interpolation=3,
+                                         value=[255, 255, 255],
+                                         p=0.3),
+                        # 安全版弹性变换
+                        A.ElasticTransform(
+                            alpha=0.5,  # 减小形变强度
+                            sigma=30,  # 增大平滑度
+                            border_mode=cv2.BORDER_REPLICATE,  # 使用边缘复制代替裁剪
+                            p=0.3)
+                    ],
+                    p=0.3),
+                A.RGBShift(r_shift_limit=15,
+                           g_shift_limit=15,
+                           b_shift_limit=15,
+                           p=0.3),
+
+                # 光学特性模拟
+                A.OneOf(
+                    [
+                        A.MotionBlur(blur_limit=(3, 5), p=0.7),
+                        A.GaussianBlur(blur_limit=(3, 5), p=0.7),
+                        # A.GlassBlur(sigma=0.05, max_delta=1, iterations=1, p=0.2)
+                    ],
+                    p=0.6),
+
+                # 色彩空间变换
+                A.RandomBrightnessContrast(brightness_limit=(-0.2, 0.2),
+                                           contrast_limit=(-0.1, 0.1),
+                                           p=0.5),
+                A.RandomGamma(gamma_limit=(80, 120), p=0.3),
+                A.CLAHE(clip_limit=2.0, p=0.2),
+
+                # 压缩伪影
+                A.ImageCompression(quality_lower=5, quality_upper=90, p=0.8),
+
+                # 传感器噪声模拟
+                A.OneOf([
+                    A.GaussNoise(var_limit=(10.0, 30.0), p=0.5),
+                    A.ISONoise(
+                        color_shift=(0.01, 0.05), intensity=(0.1, 0.3), p=0.3),
+                    A.MultiplicativeNoise(multiplier=(0.9, 1.1), p=0.2)
+                ],
+                        p=0.8),
+                A.RandomFog(fog_coef_lower=0.1, fog_coef_upper=0.3, p=0.5),
+
+                # 随机分辨率调整（保持长宽比）
+                A.Downscale(scale_min=0.5,
+                            scale_max=0.8,
+                            interpolation=cv2.INTER_LINEAR,
+                            p=0.25),
+                A.PixelDropout(dropout_prob=0.01, p=0.2)
+            ],
+            p=1)  # 100%的概率应用整个流程
+
+    def biased_random_int(self, max_value: int, exponent: float = 3.0) -> int:
+        """
+        返回 [0, max_value] 之间的一个整数，值越大概率越小。
+        exponent > 1 越大表示越偏向小值。
+        """
+        r = random.random()  # 均匀分布 [0,1)
+        biased = r**exponent  # 变成偏小
+        return int(biased * max_value)
+
+    def random_pad(self, img: Image.Image) -> Image.Image:
+        w, h = img.size
+
+        # 最大可 pad 宽度和高度
+        max_pad_w = int(w * 0.2)
+        max_pad_h = int(h * 0.2)
+
+        # 随机分配左右 padding
+        pad_left = self.biased_random_int(max_pad_w)
+        pad_right = self.biased_random_int(max_pad_w)
+
+        # 随机分配上下 padding
+        pad_top = self.biased_random_int(max_pad_h)
+        pad_bottom = self.biased_random_int(max_pad_h)
+
+        # 应用 padding
+        padded_img = ImageOps.expand(img,
+                                     border=(pad_left, pad_top, pad_right,
+                                             pad_bottom),
+                                     fill=(255, 255, 255))
+        return padded_img
+
+    def __call__(self, data):
+        image = data['image']
+        arxiv = data['arxiv']
+        # 执行数据增强
+        try:
+            if random.random() < 0.2 and arxiv:
+                # 随机填充
+                image = self.random_pad(image)
+            if not isinstance(image, np.ndarray):
+                image = np.array(image)
+
+            augmented = self.aug(image=image)['image']
+            data['image'] = Image.fromarray(augmented)
+        except Exception as e:
+            print(e)
+
+        return data
+
+
 def hsv_aug(img):
     """cvtColor."""
     hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

openocr/openrec/preprocess/resize.py

@@ -304,6 +304,63 @@ class RecDynamicResize(object):
         return data
 
 
+class NaSizeResize(object):
+
+    def __init__(self,
+                 max_side=[64 * 15, 64 * 22],
+                 divided_factor=[64, 64],
+                 **kwargs):
+        from torchvision import transforms as T
+        from torchvision.transforms import functional as F
+        self.F = F
+        self.max_side = max_side
+        self.divided_factor = divided_factor
+        self.interpolation = T.InterpolationMode.BICUBIC
+        transforms = []
+        transforms.extend([
+            T.ToTensor(),
+            T.Normalize(0.5, 0.5),
+        ])
+        self.transforms = T.Compose(transforms)
+
+    def resize_image(self, original_width, original_height, max_width,
+                     max_height):
+        # 计算宽高比
+        aspect_ratio = original_width / original_height
+
+        # 计算新的宽度和高度
+        if original_width > max_width or original_height > max_height:
+            if (max_width / max_height) >= aspect_ratio:
+                # 按高度限制比例
+                new_height = max_height
+                new_width = int(new_height * aspect_ratio)
+            else:
+                # 按宽度限制比例
+                new_width = max_width
+                new_height = int(new_width / aspect_ratio)
+        else:
+            # 如果图片已经小于或等于最大尺寸，则无需调整
+            new_width, new_height = original_width, original_height
+        return new_width, new_height
+
+    def __call__(self, data):
+        img = data['image']
+        # imgW, imgH = self.max_side
+        w, h = img.size
+        w_r, h_r = self.resize_image(w, h, self.max_side[0], self.max_side[1])
+        h_r = max(int(h_r // self.divided_factor[1] * self.divided_factor[1]),
+                  64)
+        w_r = max(int(w_r // self.divided_factor[0] * self.divided_factor[0]),
+                  64)
+        resized_image = self.F.resize(img, (h_r, w_r),
+                                      interpolation=self.interpolation)
+        img = self.transforms(resized_image)
+        valid_ratio = min(1.0, float(w_r / w))
+        data['image'] = img
+        data['valid_ratio'] = valid_ratio
+        return data
+
+
 def resize_norm_img_slice(
     img,
     image_shape,

openocr/openrec/preprocess/unirec_label_encode.py

@@ -0,0 +1,62 @@
+import numpy as np
+from openrec.preprocess.ctc_label_encode import BaseRecLabelEncode
+
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+
+class UniRecLabelEncode(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,
+                 vlmocr=False,
+                 tokenizer_path='./galactica',
+                 **kwargs):
+        super(UniRecLabelEncode,
+              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.
+        if vlmocr:
+            self.padding_idx = 1
+            self.eos_idx = 2
+            self.bos_idx = 0
+        else:
+            self.padding_idx = 0
+            self.eos_idx = 2
+            self.bos_idx = 1
+        self.batch_max_length = max_text_length + 3
+        from transformers import AutoTokenizer  # transformers==4.2.1
+        self.bpe_tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
+    def __call__(self, data):
+        text = data['label']
+        bpe_text, length = self.bpe_encode(text)
+        if bpe_text is None:
+            return None
+        data['label'] = np.array(bpe_text)
+        data['length'] = np.array(length)
+        return data
+
+    def add_special_char(self, dict_character):
+        dict_character = self.list_token + dict_character
+        return dict_character
+
+    def bpe_encode(self, text):
+        if len(text) == 0:
+            return None, None
+        token = self.bpe_tokenizer(text)['input_ids']
+        length = len(token)
+        text_list = [self.bos_idx] + token + [2]
+        if len(text_list) == 0 or len(text_list) > self.batch_max_length:
+            return None, None
+        text_list = text_list + [self.padding_idx
+                                 ] * (self.batch_max_length - len(text_list))
+        return text_list, length