first commit

ppocr/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.

ppocr/data/__init__.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import os
+import sys
+import numpy as np
+import skimage
+import paddle
+import signal
+import random
+
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.abspath(os.path.join(__dir__, "../..")))
+
+import copy
+from paddle.io import Dataset, DataLoader, BatchSampler, DistributedBatchSampler
+import paddle.distributed as dist
+
+from ppocr.data.imaug import transform, create_operators
+from ppocr.data.simple_dataset import SimpleDataSet, MultiScaleDataSet
+from ppocr.data.lmdb_dataset import LMDBDataSet, LMDBDataSetSR, LMDBDataSetTableMaster
+from ppocr.data.pgnet_dataset import PGDataSet
+from ppocr.data.pubtab_dataset import PubTabDataSet
+from ppocr.data.multi_scale_sampler import MultiScaleSampler
+from ppocr.data.latexocr_dataset import LaTeXOCRDataSet
+
+# for PaddleX dataset_type
+TextDetDataset = SimpleDataSet
+TextRecDataset = SimpleDataSet
+MSTextRecDataset = MultiScaleDataSet
+PubTabTableRecDataset = PubTabDataSet
+KieDataset = SimpleDataSet
+LaTeXOCRDataSet = LaTeXOCRDataSet
+
+__all__ = ["build_dataloader", "transform", "create_operators", "set_signal_handlers"]
+
+
+def term_mp(sig_num, frame):
+    """kill all child processes"""
+    pid = os.getpid()
+    pgid = os.getpgid(os.getpid())
+    print("main proc {} exit, kill process group " "{}".format(pid, pgid))
+    os.killpg(pgid, signal.SIGKILL)
+
+
+def set_signal_handlers():
+    pid = os.getpid()
+    try:
+        pgid = os.getpgid(pid)
+    except AttributeError:
+        # In case `os.getpgid` is not available, no signal handler will be set,
+        # because we cannot do safe cleanup.
+        pass
+    else:
+        # XXX: `term_mp` kills all processes in the process group, which in
+        # some cases includes the parent process of current process and may
+        # cause unexpected results. To solve this problem, we set signal
+        # handlers only when current process is the group leader. In the
+        # future, it would be better to consider killing only descendants of
+        # the current process.
+        if pid == pgid:
+            # support exit using ctrl+c
+            signal.signal(signal.SIGINT, term_mp)
+            signal.signal(signal.SIGTERM, term_mp)
+
+
+def build_dataloader(config, mode, device, logger, seed=None):
+    config = copy.deepcopy(config)
+
+    support_dict = [
+        "SimpleDataSet",
+        "LMDBDataSet",
+        "PGDataSet",
+        "PubTabDataSet",
+        "LMDBDataSetSR",
+        "LMDBDataSetTableMaster",
+        "MultiScaleDataSet",
+        "TextDetDataset",
+        "TextRecDataset",
+        "MSTextRecDataset",
+        "PubTabTableRecDataset",
+        "KieDataset",
+        "LaTeXOCRDataSet",
+    ]
+    module_name = config[mode]["dataset"]["name"]
+    assert module_name in support_dict, Exception(
+        "DataSet only support {}".format(support_dict)
+    )
+    assert mode in ["Train", "Eval", "Test"], "Mode should be Train, Eval or Test."
+
+    dataset = eval(module_name)(config, mode, logger, seed)
+    loader_config = config[mode]["loader"]
+    batch_size = loader_config["batch_size_per_card"]
+    drop_last = loader_config["drop_last"]
+    shuffle = loader_config["shuffle"]
+    num_workers = loader_config["num_workers"]
+    if "use_shared_memory" in loader_config.keys():
+        use_shared_memory = loader_config["use_shared_memory"]
+    else:
+        use_shared_memory = True
+
+    if mode == "Train":
+        # Distribute data to multiple cards
+        if "sampler" in config[mode]:
+            config_sampler = config[mode]["sampler"]
+            sampler_name = config_sampler.pop("name")
+            batch_sampler = eval(sampler_name)(dataset, **config_sampler)
+        else:
+            batch_sampler = DistributedBatchSampler(
+                dataset=dataset,
+                batch_size=batch_size,
+                shuffle=shuffle,
+                drop_last=drop_last,
+            )
+    else:
+        # Distribute data to single card
+        batch_sampler = BatchSampler(
+            dataset=dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last
+        )
+
+    if "collate_fn" in loader_config:
+        from . import collate_fn
+
+        collate_fn = getattr(collate_fn, loader_config["collate_fn"])()
+    else:
+        collate_fn = None
+    data_loader = DataLoader(
+        dataset=dataset,
+        batch_sampler=batch_sampler,
+        places=device,
+        num_workers=num_workers,
+        return_list=True,
+        use_shared_memory=use_shared_memory,
+        collate_fn=collate_fn,
+    )
+
+    return data_loader

ppocr/data/collate_fn.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+import paddle
+import numbers
+import numpy as np
+from collections import defaultdict
+
+
+class DictCollator(object):
+    """
+    data batch
+    """
+
+    def __call__(self, batch):
+        # todo：support batch operators
+        data_dict = defaultdict(list)
+        to_tensor_keys = []
+        for sample in batch:
+            for k, v in sample.items():
+                if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)):
+                    if k not in to_tensor_keys:
+                        to_tensor_keys.append(k)
+                data_dict[k].append(v)
+        for k in to_tensor_keys:
+            data_dict[k] = paddle.to_tensor(data_dict[k])
+        return data_dict
+
+
+class ListCollator(object):
+    """
+    data batch
+    """
+
+    def __call__(self, batch):
+        # todo：support batch operators
+        data_dict = defaultdict(list)
+        to_tensor_idxs = []
+        for sample in batch:
+            for idx, v in enumerate(sample):
+                if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)):
+                    if idx not in to_tensor_idxs:
+                        to_tensor_idxs.append(idx)
+                data_dict[idx].append(v)
+        for idx in to_tensor_idxs:
+            data_dict[idx] = paddle.to_tensor(data_dict[idx])
+        return list(data_dict.values())
+
+
+class SSLRotateCollate(object):
+    """
+    bach: [
+        [(4*3xH*W), (4,)]
+        [(4*3xH*W), (4,)]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+        output = [np.concatenate(d, axis=0) for d in zip(*batch)]
+        return output
+
+
+class DyMaskCollator(object):
+    """
+    batch: [
+        image [batch_size, channel, maxHinbatch, maxWinbatch]
+        image_mask [batch_size, channel, maxHinbatch, maxWinbatch]
+        label [batch_size, maxLabelLen]
+        label_mask [batch_size, maxLabelLen]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+        max_width, max_height, max_length = 0, 0, 0
+        bs, channel = len(batch), batch[0][0].shape[0]
+        proper_items = []
+        for item in batch:
+            if (
+                item[0].shape[1] * max_width > 1600 * 320
+                or item[0].shape[2] * max_height > 1600 * 320
+            ):
+                continue
+            max_height = (
+                item[0].shape[1] if item[0].shape[1] > max_height else max_height
+            )
+            max_width = item[0].shape[2] if item[0].shape[2] > max_width else max_width
+            max_length = len(item[1]) if len(item[1]) > max_length else max_length
+            proper_items.append(item)
+
+        images, image_masks = np.zeros(
+            (len(proper_items), channel, max_height, max_width), dtype="float32"
+        ), np.zeros((len(proper_items), 1, max_height, max_width), dtype="float32")
+        labels, label_masks = np.zeros(
+            (len(proper_items), max_length), dtype="int64"
+        ), np.zeros((len(proper_items), max_length), dtype="int64")
+
+        for i in range(len(proper_items)):
+            _, h, w = proper_items[i][0].shape
+            images[i][:, :h, :w] = proper_items[i][0]
+            image_masks[i][:, :h, :w] = 1
+            l = len(proper_items[i][1])
+            labels[i][:l] = proper_items[i][1]
+            label_masks[i][:l] = 1
+
+        return images, image_masks, labels, label_masks
+
+
+class LaTeXOCRCollator(object):
+    """
+    batch: [
+        image [batch_size, channel, maxHinbatch, maxWinbatch]
+        label [batch_size, maxLabelLen]
+        label_mask [batch_size, maxLabelLen]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+        images, labels, attention_mask = batch[0]
+        return images, labels, attention_mask
+
+
+class UniMERNetCollator(object):
+    """
+    batch: [
+        image [batch_size, channel, maxHinbatch, maxWinbatch]
+        image_mask [batch_size, channel, maxHinbatch, maxWinbatch]
+        label [batch_size, maxLabelLen]
+        label_mask [batch_size, maxLabelLen]
+        ...
+    ]
+    """
+
+    def __call__(self, batch):
+
+        max_width, max_height, max_length = 0, 0, 0
+        bs, channel = len(batch), batch[0][0].shape[0]
+        proper_items = []
+        for item in batch:
+            max_height = (
+                item[0].shape[1] if item[0].shape[1] > max_height else max_height
+            )
+            max_width = item[0].shape[2] if item[0].shape[2] > max_width else max_width
+            max_length = len(item[1]) if len(item[1]) > max_length else max_length
+            proper_items.append(item)
+
+        images = np.ones(
+            (len(proper_items), channel, max_height, max_width), dtype="float32"
+        )
+
+        labels, label_masks = np.ones(
+            (len(proper_items), max_length), dtype="int64"
+        ), np.zeros((len(proper_items), max_length), dtype="int64")
+        for i in range(len(proper_items)):
+            _, h, w = proper_items[i][0].shape
+            images[i][:, :h, :w] = proper_items[i][0]
+            l = len(proper_items[i][1])
+            labels[i][:l] = proper_items[i][1]
+            label_masks[i][:l] = proper_items[i][2]
+        return images, labels, label_masks

ppocr/data/imaug/ColorJitter.py

@@ -0,0 +1,27 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 paddle.vision.transforms import ColorJitter as pp_ColorJitter
+
+__all__ = ["ColorJitter"]
+
+
+class ColorJitter(object):
+    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0, **kwargs):
+        self.aug = pp_ColorJitter(brightness, contrast, saturation, hue)
+
+    def __call__(self, data):
+        image = data["image"]
+        image = self.aug(image)
+        data["image"] = image
+        return data

ppocr/data/imaug/__init__.py

@@ -0,0 +1,96 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from .iaa_augment import IaaAugment
+from .make_border_map import MakeBorderMap
+from .make_shrink_map import MakeShrinkMap
+from .random_crop_data import EastRandomCropData, RandomCropImgMask
+from .make_pse_gt import MakePseGt
+
+
+from .rec_img_aug import (
+    BaseDataAugmentation,
+    RecAug,
+    RecConAug,
+    RecResizeImg,
+    ClsResizeImg,
+    SRNRecResizeImg,
+    GrayRecResizeImg,
+    SARRecResizeImg,
+    PRENResizeImg,
+    ABINetRecResizeImg,
+    SVTRRecResizeImg,
+    ABINetRecAug,
+    VLRecResizeImg,
+    SPINRecResizeImg,
+    RobustScannerRecResizeImg,
+    RFLRecResizeImg,
+    SVTRRecAug,
+    ParseQRecAug,
+)
+from .ssl_img_aug import SSLRotateResize
+from .randaugment import RandAugment
+from .copy_paste import CopyPaste
+from .ColorJitter import ColorJitter
+from .operators import *
+from .label_ops import *
+
+from .east_process import *
+from .sast_process import *
+from .pg_process import *
+from .table_ops import *
+
+from .vqa import *
+
+from .fce_aug import *
+from .fce_targets import FCENetTargets
+from .ct_process import *
+from .drrg_targets import DRRGTargets
+from .latex_ocr_aug import *
+from .unimernet_aug import *
+
+
+def transform(data, ops=None):
+    """transform"""
+    if ops is None:
+        ops = []
+    for op in ops:
+        data = op(data)
+        if data is None:
+            return None
+    return data
+
+
+def create_operators(op_param_list, global_config=None):
+    """
+    create operators based on the config
+
+    Args:
+        params(list): a dict list, used to create some operators
+    """
+    assert isinstance(op_param_list, list), "operator config should be a list"
+    ops = []
+    for operator in op_param_list:
+        assert isinstance(operator, dict) and len(operator) == 1, "yaml format error"
+        op_name = list(operator)[0]
+        param = {} if operator[op_name] is None else operator[op_name]
+        if global_config is not None:
+            param.update(global_config)
+        op = eval(op_name)(**param)
+        ops.append(op)
+    return ops

ppocr/data/imaug/abinet_aug.py

@@ -0,0 +1,512 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/FangShancheng/ABINet/blob/main/transforms.py
+"""
+import math
+import numbers
+import random
+
+import cv2
+import numpy as np
+from paddle.vision.transforms import Compose, ColorJitter
+
+
+def sample_asym(magnitude, size=None):
+    return np.random.beta(1, 4, size) * magnitude
+
+
+def sample_sym(magnitude, size=None):
+    return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude
+
+
+def sample_uniform(low, high, size=None):
+    return np.random.uniform(low, high, size=size)
+
+
+def get_interpolation(type="random"):
+    if type == "random":
+        choice = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA]
+        interpolation = choice[random.randint(0, len(choice) - 1)]
+    elif type == "nearest":
+        interpolation = cv2.INTER_NEAREST
+    elif type == "linear":
+        interpolation = cv2.INTER_LINEAR
+    elif type == "cubic":
+        interpolation = cv2.INTER_CUBIC
+    elif type == "area":
+        interpolation = cv2.INTER_AREA
+    else:
+        raise TypeError(
+            "Interpolation types only nearest, linear, cubic, area are supported!"
+        )
+    return interpolation
+
+
+class CVRandomRotation(object):
+    def __init__(self, degrees=15):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+    @staticmethod
+    def get_params(degrees):
+        return sample_sym(degrees)
+
+    def __call__(self, img):
+        angle = self.get_params(self.degrees)
+        src_h, src_w = img.shape[:2]
+        M = cv2.getRotationMatrix2D(
+            center=(src_w / 2, src_h / 2), angle=angle, scale=1.0
+        )
+        abs_cos, abs_sin = abs(M[0, 0]), abs(M[0, 1])
+        dst_w = int(src_h * abs_sin + src_w * abs_cos)
+        dst_h = int(src_h * abs_cos + src_w * abs_sin)
+        M[0, 2] += (dst_w - src_w) / 2
+        M[1, 2] += (dst_h - src_h) / 2
+
+        flags = get_interpolation()
+        return cv2.warpAffine(
+            img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE
+        )
+
+
+class CVRandomAffine(object):
+    def __init__(self, degrees, translate=None, scale=None, shear=None):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+        if translate is not None:
+            assert (
+                isinstance(translate, (tuple, list)) and len(translate) == 2
+            ), "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert (
+                isinstance(scale, (tuple, list)) and len(scale) == 2
+            ), "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError(
+                        "If shear is a single number, it must be positive."
+                    )
+                self.shear = [shear]
+            else:
+                assert isinstance(shear, (tuple, list)) and (
+                    len(shear) == 2
+                ), "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+
+    def _get_inverse_affine_matrix(self, center, angle, translate, scale, shear):
+        # https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717
+        from numpy import sin, cos, tan
+
+        if isinstance(shear, numbers.Number):
+            shear = [shear, 0]
+
+        if not isinstance(shear, (tuple, list)) and len(shear) == 2:
+            raise ValueError(
+                "Shear should be a single value or a tuple/list containing "
+                + "two values. Got {}".format(shear)
+            )
+
+        rot = math.radians(angle)
+        sx, sy = [math.radians(s) for s in shear]
+
+        cx, cy = center
+        tx, ty = translate
+
+        # RSS without scaling
+        a = cos(rot - sy) / cos(sy)
+        b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
+        c = sin(rot - sy) / cos(sy)
+        d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
+
+        # Inverted rotation matrix with scale and shear
+        # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
+        M = [d, -b, 0, -c, a, 0]
+        M = [x / scale for x in M]
+
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty)
+        M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty)
+
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        M[2] += cx
+        M[5] += cy
+        return M
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, height):
+        angle = sample_sym(degrees)
+        if translate is not None:
+            max_dx = translate[0] * height
+            max_dy = translate[1] * height
+            translations = (np.round(sample_sym(max_dx)), np.round(sample_sym(max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = sample_uniform(scale_ranges[0], scale_ranges[1])
+        else:
+            scale = 1.0
+
+        if shears is not None:
+            if len(shears) == 1:
+                shear = [sample_sym(shears[0]), 0.0]
+            elif len(shears) == 2:
+                shear = [sample_sym(shears[0]), sample_sym(shears[1])]
+        else:
+            shear = 0.0
+
+        return angle, translations, scale, shear
+
+    def __call__(self, img):
+        src_h, src_w = img.shape[:2]
+        angle, translate, scale, shear = self.get_params(
+            self.degrees, self.translate, self.scale, self.shear, src_h
+        )
+
+        M = self._get_inverse_affine_matrix(
+            (src_w / 2, src_h / 2), angle, (0, 0), scale, shear
+        )
+        M = np.array(M).reshape(2, 3)
+
+        startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), (0, src_h - 1)]
+        project = lambda x, y, a, b, c: int(a * x + b * y + c)
+        endpoints = [
+            (project(x, y, *M[0]), project(x, y, *M[1])) for x, y in startpoints
+        ]
+
+        rect = cv2.minAreaRect(np.array(endpoints))
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int32)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+
+        dst_w = int(max_x - min_x)
+        dst_h = int(max_y - min_y)
+        M[0, 2] += (dst_w - src_w) / 2
+        M[1, 2] += (dst_h - src_h) / 2
+
+        # add translate
+        dst_w += int(abs(translate[0]))
+        dst_h += int(abs(translate[1]))
+        if translate[0] < 0:
+            M[0, 2] += abs(translate[0])
+        if translate[1] < 0:
+            M[1, 2] += abs(translate[1])
+
+        flags = get_interpolation()
+        return cv2.warpAffine(
+            img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE
+        )
+
+
+class CVRandomPerspective(object):
+    def __init__(self, distortion=0.5):
+        self.distortion = distortion
+
+    def get_params(self, width, height, distortion):
+        offset_h = sample_asym(distortion * height / 2, size=4).astype(dtype=np.int32)
+        offset_w = sample_asym(distortion * width / 2, size=4).astype(dtype=np.int32)
+        topleft = (offset_w[0], offset_h[0])
+        topright = (width - 1 - offset_w[1], offset_h[1])
+        botright = (width - 1 - offset_w[2], height - 1 - offset_h[2])
+        botleft = (offset_w[3], height - 1 - offset_h[3])
+
+        startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), (0, height - 1)]
+        endpoints = [topleft, topright, botright, botleft]
+        return np.array(startpoints, dtype=np.float32), np.array(
+            endpoints, dtype=np.float32
+        )
+
+    def __call__(self, img):
+        height, width = img.shape[:2]
+        startpoints, endpoints = self.get_params(width, height, self.distortion)
+        M = cv2.getPerspectiveTransform(startpoints, endpoints)
+
+        # TODO: more robust way to crop image
+        rect = cv2.minAreaRect(endpoints)
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int32)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+        min_x, min_y = max(min_x, 0), max(min_y, 0)
+
+        flags = get_interpolation()
+        img = cv2.warpPerspective(
+            img, M, (max_x, max_y), flags=flags, borderMode=cv2.BORDER_REPLICATE
+        )
+        img = img[min_y:, min_x:]
+        return img
+
+
+class CVRescale(object):
+    def __init__(self, factor=4, base_size=(128, 512)):
+        """Define image scales using gaussian pyramid and rescale image to target scale.
+
+        Args:
+            factor: the decayed factor from base size, factor=4 keeps target scale by default.
+            base_size: base size the build the bottom layer of pyramid
+        """
+        if isinstance(factor, numbers.Number):
+            self.factor = round(sample_uniform(0, factor))
+        elif isinstance(factor, (tuple, list)) and len(factor) == 2:
+            self.factor = round(sample_uniform(factor[0], factor[1]))
+        else:
+            raise Exception("factor must be number or list with length 2")
+        # assert factor is valid
+        self.base_h, self.base_w = base_size[:2]
+
+    def __call__(self, img):
+        if self.factor == 0:
+            return img
+        src_h, src_w = img.shape[:2]
+        cur_w, cur_h = self.base_w, self.base_h
+        scale_img = cv2.resize(img, (cur_w, cur_h), interpolation=get_interpolation())
+        for _ in range(self.factor):
+            scale_img = cv2.pyrDown(scale_img)
+        scale_img = cv2.resize(
+            scale_img, (src_w, src_h), interpolation=get_interpolation()
+        )
+        return scale_img
+
+
+class CVGaussianNoise(object):
+    def __init__(self, mean=0, var=20):
+        self.mean = mean
+        if isinstance(var, numbers.Number):
+            self.var = max(int(sample_asym(var)), 1)
+        elif isinstance(var, (tuple, list)) and len(var) == 2:
+            self.var = int(sample_uniform(var[0], var[1]))
+        else:
+            raise Exception("degree must be number or list with length 2")
+
+    def __call__(self, img):
+        noise = np.random.normal(self.mean, self.var**0.5, img.shape)
+        img = np.clip(img + noise, 0, 255).astype(np.uint8)
+        return img
+
+
+class CVPossionNoise(object):
+    def __init__(self, lam=20):
+        self.lam = lam
+        if isinstance(lam, numbers.Number):
+            self.lam = max(int(sample_asym(lam)), 1)
+        elif isinstance(lam, (tuple, list)) and len(lam) == 2:
+            self.lam = int(sample_uniform(lam[0], lam[1]))
+        else:
+            raise Exception("lam must be number or list with length 2")
+
+    def __call__(self, img):
+        noise = np.random.poisson(lam=self.lam, size=img.shape)
+        img = np.clip(img + noise, 0, 255).astype(np.uint8)
+        return img
+
+
+class CVGaussionBlur(object):
+    def __init__(self, radius):
+        self.radius = radius
+        if isinstance(radius, numbers.Number):
+            self.radius = max(int(sample_asym(radius)), 1)
+        elif isinstance(radius, (tuple, list)) and len(radius) == 2:
+            self.radius = int(sample_uniform(radius[0], radius[1]))
+        else:
+            raise Exception("radius must be number or list with length 2")
+
+    def __call__(self, img):
+        fil = cv2.getGaussianKernel(ksize=self.radius, sigma=1, ktype=cv2.CV_32F)
+        img = cv2.sepFilter2D(img, -1, fil, fil)
+        return img
+
+
+class CVMotionBlur(object):
+    def __init__(self, degrees=12, angle=90):
+        if isinstance(degrees, numbers.Number):
+            self.degree = max(int(sample_asym(degrees)), 1)
+        elif isinstance(degrees, (tuple, list)) and len(degrees) == 2:
+            self.degree = int(sample_uniform(degrees[0], degrees[1]))
+        else:
+            raise Exception("degree must be number or list with length 2")
+        self.angle = sample_uniform(-angle, angle)
+
+    def __call__(self, img):
+        M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), self.angle, 1)
+        motion_blur_kernel = np.zeros((self.degree, self.degree))
+        motion_blur_kernel[self.degree // 2, :] = 1
+        motion_blur_kernel = cv2.warpAffine(
+            motion_blur_kernel, M, (self.degree, self.degree)
+        )
+        motion_blur_kernel = motion_blur_kernel / self.degree
+        img = cv2.filter2D(img, -1, motion_blur_kernel)
+        img = np.clip(img, 0, 255).astype(np.uint8)
+        return img
+
+
+class CVGeometry(object):
+    def __init__(
+        self,
+        degrees=15,
+        translate=(0.3, 0.3),
+        scale=(0.5, 2.0),
+        shear=(45, 15),
+        distortion=0.5,
+        p=0.5,
+    ):
+        self.p = p
+        type_p = random.random()
+        if type_p < 0.33:
+            self.transforms = CVRandomRotation(degrees=degrees)
+        elif type_p < 0.66:
+            self.transforms = CVRandomAffine(
+                degrees=degrees, translate=translate, scale=scale, shear=shear
+            )
+        else:
+            self.transforms = CVRandomPerspective(distortion=distortion)
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return self.transforms(img)
+        else:
+            return img
+
+
+class CVDeterioration(object):
+    def __init__(self, var, degrees, factor, p=0.5):
+        self.p = p
+        transforms = []
+        if var is not None:
+            transforms.append(CVGaussianNoise(var=var))
+        if degrees is not None:
+            transforms.append(CVMotionBlur(degrees=degrees))
+        if factor is not None:
+            transforms.append(CVRescale(factor=factor))
+
+        random.shuffle(transforms)
+        transforms = Compose(transforms)
+        self.transforms = transforms
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return self.transforms(img)
+        else:
+            return img
+
+
+class CVColorJitter(object):
+    def __init__(self, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.5):
+        self.p = p
+        self.transforms = ColorJitter(
+            brightness=brightness, contrast=contrast, saturation=saturation, hue=hue
+        )
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return self.transforms(img)
+        else:
+            return img
+
+
+class SVTRDeterioration(object):
+    def __init__(self, var, degrees, factor, p=0.5):
+        self.p = p
+        transforms = []
+        if var is not None:
+            transforms.append(CVGaussianNoise(var=var))
+        if degrees is not None:
+            transforms.append(CVMotionBlur(degrees=degrees))
+        if factor is not None:
+            transforms.append(CVRescale(factor=factor))
+        self.transforms = transforms
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            random.shuffle(self.transforms)
+            transforms = Compose(self.transforms)
+            return transforms(img)
+        else:
+            return img
+
+
+class ParseQDeterioration(object):
+    def __init__(self, var, degrees, lam, radius, factor, p=0.5):
+        self.p = p
+        transforms = []
+        if var is not None:
+            transforms.append(CVGaussianNoise(var=var))
+        if degrees is not None:
+            transforms.append(CVMotionBlur(degrees=degrees))
+        if lam is not None:
+            transforms.append(CVPossionNoise(lam=lam))
+        if radius is not None:
+            transforms.append(CVGaussionBlur(radius=radius))
+        if factor is not None:
+            transforms.append(CVRescale(factor=factor))
+        self.transforms = transforms
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            random.shuffle(self.transforms)
+            transforms = Compose(self.transforms)
+            return transforms(img)
+        else:
+            return img
+
+
+class SVTRGeometry(object):
+    def __init__(
+        self,
+        aug_type=0,
+        degrees=15,
+        translate=(0.3, 0.3),
+        scale=(0.5, 2.0),
+        shear=(45, 15),
+        distortion=0.5,
+        p=0.5,
+    ):
+        self.aug_type = aug_type
+        self.p = p
+        self.transforms = []
+        self.transforms.append(CVRandomRotation(degrees=degrees))
+        self.transforms.append(
+            CVRandomAffine(
+                degrees=degrees, translate=translate, scale=scale, shear=shear
+            )
+        )
+        self.transforms.append(CVRandomPerspective(distortion=distortion))
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            if self.aug_type:
+                random.shuffle(self.transforms)
+                transforms = Compose(self.transforms[: random.randint(1, 3)])
+                img = transforms(img)
+            else:
+                img = self.transforms[random.randint(0, 2)](img)
+            return img
+        else:
+            return img

ppocr/data/imaug/copy_paste.py

@@ -0,0 +1,179 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import copy
+import cv2
+import random
+import numpy as np
+from PIL import Image
+from shapely.geometry import Polygon
+
+from ppocr.data.imaug.iaa_augment import IaaAugment
+from ppocr.data.imaug.random_crop_data import is_poly_outside_rect
+from tools.infer.utility import get_rotate_crop_image
+
+
+class CopyPaste(object):
+    def __init__(self, objects_paste_ratio=0.2, limit_paste=True, **kwargs):
+        self.ext_data_num = 1
+        self.objects_paste_ratio = objects_paste_ratio
+        self.limit_paste = limit_paste
+        augmenter_args = [{"type": "Resize", "args": {"size": [0.5, 3]}}]
+        self.aug = IaaAugment(augmenter_args)
+
+    def __call__(self, data):
+        point_num = data["polys"].shape[1]
+        src_img = data["image"]
+        src_polys = data["polys"].tolist()
+        src_texts = data["texts"]
+        src_ignores = data["ignore_tags"].tolist()
+        ext_data = data["ext_data"][0]
+        ext_image = ext_data["image"]
+        ext_polys = ext_data["polys"]
+        ext_texts = ext_data["texts"]
+        ext_ignores = ext_data["ignore_tags"]
+
+        indexes = [i for i in range(len(ext_ignores)) if not ext_ignores[i]]
+        select_num = max(1, min(int(self.objects_paste_ratio * len(ext_polys)), 30))
+
+        random.shuffle(indexes)
+        select_idxs = indexes[:select_num]
+        select_polys = ext_polys[select_idxs]
+        select_ignores = ext_ignores[select_idxs]
+
+        src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
+        ext_image = cv2.cvtColor(ext_image, cv2.COLOR_BGR2RGB)
+        src_img = Image.fromarray(src_img).convert("RGBA")
+        for idx, poly, tag in zip(select_idxs, select_polys, select_ignores):
+            box_img = get_rotate_crop_image(ext_image, poly)
+
+            src_img, box = self.paste_img(src_img, box_img, src_polys)
+            if box is not None:
+                box = box.tolist()
+                for _ in range(len(box), point_num):
+                    box.append(box[-1])
+                src_polys.append(box)
+                src_texts.append(ext_texts[idx])
+                src_ignores.append(tag)
+        src_img = cv2.cvtColor(np.array(src_img), cv2.COLOR_RGB2BGR)
+        h, w = src_img.shape[:2]
+        src_polys = np.array(src_polys)
+        src_polys[:, :, 0] = np.clip(src_polys[:, :, 0], 0, w)
+        src_polys[:, :, 1] = np.clip(src_polys[:, :, 1], 0, h)
+        data["image"] = src_img
+        data["polys"] = src_polys
+        data["texts"] = src_texts
+        data["ignore_tags"] = np.array(src_ignores)
+        return data
+
+    def paste_img(self, src_img, box_img, src_polys):
+        box_img_pil = Image.fromarray(box_img).convert("RGBA")
+        src_w, src_h = src_img.size
+        box_w, box_h = box_img_pil.size
+
+        angle = np.random.randint(0, 360)
+        box = np.array([[[0, 0], [box_w, 0], [box_w, box_h], [0, box_h]]])
+        box = rotate_bbox(box_img, box, angle)[0]
+        box_img_pil = box_img_pil.rotate(angle, expand=1)
+        box_w, box_h = box_img_pil.width, box_img_pil.height
+        if src_w - box_w < 0 or src_h - box_h < 0:
+            return src_img, None
+
+        paste_x, paste_y = self.select_coord(
+            src_polys, box, src_w - box_w, src_h - box_h
+        )
+        if paste_x is None:
+            return src_img, None
+        box[:, 0] += paste_x
+        box[:, 1] += paste_y
+        r, g, b, A = box_img_pil.split()
+        src_img.paste(box_img_pil, (paste_x, paste_y), mask=A)
+
+        return src_img, box
+
+    def select_coord(self, src_polys, box, endx, endy):
+        if self.limit_paste:
+            xmin, ymin, xmax, ymax = (
+                box[:, 0].min(),
+                box[:, 1].min(),
+                box[:, 0].max(),
+                box[:, 1].max(),
+            )
+            for _ in range(50):
+                paste_x = random.randint(0, endx)
+                paste_y = random.randint(0, endy)
+                xmin1 = xmin + paste_x
+                xmax1 = xmax + paste_x
+                ymin1 = ymin + paste_y
+                ymax1 = ymax + paste_y
+
+                num_poly_in_rect = 0
+                for poly in src_polys:
+                    if not is_poly_outside_rect(
+                        poly, xmin1, ymin1, xmax1 - xmin1, ymax1 - ymin1
+                    ):
+                        num_poly_in_rect += 1
+                        break
+                if num_poly_in_rect == 0:
+                    return paste_x, paste_y
+            return None, None
+        else:
+            paste_x = random.randint(0, endx)
+            paste_y = random.randint(0, endy)
+            return paste_x, paste_y
+
+
+def get_union(pD, pG):
+    return Polygon(pD).union(Polygon(pG)).area
+
+
+def get_intersection_over_union(pD, pG):
+    return get_intersection(pD, pG) / get_union(pD, pG)
+
+
+def get_intersection(pD, pG):
+    return Polygon(pD).intersection(Polygon(pG)).area
+
+
+def rotate_bbox(img, text_polys, angle, scale=1):
+    """
+    from https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/augment.py
+    Args:
+        img: np.ndarray
+        text_polys: np.ndarray N*4*2
+        angle: int
+        scale: int
+
+    Returns:
+
+    """
+    w = img.shape[1]
+    h = img.shape[0]
+
+    rangle = np.deg2rad(angle)
+    nw = abs(np.sin(rangle) * h) + abs(np.cos(rangle) * w)
+    nh = abs(np.cos(rangle) * h) + abs(np.sin(rangle) * w)
+    rot_mat = cv2.getRotationMatrix2D((nw * 0.5, nh * 0.5), angle, scale)
+    rot_move = np.dot(rot_mat, np.array([(nw - w) * 0.5, (nh - h) * 0.5, 0]))
+    rot_mat[0, 2] += rot_move[0]
+    rot_mat[1, 2] += rot_move[1]
+
+    # ---------------------- rotate box ----------------------
+    rot_text_polys = list()
+    for bbox in text_polys:
+        point1 = np.dot(rot_mat, np.array([bbox[0, 0], bbox[0, 1], 1]))
+        point2 = np.dot(rot_mat, np.array([bbox[1, 0], bbox[1, 1], 1]))
+        point3 = np.dot(rot_mat, np.array([bbox[2, 0], bbox[2, 1], 1]))
+        point4 = np.dot(rot_mat, np.array([bbox[3, 0], bbox[3, 1], 1]))
+        rot_text_polys.append([point1, point2, point3, point4])
+    return np.array(rot_text_polys, dtype=np.float32)

ppocr/data/imaug/ct_process.py

@@ -0,0 +1,376 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import os
+import cv2
+import paddle
+import random
+import pyclipper
+import numpy as np
+from PIL import Image
+
+import paddle.vision.transforms as transforms
+
+from ppocr.utils.utility import check_install
+
+
+class RandomScale:
+    def __init__(self, short_size=640, **kwargs):
+        self.short_size = short_size
+
+    def scale_aligned(self, img, scale):
+        oh, ow = img.shape[0:2]
+        h = int(oh * scale + 0.5)
+        w = int(ow * scale + 0.5)
+        if h % 32 != 0:
+            h = h + (32 - h % 32)
+        if w % 32 != 0:
+            w = w + (32 - w % 32)
+        img = cv2.resize(img, dsize=(w, h))
+        factor_h = h / oh
+        factor_w = w / ow
+        return img, factor_h, factor_w
+
+    def __call__(self, data):
+        img = data["image"]
+
+        h, w = img.shape[0:2]
+        random_scale = np.array([0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3])
+        scale = (np.random.choice(random_scale) * self.short_size) / min(h, w)
+        img, factor_h, factor_w = self.scale_aligned(img, scale)
+
+        data["scale_factor"] = (factor_w, factor_h)
+        data["image"] = img
+        return data
+
+
+class MakeShrink:
+    def __init__(self, kernel_scale=0.7, **kwargs):
+        self.kernel_scale = kernel_scale
+
+    def dist(self, a, b):
+        return np.linalg.norm((a - b), ord=2, axis=0)
+
+    def perimeter(self, bbox):
+        peri = 0.0
+        for i in range(bbox.shape[0]):
+            peri += self.dist(bbox[i], bbox[(i + 1) % bbox.shape[0]])
+        return peri
+
+    def shrink(self, bboxes, rate, max_shr=20):
+        check_install("Polygon", "Polygon3")
+        import Polygon as plg
+
+        rate = rate * rate
+        shrinked_bboxes = []
+        for bbox in bboxes:
+            area = plg.Polygon(bbox).area()
+            peri = self.perimeter(bbox)
+
+            try:
+                pco = pyclipper.PyclipperOffset()
+                pco.AddPath(bbox, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+                offset = min(int(area * (1 - rate) / (peri + 0.001) + 0.5), max_shr)
+
+                shrinked_bbox = pco.Execute(-offset)
+                if len(shrinked_bbox) == 0:
+                    shrinked_bboxes.append(bbox)
+                    continue
+
+                shrinked_bbox = np.array(shrinked_bbox[0])
+                if shrinked_bbox.shape[0] <= 2:
+                    shrinked_bboxes.append(bbox)
+                    continue
+
+                shrinked_bboxes.append(shrinked_bbox)
+            except Exception as e:
+                shrinked_bboxes.append(bbox)
+
+        return shrinked_bboxes
+
+    def __call__(self, data):
+        img = data["image"]
+        bboxes = data["polys"]
+        words = data["texts"]
+        scale_factor = data["scale_factor"]
+
+        gt_instance = np.zeros(img.shape[0:2], dtype="uint8")  # h,w
+        training_mask = np.ones(img.shape[0:2], dtype="uint8")
+        training_mask_distance = np.ones(img.shape[0:2], dtype="uint8")
+
+        for i in range(len(bboxes)):
+            bboxes[i] = np.reshape(
+                bboxes[i]
+                * ([scale_factor[0], scale_factor[1]] * (bboxes[i].shape[0] // 2)),
+                (bboxes[i].shape[0] // 2, 2),
+            ).astype("int32")
+
+        for i in range(len(bboxes)):
+            # different value for different bbox
+            cv2.drawContours(gt_instance, [bboxes[i]], -1, i + 1, -1)
+
+            # set training mask to 0
+            cv2.drawContours(training_mask, [bboxes[i]], -1, 0, -1)
+
+            # for not accurate annotation, use training_mask_distance
+            if words[i] == "###" or words[i] == "???":
+                cv2.drawContours(training_mask_distance, [bboxes[i]], -1, 0, -1)
+
+        # make shrink
+        gt_kernel_instance = np.zeros(img.shape[0:2], dtype="uint8")
+        kernel_bboxes = self.shrink(bboxes, self.kernel_scale)
+        for i in range(len(bboxes)):
+            cv2.drawContours(gt_kernel_instance, [kernel_bboxes[i]], -1, i + 1, -1)
+
+            # for training mask, kernel and background= 1, box region=0
+            if words[i] != "###" and words[i] != "???":
+                cv2.drawContours(training_mask, [kernel_bboxes[i]], -1, 1, -1)
+
+        gt_kernel = gt_kernel_instance.copy()
+        # for gt_kernel, kernel = 1
+        gt_kernel[gt_kernel > 0] = 1
+
+        # shrink 2 times
+        tmp1 = gt_kernel_instance.copy()
+        erode_kernel = np.ones((3, 3), np.uint8)
+        tmp1 = cv2.erode(tmp1, erode_kernel, iterations=1)
+        tmp2 = tmp1.copy()
+        tmp2 = cv2.erode(tmp2, erode_kernel, iterations=1)
+
+        # compute text region
+        gt_kernel_inner = tmp1 - tmp2
+
+        # gt_instance: text instance, bg=0, diff word use diff value
+        # training_mask: text instance mask, word=0，kernel and bg=1
+        # gt_kernel_instance: text kernel instance, bg=0, diff word use diff value
+        # gt_kernel: text_kernel, bg=0，diff word use same value
+        # gt_kernel_inner: text kernel reference
+        # training_mask_distance: word without anno = 0, else 1
+
+        data["image"] = [
+            img,
+            gt_instance,
+            training_mask,
+            gt_kernel_instance,
+            gt_kernel,
+            gt_kernel_inner,
+            training_mask_distance,
+        ]
+        return data
+
+
+class GroupRandomHorizontalFlip:
+    def __init__(self, p=0.5, **kwargs):
+        self.p = p
+
+    def __call__(self, data):
+        imgs = data["image"]
+
+        if random.random() < self.p:
+            for i in range(len(imgs)):
+                imgs[i] = np.flip(imgs[i], axis=1).copy()
+        data["image"] = imgs
+        return data
+
+
+class GroupRandomRotate:
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        imgs = data["image"]
+
+        max_angle = 10
+        angle = random.random() * 2 * max_angle - max_angle
+        for i in range(len(imgs)):
+            img = imgs[i]
+            w, h = img.shape[:2]
+            rotation_matrix = cv2.getRotationMatrix2D((h / 2, w / 2), angle, 1)
+            img_rotation = cv2.warpAffine(
+                img, rotation_matrix, (h, w), flags=cv2.INTER_NEAREST
+            )
+            imgs[i] = img_rotation
+
+        data["image"] = imgs
+        return data
+
+
+class GroupRandomCropPadding:
+    def __init__(self, target_size=(640, 640), **kwargs):
+        self.target_size = target_size
+
+    def __call__(self, data):
+        imgs = data["image"]
+
+        h, w = imgs[0].shape[0:2]
+        t_w, t_h = self.target_size
+        p_w, p_h = self.target_size
+        if w == t_w and h == t_h:
+            return data
+
+        t_h = t_h if t_h < h else h
+        t_w = t_w if t_w < w else w
+
+        if random.random() > 3.0 / 8.0 and np.max(imgs[1]) > 0:
+            # make sure to crop the text region
+            tl = np.min(np.where(imgs[1] > 0), axis=1) - (t_h, t_w)
+            tl[tl < 0] = 0
+            br = np.max(np.where(imgs[1] > 0), axis=1) - (t_h, t_w)
+            br[br < 0] = 0
+            br[0] = min(br[0], h - t_h)
+            br[1] = min(br[1], w - t_w)
+
+            i = random.randint(tl[0], br[0]) if tl[0] < br[0] else 0
+            j = random.randint(tl[1], br[1]) if tl[1] < br[1] else 0
+        else:
+            i = random.randint(0, h - t_h) if h - t_h > 0 else 0
+            j = random.randint(0, w - t_w) if w - t_w > 0 else 0
+
+        n_imgs = []
+        for idx in range(len(imgs)):
+            if len(imgs[idx].shape) == 3:
+                s3_length = int(imgs[idx].shape[-1])
+                img = imgs[idx][i : i + t_h, j : j + t_w, :]
+                img_p = cv2.copyMakeBorder(
+                    img,
+                    0,
+                    p_h - t_h,
+                    0,
+                    p_w - t_w,
+                    borderType=cv2.BORDER_CONSTANT,
+                    value=tuple(0 for i in range(s3_length)),
+                )
+            else:
+                img = imgs[idx][i : i + t_h, j : j + t_w]
+                img_p = cv2.copyMakeBorder(
+                    img,
+                    0,
+                    p_h - t_h,
+                    0,
+                    p_w - t_w,
+                    borderType=cv2.BORDER_CONSTANT,
+                    value=(0,),
+                )
+            n_imgs.append(img_p)
+
+        data["image"] = n_imgs
+        return data
+
+
+class MakeCentripetalShift:
+    def __init__(self, **kwargs):
+        pass
+
+    def jaccard(self, As, Bs):
+        A = As.shape[0]  # small
+        B = Bs.shape[0]  # large
+
+        dis = np.sqrt(
+            np.sum(
+                (
+                    As[:, np.newaxis, :].repeat(B, axis=1)
+                    - Bs[np.newaxis, :, :].repeat(A, axis=0)
+                )
+                ** 2,
+                axis=-1,
+            )
+        )
+
+        ind = np.argmin(dis, axis=-1)
+
+        return ind
+
+    def __call__(self, data):
+        imgs = data["image"]
+
+        (
+            img,
+            gt_instance,
+            training_mask,
+            gt_kernel_instance,
+            gt_kernel,
+            gt_kernel_inner,
+            training_mask_distance,
+        ) = (imgs[0], imgs[1], imgs[2], imgs[3], imgs[4], imgs[5], imgs[6])
+
+        max_instance = np.max(gt_instance)  # num bbox
+
+        # make centripetal shift
+        gt_distance = np.zeros((2, *img.shape[0:2]), dtype=np.float32)
+        for i in range(1, max_instance + 1):
+            # kernel_reference
+            ind = gt_kernel_inner == i
+
+            if np.sum(ind) == 0:
+                training_mask[gt_instance == i] = 0
+                training_mask_distance[gt_instance == i] = 0
+                continue
+
+            kpoints = (
+                np.array(np.where(ind)).transpose((1, 0))[:, ::-1].astype("float32")
+            )
+
+            ind = (gt_instance == i) * (gt_kernel_instance == 0)
+            if np.sum(ind) == 0:
+                continue
+            pixels = np.where(ind)
+
+            points = np.array(pixels).transpose((1, 0))[:, ::-1].astype("float32")
+
+            bbox_ind = self.jaccard(points, kpoints)
+
+            offset_gt = kpoints[bbox_ind] - points
+
+            gt_distance[:, pixels[0], pixels[1]] = offset_gt.T * 0.1
+
+        img = Image.fromarray(img)
+        img = img.convert("RGB")
+
+        data["image"] = img
+        data["gt_kernel"] = gt_kernel.astype("int64")
+        data["training_mask"] = training_mask.astype("int64")
+        data["gt_instance"] = gt_instance.astype("int64")
+        data["gt_kernel_instance"] = gt_kernel_instance.astype("int64")
+        data["training_mask_distance"] = training_mask_distance.astype("int64")
+        data["gt_distance"] = gt_distance.astype("float32")
+
+        return data
+
+
+class ScaleAlignedShort:
+    def __init__(self, short_size=640, **kwargs):
+        self.short_size = short_size
+
+    def __call__(self, data):
+        img = data["image"]
+
+        org_img_shape = img.shape
+
+        h, w = img.shape[0:2]
+        scale = self.short_size * 1.0 / min(h, w)
+        h = int(h * scale + 0.5)
+        w = int(w * scale + 0.5)
+        if h % 32 != 0:
+            h = h + (32 - h % 32)
+        if w % 32 != 0:
+            w = w + (32 - w % 32)
+        img = cv2.resize(img, dsize=(w, h))
+
+        new_img_shape = img.shape
+        img_shape = np.array(org_img_shape + new_img_shape)
+
+        data["shape"] = img_shape
+        data["image"] = img
+
+        return data

ppocr/data/imaug/drrg_targets.py

@@ -0,0 +1,770 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/drrg_targets.py
+"""
+
+import cv2
+import numpy as np
+from ppocr.utils.utility import check_install
+from numpy.linalg import norm
+
+
+class DRRGTargets(object):
+    def __init__(
+        self,
+        orientation_thr=2.0,
+        resample_step=8.0,
+        num_min_comps=9,
+        num_max_comps=600,
+        min_width=8.0,
+        max_width=24.0,
+        center_region_shrink_ratio=0.3,
+        comp_shrink_ratio=1.0,
+        comp_w_h_ratio=0.3,
+        text_comp_nms_thr=0.25,
+        min_rand_half_height=8.0,
+        max_rand_half_height=24.0,
+        jitter_level=0.2,
+        **kwargs,
+    ):
+        super().__init__()
+        self.orientation_thr = orientation_thr
+        self.resample_step = resample_step
+        self.num_max_comps = num_max_comps
+        self.num_min_comps = num_min_comps
+        self.min_width = min_width
+        self.max_width = max_width
+        self.center_region_shrink_ratio = center_region_shrink_ratio
+        self.comp_shrink_ratio = comp_shrink_ratio
+        self.comp_w_h_ratio = comp_w_h_ratio
+        self.text_comp_nms_thr = text_comp_nms_thr
+        self.min_rand_half_height = min_rand_half_height
+        self.max_rand_half_height = max_rand_half_height
+        self.jitter_level = jitter_level
+        self.eps = 1e-8
+
+    def vector_angle(self, vec1, vec2):
+        if vec1.ndim > 1:
+            unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps).reshape((-1, 1))
+        else:
+            unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps)
+        if vec2.ndim > 1:
+            unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps).reshape((-1, 1))
+        else:
+            unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps)
+        return np.arccos(np.clip(np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0))
+
+    def vector_slope(self, vec):
+        assert len(vec) == 2
+        return abs(vec[1] / (vec[0] + self.eps))
+
+    def vector_sin(self, vec):
+        assert len(vec) == 2
+        return vec[1] / (norm(vec) + self.eps)
+
+    def vector_cos(self, vec):
+        assert len(vec) == 2
+        return vec[0] / (norm(vec) + self.eps)
+
+    def find_head_tail(self, points, orientation_thr):
+        assert points.ndim == 2
+        assert points.shape[0] >= 4
+        assert points.shape[1] == 2
+        assert isinstance(orientation_thr, float)
+
+        if len(points) > 4:
+            pad_points = np.vstack([points, points[0]])
+            edge_vec = pad_points[1:] - pad_points[:-1]
+
+            theta_sum = []
+            adjacent_vec_theta = []
+            for i, edge_vec1 in enumerate(edge_vec):
+                adjacent_ind = [x % len(edge_vec) for x in [i - 1, i + 1]]
+                adjacent_edge_vec = edge_vec[adjacent_ind]
+                temp_theta_sum = np.sum(self.vector_angle(edge_vec1, adjacent_edge_vec))
+                temp_adjacent_theta = self.vector_angle(
+                    adjacent_edge_vec[0], adjacent_edge_vec[1]
+                )
+                theta_sum.append(temp_theta_sum)
+                adjacent_vec_theta.append(temp_adjacent_theta)
+            theta_sum_score = np.array(theta_sum) / np.pi
+            adjacent_theta_score = np.array(adjacent_vec_theta) / np.pi
+            poly_center = np.mean(points, axis=0)
+            edge_dist = np.maximum(
+                norm(pad_points[1:] - poly_center, axis=-1),
+                norm(pad_points[:-1] - poly_center, axis=-1),
+            )
+            dist_score = edge_dist / (np.max(edge_dist) + self.eps)
+            position_score = np.zeros(len(edge_vec))
+            score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score
+            score += 0.35 * dist_score
+            if len(points) % 2 == 0:
+                position_score[(len(score) // 2 - 1)] += 1
+                position_score[-1] += 1
+            score += 0.1 * position_score
+            pad_score = np.concatenate([score, score])
+            score_matrix = np.zeros((len(score), len(score) - 3))
+            x = np.arange(len(score) - 3) / float(len(score) - 4)
+            gaussian = (
+                1.0
+                / (np.sqrt(2.0 * np.pi) * 0.5)
+                * np.exp(-np.power((x - 0.5) / 0.5, 2.0) / 2)
+            )
+            gaussian = gaussian / np.max(gaussian)
+            for i in range(len(score)):
+                score_matrix[i, :] = (
+                    score[i]
+                    + pad_score[(i + 2) : (i + len(score) - 1)] * gaussian * 0.3
+                )
+
+            head_start, tail_increment = np.unravel_index(
+                score_matrix.argmax(), score_matrix.shape
+            )
+            tail_start = (head_start + tail_increment + 2) % len(points)
+            head_end = (head_start + 1) % len(points)
+            tail_end = (tail_start + 1) % len(points)
+
+            if head_end > tail_end:
+                head_start, tail_start = tail_start, head_start
+                head_end, tail_end = tail_end, head_end
+            head_inds = [head_start, head_end]
+            tail_inds = [tail_start, tail_end]
+        else:
+            if self.vector_slope(points[1] - points[0]) + self.vector_slope(
+                points[3] - points[2]
+            ) < self.vector_slope(points[2] - points[1]) + self.vector_slope(
+                points[0] - points[3]
+            ):
+                horizontal_edge_inds = [[0, 1], [2, 3]]
+                vertical_edge_inds = [[3, 0], [1, 2]]
+            else:
+                horizontal_edge_inds = [[3, 0], [1, 2]]
+                vertical_edge_inds = [[0, 1], [2, 3]]
+
+            vertical_len_sum = norm(
+                points[vertical_edge_inds[0][0]] - points[vertical_edge_inds[0][1]]
+            ) + norm(
+                points[vertical_edge_inds[1][0]] - points[vertical_edge_inds[1][1]]
+            )
+            horizontal_len_sum = norm(
+                points[horizontal_edge_inds[0][0]] - points[horizontal_edge_inds[0][1]]
+            ) + norm(
+                points[horizontal_edge_inds[1][0]] - points[horizontal_edge_inds[1][1]]
+            )
+
+            if vertical_len_sum > horizontal_len_sum * orientation_thr:
+                head_inds = horizontal_edge_inds[0]
+                tail_inds = horizontal_edge_inds[1]
+            else:
+                head_inds = vertical_edge_inds[0]
+                tail_inds = vertical_edge_inds[1]
+
+        return head_inds, tail_inds
+
+    def reorder_poly_edge(self, points):
+        assert points.ndim == 2
+        assert points.shape[0] >= 4
+        assert points.shape[1] == 2
+
+        head_inds, tail_inds = self.find_head_tail(points, self.orientation_thr)
+        head_edge, tail_edge = points[head_inds], points[tail_inds]
+
+        pad_points = np.vstack([points, points])
+        if tail_inds[1] < 1:
+            tail_inds[1] = len(points)
+        sideline1 = pad_points[head_inds[1] : tail_inds[1]]
+        sideline2 = pad_points[tail_inds[1] : (head_inds[1] + len(points))]
+        sideline_mean_shift = np.mean(sideline1, axis=0) - np.mean(sideline2, axis=0)
+
+        if sideline_mean_shift[1] > 0:
+            top_sideline, bot_sideline = sideline2, sideline1
+        else:
+            top_sideline, bot_sideline = sideline1, sideline2
+
+        return head_edge, tail_edge, top_sideline, bot_sideline
+
+    def cal_curve_length(self, line):
+        assert line.ndim == 2
+        assert len(line) >= 2
+
+        edges_length = np.sqrt(
+            (line[1:, 0] - line[:-1, 0]) ** 2 + (line[1:, 1] - line[:-1, 1]) ** 2
+        )
+        total_length = np.sum(edges_length)
+        return edges_length, total_length
+
+    def resample_line(self, line, n):
+        assert line.ndim == 2
+        assert line.shape[0] >= 2
+        assert line.shape[1] == 2
+        assert isinstance(n, int)
+        assert n > 2
+
+        edges_length, total_length = self.cal_curve_length(line)
+        t_org = np.insert(np.cumsum(edges_length), 0, 0)
+        unit_t = total_length / (n - 1)
+        t_equidistant = np.arange(1, n - 1, dtype=np.float32) * unit_t
+        edge_ind = 0
+        points = [line[0]]
+        for t in t_equidistant:
+            while edge_ind < len(edges_length) - 1 and t > t_org[edge_ind + 1]:
+                edge_ind += 1
+            t_l, t_r = t_org[edge_ind], t_org[edge_ind + 1]
+            weight = np.array([t_r - t, t - t_l], dtype=np.float32) / (
+                t_r - t_l + self.eps
+            )
+            p_coords = np.dot(weight, line[[edge_ind, edge_ind + 1]])
+            points.append(p_coords)
+        points.append(line[-1])
+        resampled_line = np.vstack(points)
+
+        return resampled_line
+
+    def resample_sidelines(self, sideline1, sideline2, resample_step):
+        assert sideline1.ndim == sideline2.ndim == 2
+        assert sideline1.shape[1] == sideline2.shape[1] == 2
+        assert sideline1.shape[0] >= 2
+        assert sideline2.shape[0] >= 2
+        assert isinstance(resample_step, float)
+
+        _, length1 = self.cal_curve_length(sideline1)
+        _, length2 = self.cal_curve_length(sideline2)
+
+        avg_length = (length1 + length2) / 2
+        resample_point_num = max(int(float(avg_length) / resample_step) + 1, 3)
+
+        resampled_line1 = self.resample_line(sideline1, resample_point_num)
+        resampled_line2 = self.resample_line(sideline2, resample_point_num)
+
+        return resampled_line1, resampled_line2
+
+    def dist_point2line(self, point, line):
+        assert isinstance(line, tuple)
+        point1, point2 = line
+        d = abs(np.cross(point2 - point1, point - point1)) / (
+            norm(point2 - point1) + 1e-8
+        )
+        return d
+
+    def draw_center_region_maps(
+        self,
+        top_line,
+        bot_line,
+        center_line,
+        center_region_mask,
+        top_height_map,
+        bot_height_map,
+        sin_map,
+        cos_map,
+        region_shrink_ratio,
+    ):
+        assert top_line.shape == bot_line.shape == center_line.shape
+        assert (
+            center_region_mask.shape
+            == top_height_map.shape
+            == bot_height_map.shape
+            == sin_map.shape
+            == cos_map.shape
+        )
+        assert isinstance(region_shrink_ratio, float)
+
+        h, w = center_region_mask.shape
+        for i in range(0, len(center_line) - 1):
+            top_mid_point = (top_line[i] + top_line[i + 1]) / 2
+            bot_mid_point = (bot_line[i] + bot_line[i + 1]) / 2
+
+            sin_theta = self.vector_sin(top_mid_point - bot_mid_point)
+            cos_theta = self.vector_cos(top_mid_point - bot_mid_point)
+
+            tl = center_line[i] + (top_line[i] - center_line[i]) * region_shrink_ratio
+            tr = (
+                center_line[i + 1]
+                + (top_line[i + 1] - center_line[i + 1]) * region_shrink_ratio
+            )
+            br = (
+                center_line[i + 1]
+                + (bot_line[i + 1] - center_line[i + 1]) * region_shrink_ratio
+            )
+            bl = center_line[i] + (bot_line[i] - center_line[i]) * region_shrink_ratio
+            current_center_box = np.vstack([tl, tr, br, bl]).astype(np.int32)
+
+            cv2.fillPoly(center_region_mask, [current_center_box], color=1)
+            cv2.fillPoly(sin_map, [current_center_box], color=sin_theta)
+            cv2.fillPoly(cos_map, [current_center_box], color=cos_theta)
+
+            current_center_box[:, 0] = np.clip(current_center_box[:, 0], 0, w - 1)
+            current_center_box[:, 1] = np.clip(current_center_box[:, 1], 0, h - 1)
+            min_coord = np.min(current_center_box, axis=0).astype(np.int32)
+            max_coord = np.max(current_center_box, axis=0).astype(np.int32)
+            current_center_box = current_center_box - min_coord
+            box_sz = max_coord - min_coord + 1
+
+            center_box_mask = np.zeros((box_sz[1], box_sz[0]), dtype=np.uint8)
+            cv2.fillPoly(center_box_mask, [current_center_box], color=1)
+
+            inds = np.argwhere(center_box_mask > 0)
+            inds = inds + (min_coord[1], min_coord[0])
+            inds_xy = np.fliplr(inds)
+            top_height_map[(inds[:, 0], inds[:, 1])] = self.dist_point2line(
+                inds_xy, (top_line[i], top_line[i + 1])
+            )
+            bot_height_map[(inds[:, 0], inds[:, 1])] = self.dist_point2line(
+                inds_xy, (bot_line[i], bot_line[i + 1])
+            )
+
+    def generate_center_mask_attrib_maps(self, img_size, text_polys):
+        assert isinstance(img_size, tuple)
+
+        h, w = img_size
+
+        center_lines = []
+        center_region_mask = np.zeros((h, w), np.uint8)
+        top_height_map = np.zeros((h, w), dtype=np.float32)
+        bot_height_map = np.zeros((h, w), dtype=np.float32)
+        sin_map = np.zeros((h, w), dtype=np.float32)
+        cos_map = np.zeros((h, w), dtype=np.float32)
+
+        for poly in text_polys:
+            polygon_points = poly
+            _, _, top_line, bot_line = self.reorder_poly_edge(polygon_points)
+            resampled_top_line, resampled_bot_line = self.resample_sidelines(
+                top_line, bot_line, self.resample_step
+            )
+            resampled_bot_line = resampled_bot_line[::-1]
+            center_line = (resampled_top_line + resampled_bot_line) / 2
+
+            if self.vector_slope(center_line[-1] - center_line[0]) > 2:
+                if (center_line[-1] - center_line[0])[1] < 0:
+                    center_line = center_line[::-1]
+                    resampled_top_line = resampled_top_line[::-1]
+                    resampled_bot_line = resampled_bot_line[::-1]
+            else:
+                if (center_line[-1] - center_line[0])[0] < 0:
+                    center_line = center_line[::-1]
+                    resampled_top_line = resampled_top_line[::-1]
+                    resampled_bot_line = resampled_bot_line[::-1]
+
+            line_head_shrink_len = (
+                np.clip(
+                    (norm(top_line[0] - bot_line[0]) * self.comp_w_h_ratio),
+                    self.min_width,
+                    self.max_width,
+                )
+                / 2
+            )
+            line_tail_shrink_len = (
+                np.clip(
+                    (norm(top_line[-1] - bot_line[-1]) * self.comp_w_h_ratio),
+                    self.min_width,
+                    self.max_width,
+                )
+                / 2
+            )
+            num_head_shrink = int(line_head_shrink_len // self.resample_step)
+            num_tail_shrink = int(line_tail_shrink_len // self.resample_step)
+            if len(center_line) > num_head_shrink + num_tail_shrink + 2:
+                center_line = center_line[
+                    num_head_shrink : len(center_line) - num_tail_shrink
+                ]
+                resampled_top_line = resampled_top_line[
+                    num_head_shrink : len(resampled_top_line) - num_tail_shrink
+                ]
+                resampled_bot_line = resampled_bot_line[
+                    num_head_shrink : len(resampled_bot_line) - num_tail_shrink
+                ]
+            center_lines.append(center_line.astype(np.int32))
+
+            self.draw_center_region_maps(
+                resampled_top_line,
+                resampled_bot_line,
+                center_line,
+                center_region_mask,
+                top_height_map,
+                bot_height_map,
+                sin_map,
+                cos_map,
+                self.center_region_shrink_ratio,
+            )
+
+        return (
+            center_lines,
+            center_region_mask,
+            top_height_map,
+            bot_height_map,
+            sin_map,
+            cos_map,
+        )
+
+    def generate_rand_comp_attribs(self, num_rand_comps, center_sample_mask):
+        assert isinstance(num_rand_comps, int)
+        assert num_rand_comps > 0
+        assert center_sample_mask.ndim == 2
+
+        h, w = center_sample_mask.shape
+
+        max_rand_half_height = self.max_rand_half_height
+        min_rand_half_height = self.min_rand_half_height
+        max_rand_height = max_rand_half_height * 2
+        max_rand_width = np.clip(
+            max_rand_height * self.comp_w_h_ratio, self.min_width, self.max_width
+        )
+        margin = (
+            int(np.sqrt((max_rand_height / 2) ** 2 + (max_rand_width / 2) ** 2)) + 1
+        )
+
+        if 2 * margin + 1 > min(h, w):
+            assert min(h, w) > (np.sqrt(2) * (self.min_width + 1))
+            max_rand_half_height = max(min(h, w) / 4, self.min_width / 2 + 1)
+            min_rand_half_height = max(max_rand_half_height / 4, self.min_width / 2)
+
+            max_rand_height = max_rand_half_height * 2
+            max_rand_width = np.clip(
+                max_rand_height * self.comp_w_h_ratio, self.min_width, self.max_width
+            )
+            margin = (
+                int(np.sqrt((max_rand_height / 2) ** 2 + (max_rand_width / 2) ** 2)) + 1
+            )
+
+        inner_center_sample_mask = np.zeros_like(center_sample_mask)
+        inner_center_sample_mask[margin : h - margin, margin : w - margin] = (
+            center_sample_mask[margin : h - margin, margin : w - margin]
+        )
+        kernel_size = int(np.clip(max_rand_half_height, 7, 21))
+        inner_center_sample_mask = cv2.erode(
+            inner_center_sample_mask, np.ones((kernel_size, kernel_size), np.uint8)
+        )
+
+        center_candidates = np.argwhere(inner_center_sample_mask > 0)
+        num_center_candidates = len(center_candidates)
+        sample_inds = np.random.choice(num_center_candidates, num_rand_comps)
+        rand_centers = center_candidates[sample_inds]
+
+        rand_top_height = np.random.randint(
+            min_rand_half_height, max_rand_half_height, size=(len(rand_centers), 1)
+        )
+        rand_bot_height = np.random.randint(
+            min_rand_half_height, max_rand_half_height, size=(len(rand_centers), 1)
+        )
+
+        rand_cos = 2 * np.random.random(size=(len(rand_centers), 1)) - 1
+        rand_sin = 2 * np.random.random(size=(len(rand_centers), 1)) - 1
+        scale = np.sqrt(1.0 / (rand_cos**2 + rand_sin**2 + 1e-8))
+        rand_cos = rand_cos * scale
+        rand_sin = rand_sin * scale
+
+        height = rand_top_height + rand_bot_height
+        width = np.clip(height * self.comp_w_h_ratio, self.min_width, self.max_width)
+
+        rand_comp_attribs = np.hstack(
+            [
+                rand_centers[:, ::-1],
+                height,
+                width,
+                rand_cos,
+                rand_sin,
+                np.zeros_like(rand_sin),
+            ]
+        ).astype(np.float32)
+
+        return rand_comp_attribs
+
+    def jitter_comp_attribs(self, comp_attribs, jitter_level):
+        """Jitter text components attributes.
+
+        Args:
+            comp_attribs (ndarray): The text component attributes.
+            jitter_level (float): The jitter level of text components
+                attributes.
+
+        Returns:
+            jittered_comp_attribs (ndarray): The jittered text component
+                attributes (x, y, h, w, cos, sin, comp_label).
+        """
+
+        assert comp_attribs.shape[1] == 7
+        assert comp_attribs.shape[0] > 0
+        assert isinstance(jitter_level, float)
+
+        x = comp_attribs[:, 0].reshape((-1, 1))
+        y = comp_attribs[:, 1].reshape((-1, 1))
+        h = comp_attribs[:, 2].reshape((-1, 1))
+        w = comp_attribs[:, 3].reshape((-1, 1))
+        cos = comp_attribs[:, 4].reshape((-1, 1))
+        sin = comp_attribs[:, 5].reshape((-1, 1))
+        comp_labels = comp_attribs[:, 6].reshape((-1, 1))
+
+        x += (
+            (np.random.random(size=(len(comp_attribs), 1)) - 0.5)
+            * (h * np.abs(cos) + w * np.abs(sin))
+            * jitter_level
+        )
+        y += (
+            (np.random.random(size=(len(comp_attribs), 1)) - 0.5)
+            * (h * np.abs(sin) + w * np.abs(cos))
+            * jitter_level
+        )
+
+        h += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * h * jitter_level
+        w += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * w * jitter_level
+
+        cos += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * 2 * jitter_level
+        sin += (np.random.random(size=(len(comp_attribs), 1)) - 0.5) * 2 * jitter_level
+
+        scale = np.sqrt(1.0 / (cos**2 + sin**2 + 1e-8))
+        cos = cos * scale
+        sin = sin * scale
+
+        jittered_comp_attribs = np.hstack([x, y, h, w, cos, sin, comp_labels])
+
+        return jittered_comp_attribs
+
+    def generate_comp_attribs(
+        self,
+        center_lines,
+        text_mask,
+        center_region_mask,
+        top_height_map,
+        bot_height_map,
+        sin_map,
+        cos_map,
+    ):
+        """Generate text component attributes.
+
+        Args:
+            center_lines (list[ndarray]): The list of text center lines .
+            text_mask (ndarray): The text region mask.
+            center_region_mask (ndarray): The text center region mask.
+            top_height_map (ndarray): The map on which the distance from points
+                to top side lines will be drawn for each pixel in text center
+                regions.
+            bot_height_map (ndarray): The map on which the distance from points
+                to bottom side lines will be drawn for each pixel in text
+                center regions.
+            sin_map (ndarray): The sin(theta) map where theta is the angle
+                between vector (top point - bottom point) and vector (1, 0).
+            cos_map (ndarray): The cos(theta) map where theta is the angle
+                between vector (top point - bottom point) and vector (1, 0).
+
+        Returns:
+            pad_comp_attribs (ndarray): The padded text component attributes
+                of a fixed size.
+        """
+
+        assert isinstance(center_lines, list)
+        assert (
+            text_mask.shape
+            == center_region_mask.shape
+            == top_height_map.shape
+            == bot_height_map.shape
+            == sin_map.shape
+            == cos_map.shape
+        )
+
+        center_lines_mask = np.zeros_like(center_region_mask)
+        cv2.polylines(center_lines_mask, center_lines, 0, 1, 1)
+        center_lines_mask = center_lines_mask * center_region_mask
+        comp_centers = np.argwhere(center_lines_mask > 0)
+
+        y = comp_centers[:, 0]
+        x = comp_centers[:, 1]
+
+        top_height = top_height_map[y, x].reshape((-1, 1)) * self.comp_shrink_ratio
+        bot_height = bot_height_map[y, x].reshape((-1, 1)) * self.comp_shrink_ratio
+        sin = sin_map[y, x].reshape((-1, 1))
+        cos = cos_map[y, x].reshape((-1, 1))
+
+        top_mid_points = comp_centers + np.hstack([top_height * sin, top_height * cos])
+        bot_mid_points = comp_centers - np.hstack([bot_height * sin, bot_height * cos])
+
+        width = (top_height + bot_height) * self.comp_w_h_ratio
+        width = np.clip(width, self.min_width, self.max_width)
+        r = width / 2
+
+        tl = top_mid_points[:, ::-1] - np.hstack([-r * sin, r * cos])
+        tr = top_mid_points[:, ::-1] + np.hstack([-r * sin, r * cos])
+        br = bot_mid_points[:, ::-1] + np.hstack([-r * sin, r * cos])
+        bl = bot_mid_points[:, ::-1] - np.hstack([-r * sin, r * cos])
+        text_comps = np.hstack([tl, tr, br, bl]).astype(np.float32)
+
+        score = np.ones((text_comps.shape[0], 1), dtype=np.float32)
+        text_comps = np.hstack([text_comps, score])
+        check_install("lanms", "lanms-neo")
+        from lanms import merge_quadrangle_n9 as la_nms
+
+        text_comps = la_nms(text_comps, self.text_comp_nms_thr)
+
+        if text_comps.shape[0] >= 1:
+            img_h, img_w = center_region_mask.shape
+            text_comps[:, 0:8:2] = np.clip(text_comps[:, 0:8:2], 0, img_w - 1)
+            text_comps[:, 1:8:2] = np.clip(text_comps[:, 1:8:2], 0, img_h - 1)
+
+            comp_centers = np.mean(
+                text_comps[:, 0:8].reshape((-1, 4, 2)), axis=1
+            ).astype(np.int32)
+            x = comp_centers[:, 0]
+            y = comp_centers[:, 1]
+
+            height = (top_height_map[y, x] + bot_height_map[y, x]).reshape((-1, 1))
+            width = np.clip(
+                height * self.comp_w_h_ratio, self.min_width, self.max_width
+            )
+
+            cos = cos_map[y, x].reshape((-1, 1))
+            sin = sin_map[y, x].reshape((-1, 1))
+
+            _, comp_label_mask = cv2.connectedComponents(
+                center_region_mask, connectivity=8
+            )
+            comp_labels = comp_label_mask[y, x].reshape((-1, 1)).astype(np.float32)
+
+            x = x.reshape((-1, 1)).astype(np.float32)
+            y = y.reshape((-1, 1)).astype(np.float32)
+            comp_attribs = np.hstack([x, y, height, width, cos, sin, comp_labels])
+            comp_attribs = self.jitter_comp_attribs(comp_attribs, self.jitter_level)
+
+            if comp_attribs.shape[0] < self.num_min_comps:
+                num_rand_comps = self.num_min_comps - comp_attribs.shape[0]
+                rand_comp_attribs = self.generate_rand_comp_attribs(
+                    num_rand_comps, 1 - text_mask
+                )
+                comp_attribs = np.vstack([comp_attribs, rand_comp_attribs])
+        else:
+            comp_attribs = self.generate_rand_comp_attribs(
+                self.num_min_comps, 1 - text_mask
+            )
+
+        num_comps = (
+            np.ones((comp_attribs.shape[0], 1), dtype=np.float32)
+            * comp_attribs.shape[0]
+        )
+        comp_attribs = np.hstack([num_comps, comp_attribs])
+
+        if comp_attribs.shape[0] > self.num_max_comps:
+            comp_attribs = comp_attribs[: self.num_max_comps, :]
+            comp_attribs[:, 0] = self.num_max_comps
+
+        pad_comp_attribs = np.zeros(
+            (self.num_max_comps, comp_attribs.shape[1]), dtype=np.float32
+        )
+        pad_comp_attribs[: comp_attribs.shape[0], :] = comp_attribs
+
+        return pad_comp_attribs
+
+    def generate_text_region_mask(self, img_size, text_polys):
+        """Generate text center region mask and geometry attribute maps.
+
+        Args:
+            img_size (tuple): The image size (height, width).
+            text_polys (list[list[ndarray]]): The list of text polygons.
+
+        Returns:
+            text_region_mask (ndarray): The text region mask.
+        """
+
+        assert isinstance(img_size, tuple)
+
+        h, w = img_size
+        text_region_mask = np.zeros((h, w), dtype=np.uint8)
+
+        for poly in text_polys:
+            polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
+            cv2.fillPoly(text_region_mask, polygon, 1)
+
+        return text_region_mask
+
+    def generate_effective_mask(self, mask_size: tuple, polygons_ignore):
+        """Generate effective mask by setting the ineffective regions to 0 and
+        effective regions to 1.
+
+        Args:
+            mask_size (tuple): The mask size.
+            polygons_ignore (list[[ndarray]]: The list of ignored text
+                polygons.
+
+        Returns:
+            mask (ndarray): The effective mask of (height, width).
+        """
+        mask = np.ones(mask_size, dtype=np.uint8)
+
+        for poly in polygons_ignore:
+            instance = poly.astype(np.int32).reshape(1, -1, 2)
+            cv2.fillPoly(mask, instance, 0)
+
+        return mask
+
+    def generate_targets(self, data):
+        """Generate the gt targets for DRRG.
+
+        Args:
+            data (dict): The input result dictionary.
+
+        Returns:
+            data (dict): The output result dictionary.
+        """
+
+        assert isinstance(data, dict)
+
+        image = data["image"]
+        polygons = data["polys"]
+        ignore_tags = data["ignore_tags"]
+        h, w, _ = image.shape
+
+        polygon_masks = []
+        polygon_masks_ignore = []
+        for tag, polygon in zip(ignore_tags, polygons):
+            if tag is True:
+                polygon_masks_ignore.append(polygon)
+            else:
+                polygon_masks.append(polygon)
+
+        gt_text_mask = self.generate_text_region_mask((h, w), polygon_masks)
+        gt_mask = self.generate_effective_mask((h, w), polygon_masks_ignore)
+        (
+            center_lines,
+            gt_center_region_mask,
+            gt_top_height_map,
+            gt_bot_height_map,
+            gt_sin_map,
+            gt_cos_map,
+        ) = self.generate_center_mask_attrib_maps((h, w), polygon_masks)
+
+        gt_comp_attribs = self.generate_comp_attribs(
+            center_lines,
+            gt_text_mask,
+            gt_center_region_mask,
+            gt_top_height_map,
+            gt_bot_height_map,
+            gt_sin_map,
+            gt_cos_map,
+        )
+
+        mapping = {
+            "gt_text_mask": gt_text_mask,
+            "gt_center_region_mask": gt_center_region_mask,
+            "gt_mask": gt_mask,
+            "gt_top_height_map": gt_top_height_map,
+            "gt_bot_height_map": gt_bot_height_map,
+            "gt_sin_map": gt_sin_map,
+            "gt_cos_map": gt_cos_map,
+        }
+
+        data.update(mapping)
+        data["gt_comp_attribs"] = gt_comp_attribs
+        return data
+
+    def __call__(self, data):
+        data = self.generate_targets(data)
+        return data

ppocr/data/imaug/east_process.py

@@ -0,0 +1,446 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is referred from:
+https://github.com/songdejia/EAST/blob/master/data_utils.py
+"""
+import math
+import cv2
+import numpy as np
+import json
+import sys
+import os
+
+__all__ = ["EASTProcessTrain"]
+
+
+class EASTProcessTrain(object):
+    def __init__(
+        self,
+        image_shape=[512, 512],
+        background_ratio=0.125,
+        min_crop_side_ratio=0.1,
+        min_text_size=10,
+        **kwargs,
+    ):
+        self.input_size = image_shape[1]
+        self.random_scale = np.array([0.5, 1, 2.0, 3.0])
+        self.background_ratio = background_ratio
+        self.min_crop_side_ratio = min_crop_side_ratio
+        self.min_text_size = min_text_size
+
+    def preprocess(self, im):
+        input_size = self.input_size
+        im_shape = im.shape
+        im_size_min = np.min(im_shape[0:2])
+        im_size_max = np.max(im_shape[0:2])
+        im_scale = float(input_size) / float(im_size_max)
+        im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale)
+        img_mean = [0.485, 0.456, 0.406]
+        img_std = [0.229, 0.224, 0.225]
+        # im = im[:, :, ::-1].astype(np.float32)
+        im = im / 255
+        im -= img_mean
+        im /= img_std
+        new_h, new_w, _ = im.shape
+        im_padded = np.zeros((input_size, input_size, 3), dtype=np.float32)
+        im_padded[:new_h, :new_w, :] = im
+        im_padded = im_padded.transpose((2, 0, 1))
+        im_padded = im_padded[np.newaxis, :]
+        return im_padded, im_scale
+
+    def rotate_im_poly(self, im, text_polys):
+        """
+        rotate image with 90 / 180 / 270 degre
+        """
+        im_w, im_h = im.shape[1], im.shape[0]
+        dst_im = im.copy()
+        dst_polys = []
+        rand_degree_ratio = np.random.rand()
+        rand_degree_cnt = 1
+        if 0.333 < rand_degree_ratio < 0.666:
+            rand_degree_cnt = 2
+        elif rand_degree_ratio > 0.666:
+            rand_degree_cnt = 3
+        for i in range(rand_degree_cnt):
+            dst_im = np.rot90(dst_im)
+        rot_degree = -90 * rand_degree_cnt
+        rot_angle = rot_degree * math.pi / 180.0
+        n_poly = text_polys.shape[0]
+        cx, cy = 0.5 * im_w, 0.5 * im_h
+        ncx, ncy = 0.5 * dst_im.shape[1], 0.5 * dst_im.shape[0]
+        for i in range(n_poly):
+            wordBB = text_polys[i]
+            poly = []
+            for j in range(4):
+                sx, sy = wordBB[j][0], wordBB[j][1]
+                dx = (
+                    math.cos(rot_angle) * (sx - cx)
+                    - math.sin(rot_angle) * (sy - cy)
+                    + ncx
+                )
+                dy = (
+                    math.sin(rot_angle) * (sx - cx)
+                    + math.cos(rot_angle) * (sy - cy)
+                    + ncy
+                )
+                poly.append([dx, dy])
+            dst_polys.append(poly)
+        dst_polys = np.array(dst_polys, dtype=np.float32)
+        return dst_im, dst_polys
+
+    def polygon_area(self, poly):
+        """
+        compute area of a polygon
+        :param poly:
+        :return:
+        """
+        edge = [
+            (poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
+            (poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
+            (poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
+            (poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1]),
+        ]
+        return np.sum(edge) / 2.0
+
+    def check_and_validate_polys(self, polys, tags, img_height, img_width):
+        """
+        check so that the text poly is in the same direction,
+        and also filter some invalid polygons
+        :param polys:
+        :param tags:
+        :return:
+        """
+        h, w = img_height, img_width
+        if polys.shape[0] == 0:
+            return polys
+        polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
+        polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
+
+        validated_polys = []
+        validated_tags = []
+        for poly, tag in zip(polys, tags):
+            p_area = self.polygon_area(poly)
+            # invalid poly
+            if abs(p_area) < 1:
+                continue
+            if p_area > 0:
+                #'poly in wrong direction'
+                if not tag:
+                    tag = True  # reversed cases should be ignore
+                poly = poly[(0, 3, 2, 1), :]
+            validated_polys.append(poly)
+            validated_tags.append(tag)
+        return np.array(validated_polys), np.array(validated_tags)
+
+    def draw_img_polys(self, img, polys):
+        if len(img.shape) == 4:
+            img = np.squeeze(img, axis=0)
+        if img.shape[0] == 3:
+            img = img.transpose((1, 2, 0))
+            img[:, :, 2] += 123.68
+            img[:, :, 1] += 116.78
+            img[:, :, 0] += 103.94
+        cv2.imwrite("tmp.jpg", img)
+        img = cv2.imread("tmp.jpg")
+        for box in polys:
+            box = box.astype(np.int32).reshape((-1, 1, 2))
+            cv2.polylines(img, [box], True, color=(255, 255, 0), thickness=2)
+        import random
+
+        ino = random.randint(0, 100)
+        cv2.imwrite("tmp_%d.jpg" % ino, img)
+        return
+
+    def shrink_poly(self, poly, r):
+        """
+        fit a poly inside the origin poly, maybe bugs here...
+        used for generate the score map
+        :param poly: the text poly
+        :param r: r in the paper
+        :return: the shrunk poly
+        """
+        # shrink ratio
+        R = 0.3
+        # find the longer pair
+        dist0 = np.linalg.norm(poly[0] - poly[1])
+        dist1 = np.linalg.norm(poly[2] - poly[3])
+        dist2 = np.linalg.norm(poly[0] - poly[3])
+        dist3 = np.linalg.norm(poly[1] - poly[2])
+        if dist0 + dist1 > dist2 + dist3:
+            # first move (p0, p1), (p2, p3), then (p0, p3), (p1, p2)
+            ## p0, p1
+            theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
+            poly[0][0] += R * r[0] * np.cos(theta)
+            poly[0][1] += R * r[0] * np.sin(theta)
+            poly[1][0] -= R * r[1] * np.cos(theta)
+            poly[1][1] -= R * r[1] * np.sin(theta)
+            ## p2, p3
+            theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
+            poly[3][0] += R * r[3] * np.cos(theta)
+            poly[3][1] += R * r[3] * np.sin(theta)
+            poly[2][0] -= R * r[2] * np.cos(theta)
+            poly[2][1] -= R * r[2] * np.sin(theta)
+            ## p0, p3
+            theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
+            poly[0][0] += R * r[0] * np.sin(theta)
+            poly[0][1] += R * r[0] * np.cos(theta)
+            poly[3][0] -= R * r[3] * np.sin(theta)
+            poly[3][1] -= R * r[3] * np.cos(theta)
+            ## p1, p2
+            theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
+            poly[1][0] += R * r[1] * np.sin(theta)
+            poly[1][1] += R * r[1] * np.cos(theta)
+            poly[2][0] -= R * r[2] * np.sin(theta)
+            poly[2][1] -= R * r[2] * np.cos(theta)
+        else:
+            ## p0, p3
+            # print poly
+            theta = np.arctan2((poly[3][0] - poly[0][0]), (poly[3][1] - poly[0][1]))
+            poly[0][0] += R * r[0] * np.sin(theta)
+            poly[0][1] += R * r[0] * np.cos(theta)
+            poly[3][0] -= R * r[3] * np.sin(theta)
+            poly[3][1] -= R * r[3] * np.cos(theta)
+            ## p1, p2
+            theta = np.arctan2((poly[2][0] - poly[1][0]), (poly[2][1] - poly[1][1]))
+            poly[1][0] += R * r[1] * np.sin(theta)
+            poly[1][1] += R * r[1] * np.cos(theta)
+            poly[2][0] -= R * r[2] * np.sin(theta)
+            poly[2][1] -= R * r[2] * np.cos(theta)
+            ## p0, p1
+            theta = np.arctan2((poly[1][1] - poly[0][1]), (poly[1][0] - poly[0][0]))
+            poly[0][0] += R * r[0] * np.cos(theta)
+            poly[0][1] += R * r[0] * np.sin(theta)
+            poly[1][0] -= R * r[1] * np.cos(theta)
+            poly[1][1] -= R * r[1] * np.sin(theta)
+            ## p2, p3
+            theta = np.arctan2((poly[2][1] - poly[3][1]), (poly[2][0] - poly[3][0]))
+            poly[3][0] += R * r[3] * np.cos(theta)
+            poly[3][1] += R * r[3] * np.sin(theta)
+            poly[2][0] -= R * r[2] * np.cos(theta)
+            poly[2][1] -= R * r[2] * np.sin(theta)
+        return poly
+
+    def generate_quad(self, im_size, polys, tags):
+        """
+        Generate quadrangle.
+        """
+        h, w = im_size
+        poly_mask = np.zeros((h, w), dtype=np.uint8)
+        score_map = np.zeros((h, w), dtype=np.uint8)
+        # (x1, y1, ..., x4, y4, short_edge_norm)
+        geo_map = np.zeros((h, w, 9), dtype=np.float32)
+        # mask used during training, to ignore some hard areas
+        training_mask = np.ones((h, w), dtype=np.uint8)
+        for poly_idx, poly_tag in enumerate(zip(polys, tags)):
+            poly = poly_tag[0]
+            tag = poly_tag[1]
+
+            r = [None, None, None, None]
+            for i in range(4):
+                dist1 = np.linalg.norm(poly[i] - poly[(i + 1) % 4])
+                dist2 = np.linalg.norm(poly[i] - poly[(i - 1) % 4])
+                r[i] = min(dist1, dist2)
+            # score map
+            shrinked_poly = self.shrink_poly(poly.copy(), r).astype(np.int32)[
+                np.newaxis, :, :
+            ]
+            cv2.fillPoly(score_map, shrinked_poly, 1)
+            cv2.fillPoly(poly_mask, shrinked_poly, poly_idx + 1)
+            # if the poly is too small, then ignore it during training
+            poly_h = min(
+                np.linalg.norm(poly[0] - poly[3]), np.linalg.norm(poly[1] - poly[2])
+            )
+            poly_w = min(
+                np.linalg.norm(poly[0] - poly[1]), np.linalg.norm(poly[2] - poly[3])
+            )
+            if min(poly_h, poly_w) < self.min_text_size:
+                cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
+
+            if tag:
+                cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
+
+            xy_in_poly = np.argwhere(poly_mask == (poly_idx + 1))
+            # geo map.
+            y_in_poly = xy_in_poly[:, 0]
+            x_in_poly = xy_in_poly[:, 1]
+            poly[:, 0] = np.minimum(np.maximum(poly[:, 0], 0), w)
+            poly[:, 1] = np.minimum(np.maximum(poly[:, 1], 0), h)
+            for pno in range(4):
+                geo_channel_beg = pno * 2
+                geo_map[y_in_poly, x_in_poly, geo_channel_beg] = (
+                    x_in_poly - poly[pno, 0]
+                )
+                geo_map[y_in_poly, x_in_poly, geo_channel_beg + 1] = (
+                    y_in_poly - poly[pno, 1]
+                )
+            geo_map[y_in_poly, x_in_poly, 8] = 1.0 / max(min(poly_h, poly_w), 1.0)
+        return score_map, geo_map, training_mask
+
+    def crop_area(self, im, polys, tags, crop_background=False, max_tries=50):
+        """
+        make random crop from the input image
+        :param im:
+        :param polys:
+        :param tags:
+        :param crop_background:
+        :param max_tries:
+        :return:
+        """
+        h, w, _ = im.shape
+        pad_h = h // 10
+        pad_w = w // 10
+        h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
+        w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
+        for poly in polys:
+            poly = np.round(poly, decimals=0).astype(np.int32)
+            minx = np.min(poly[:, 0])
+            maxx = np.max(poly[:, 0])
+            w_array[minx + pad_w : maxx + pad_w] = 1
+            miny = np.min(poly[:, 1])
+            maxy = np.max(poly[:, 1])
+            h_array[miny + pad_h : maxy + pad_h] = 1
+        # ensure the cropped area not across a text
+        h_axis = np.where(h_array == 0)[0]
+        w_axis = np.where(w_array == 0)[0]
+        if len(h_axis) == 0 or len(w_axis) == 0:
+            return im, polys, tags
+
+        for i in range(max_tries):
+            xx = np.random.choice(w_axis, size=2)
+            xmin = np.min(xx) - pad_w
+            xmax = np.max(xx) - pad_w
+            xmin = np.clip(xmin, 0, w - 1)
+            xmax = np.clip(xmax, 0, w - 1)
+            yy = np.random.choice(h_axis, size=2)
+            ymin = np.min(yy) - pad_h
+            ymax = np.max(yy) - pad_h
+            ymin = np.clip(ymin, 0, h - 1)
+            ymax = np.clip(ymax, 0, h - 1)
+            if (
+                xmax - xmin < self.min_crop_side_ratio * w
+                or ymax - ymin < self.min_crop_side_ratio * h
+            ):
+                # area too small
+                continue
+            if polys.shape[0] != 0:
+                poly_axis_in_area = (
+                    (polys[:, :, 0] >= xmin)
+                    & (polys[:, :, 0] <= xmax)
+                    & (polys[:, :, 1] >= ymin)
+                    & (polys[:, :, 1] <= ymax)
+                )
+                selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
+            else:
+                selected_polys = []
+
+            if len(selected_polys) == 0:
+                # no text in this area
+                if crop_background:
+                    im = im[ymin : ymax + 1, xmin : xmax + 1, :]
+                    polys = []
+                    tags = []
+                    return im, polys, tags
+                else:
+                    continue
+
+            im = im[ymin : ymax + 1, xmin : xmax + 1, :]
+            polys = polys[selected_polys]
+            tags = tags[selected_polys]
+            polys[:, :, 0] -= xmin
+            polys[:, :, 1] -= ymin
+            return im, polys, tags
+        return im, polys, tags
+
+    def crop_background_infor(self, im, text_polys, text_tags):
+        im, text_polys, text_tags = self.crop_area(
+            im, text_polys, text_tags, crop_background=True
+        )
+
+        if len(text_polys) > 0:
+            return None
+        # pad and resize image
+        input_size = self.input_size
+        im, ratio = self.preprocess(im)
+        score_map = np.zeros((input_size, input_size), dtype=np.float32)
+        geo_map = np.zeros((input_size, input_size, 9), dtype=np.float32)
+        training_mask = np.ones((input_size, input_size), dtype=np.float32)
+        return im, score_map, geo_map, training_mask
+
+    def crop_foreground_infor(self, im, text_polys, text_tags):
+        im, text_polys, text_tags = self.crop_area(
+            im, text_polys, text_tags, crop_background=False
+        )
+
+        if text_polys.shape[0] == 0:
+            return None
+        # continue for all ignore case
+        if np.sum((text_tags * 1.0)) >= text_tags.size:
+            return None
+        # pad and resize image
+        input_size = self.input_size
+        im, ratio = self.preprocess(im)
+        text_polys[:, :, 0] *= ratio
+        text_polys[:, :, 1] *= ratio
+        _, _, new_h, new_w = im.shape
+        #         print(im.shape)
+        #         self.draw_img_polys(im, text_polys)
+        score_map, geo_map, training_mask = self.generate_quad(
+            (new_h, new_w), text_polys, text_tags
+        )
+        return im, score_map, geo_map, training_mask
+
+    def __call__(self, data):
+        im = data["image"]
+        text_polys = data["polys"]
+        text_tags = data["ignore_tags"]
+        if im is None:
+            return None
+        if text_polys.shape[0] == 0:
+            return None
+
+        # add rotate cases
+        if np.random.rand() < 0.5:
+            im, text_polys = self.rotate_im_poly(im, text_polys)
+        h, w, _ = im.shape
+        text_polys, text_tags = self.check_and_validate_polys(
+            text_polys, text_tags, h, w
+        )
+        if text_polys.shape[0] == 0:
+            return None
+
+        # random scale this image
+        rd_scale = np.random.choice(self.random_scale)
+        im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
+        text_polys *= rd_scale
+        if np.random.rand() < self.background_ratio:
+            outs = self.crop_background_infor(im, text_polys, text_tags)
+        else:
+            outs = self.crop_foreground_infor(im, text_polys, text_tags)
+
+        if outs is None:
+            return None
+        im, score_map, geo_map, training_mask = outs
+        score_map = score_map[np.newaxis, ::4, ::4].astype(np.float32)
+        geo_map = np.swapaxes(geo_map, 1, 2)
+        geo_map = np.swapaxes(geo_map, 1, 0)
+        geo_map = geo_map[:, ::4, ::4].astype(np.float32)
+        training_mask = training_mask[np.newaxis, ::4, ::4]
+        training_mask = training_mask.astype(np.float32)
+
+        data["image"] = im[0]
+        data["score_map"] = score_map
+        data["geo_map"] = geo_map
+        data["training_mask"] = training_mask
+        return data

ppocr/data/imaug/fce_aug.py

@@ -0,0 +1,575 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/transforms.py
+"""
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+from shapely.geometry import Polygon
+import math
+from ppocr.utils.poly_nms import poly_intersection
+
+
+class RandomScaling:
+    def __init__(self, size=800, scale=(3.0 / 4, 5.0 / 2), **kwargs):
+        """Random scale the image while keeping aspect.
+
+        Args:
+            size (int) : Base size before scaling.
+            scale (tuple(float)) : The range of scaling.
+        """
+        assert isinstance(size, int)
+        assert isinstance(scale, float) or isinstance(scale, tuple)
+        self.size = size
+        self.scale = scale if isinstance(scale, tuple) else (1 - scale, 1 + scale)
+
+    def __call__(self, data):
+        image = data["image"]
+        text_polys = data["polys"]
+        h, w, _ = image.shape
+
+        aspect_ratio = np.random.uniform(min(self.scale), max(self.scale))
+        scales = self.size * 1.0 / max(h, w) * aspect_ratio
+        scales = np.array([scales, scales])
+        out_size = (int(h * scales[1]), int(w * scales[0]))
+        image = cv2.resize(image, out_size[::-1])
+
+        data["image"] = image
+        text_polys[:, :, 0::2] = text_polys[:, :, 0::2] * scales[1]
+        text_polys[:, :, 1::2] = text_polys[:, :, 1::2] * scales[0]
+        data["polys"] = text_polys
+
+        return data
+
+
+class RandomCropFlip:
+    def __init__(
+        self, pad_ratio=0.1, crop_ratio=0.5, iter_num=1, min_area_ratio=0.2, **kwargs
+    ):
+        """Random crop and flip a patch of the image.
+
+        Args:
+            crop_ratio (float): The ratio of cropping.
+            iter_num (int): Number of operations.
+            min_area_ratio (float): Minimal area ratio between cropped patch
+                and original image.
+        """
+        assert isinstance(crop_ratio, float)
+        assert isinstance(iter_num, int)
+        assert isinstance(min_area_ratio, float)
+
+        self.pad_ratio = pad_ratio
+        self.epsilon = 1e-2
+        self.crop_ratio = crop_ratio
+        self.iter_num = iter_num
+        self.min_area_ratio = min_area_ratio
+
+    def __call__(self, results):
+        for i in range(self.iter_num):
+            results = self.random_crop_flip(results)
+
+        return results
+
+    def random_crop_flip(self, results):
+        image = results["image"]
+        polygons = results["polys"]
+        ignore_tags = results["ignore_tags"]
+        if len(polygons) == 0:
+            return results
+
+        if np.random.random() >= self.crop_ratio:
+            return results
+
+        h, w, _ = image.shape
+        area = h * w
+        pad_h = int(h * self.pad_ratio)
+        pad_w = int(w * self.pad_ratio)
+        h_axis, w_axis = self.generate_crop_target(image, polygons, pad_h, pad_w)
+        if len(h_axis) == 0 or len(w_axis) == 0:
+            return results
+
+        attempt = 0
+        while attempt < 50:
+            attempt += 1
+            polys_keep = []
+            polys_new = []
+            ignore_tags_keep = []
+            ignore_tags_new = []
+            xx = np.random.choice(w_axis, size=2)
+            xmin = np.min(xx) - pad_w
+            xmax = np.max(xx) - pad_w
+            xmin = np.clip(xmin, 0, w - 1)
+            xmax = np.clip(xmax, 0, w - 1)
+            yy = np.random.choice(h_axis, size=2)
+            ymin = np.min(yy) - pad_h
+            ymax = np.max(yy) - pad_h
+            ymin = np.clip(ymin, 0, h - 1)
+            ymax = np.clip(ymax, 0, h - 1)
+            if (xmax - xmin) * (ymax - ymin) < area * self.min_area_ratio:
+                # area too small
+                continue
+
+            pts = np.stack(
+                [[xmin, xmax, xmax, xmin], [ymin, ymin, ymax, ymax]]
+            ).T.astype(np.int32)
+            pp = Polygon(pts)
+            fail_flag = False
+            for polygon, ignore_tag in zip(polygons, ignore_tags):
+                ppi = Polygon(polygon.reshape(-1, 2))
+                ppiou, _ = poly_intersection(ppi, pp, buffer=0)
+                if (
+                    np.abs(ppiou - float(ppi.area)) > self.epsilon
+                    and np.abs(ppiou) > self.epsilon
+                ):
+                    fail_flag = True
+                    break
+                elif np.abs(ppiou - float(ppi.area)) < self.epsilon:
+                    polys_new.append(polygon)
+                    ignore_tags_new.append(ignore_tag)
+                else:
+                    polys_keep.append(polygon)
+                    ignore_tags_keep.append(ignore_tag)
+
+            if fail_flag:
+                continue
+            else:
+                break
+
+        cropped = image[ymin:ymax, xmin:xmax, :]
+        select_type = np.random.randint(3)
+        if select_type == 0:
+            img = np.ascontiguousarray(cropped[:, ::-1])
+        elif select_type == 1:
+            img = np.ascontiguousarray(cropped[::-1, :])
+        else:
+            img = np.ascontiguousarray(cropped[::-1, ::-1])
+        image[ymin:ymax, xmin:xmax, :] = img
+        results["img"] = image
+
+        if len(polys_new) != 0:
+            height, width, _ = cropped.shape
+            if select_type == 0:
+                for idx, polygon in enumerate(polys_new):
+                    poly = polygon.reshape(-1, 2)
+                    poly[:, 0] = width - poly[:, 0] + 2 * xmin
+                    polys_new[idx] = poly
+            elif select_type == 1:
+                for idx, polygon in enumerate(polys_new):
+                    poly = polygon.reshape(-1, 2)
+                    poly[:, 1] = height - poly[:, 1] + 2 * ymin
+                    polys_new[idx] = poly
+            else:
+                for idx, polygon in enumerate(polys_new):
+                    poly = polygon.reshape(-1, 2)
+                    poly[:, 0] = width - poly[:, 0] + 2 * xmin
+                    poly[:, 1] = height - poly[:, 1] + 2 * ymin
+                    polys_new[idx] = poly
+            polygons = polys_keep + polys_new
+            ignore_tags = ignore_tags_keep + ignore_tags_new
+            results["polys"] = np.array(polygons)
+            results["ignore_tags"] = ignore_tags
+
+        return results
+
+    def generate_crop_target(self, image, all_polys, pad_h, pad_w):
+        """Generate crop target and make sure not to crop the polygon
+        instances.
+
+        Args:
+            image (ndarray): The image waited to be crop.
+            all_polys (list[list[ndarray]]): All polygons including ground
+                truth polygons and ground truth ignored polygons.
+            pad_h (int): Padding length of height.
+            pad_w (int): Padding length of width.
+        Returns:
+            h_axis (ndarray): Vertical cropping range.
+            w_axis (ndarray): Horizontal cropping range.
+        """
+        h, w, _ = image.shape
+        h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
+        w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
+
+        text_polys = []
+        for polygon in all_polys:
+            rect = cv2.minAreaRect(polygon.astype(np.int32).reshape(-1, 2))
+            box = cv2.boxPoints(rect)
+            box = np.int64(box)
+            text_polys.append([box[0], box[1], box[2], box[3]])
+
+        polys = np.array(text_polys, dtype=np.int32)
+        for poly in polys:
+            poly = np.round(poly, decimals=0).astype(np.int32)
+            minx = np.min(poly[:, 0])
+            maxx = np.max(poly[:, 0])
+            w_array[minx + pad_w : maxx + pad_w] = 1
+            miny = np.min(poly[:, 1])
+            maxy = np.max(poly[:, 1])
+            h_array[miny + pad_h : maxy + pad_h] = 1
+
+        h_axis = np.where(h_array == 0)[0]
+        w_axis = np.where(w_array == 0)[0]
+        return h_axis, w_axis
+
+
+class RandomCropPolyInstances:
+    """Randomly crop images and make sure to contain at least one intact
+    instance."""
+
+    def __init__(self, crop_ratio=5.0 / 8.0, min_side_ratio=0.4, **kwargs):
+        super().__init__()
+        self.crop_ratio = crop_ratio
+        self.min_side_ratio = min_side_ratio
+
+    def sample_valid_start_end(self, valid_array, min_len, max_start, min_end):
+        assert isinstance(min_len, int)
+        assert len(valid_array) > min_len
+
+        start_array = valid_array.copy()
+        max_start = min(len(start_array) - min_len, max_start)
+        start_array[max_start:] = 0
+        start_array[0] = 1
+        diff_array = np.hstack([0, start_array]) - np.hstack([start_array, 0])
+        region_starts = np.where(diff_array < 0)[0]
+        region_ends = np.where(diff_array > 0)[0]
+        region_ind = np.random.randint(0, len(region_starts))
+        start = np.random.randint(region_starts[region_ind], region_ends[region_ind])
+
+        end_array = valid_array.copy()
+        min_end = max(start + min_len, min_end)
+        end_array[:min_end] = 0
+        end_array[-1] = 1
+        diff_array = np.hstack([0, end_array]) - np.hstack([end_array, 0])
+        region_starts = np.where(diff_array < 0)[0]
+        region_ends = np.where(diff_array > 0)[0]
+        region_ind = np.random.randint(0, len(region_starts))
+        end = np.random.randint(region_starts[region_ind], region_ends[region_ind])
+        return start, end
+
+    def sample_crop_box(self, img_size, results):
+        """Generate crop box and make sure not to crop the polygon instances.
+
+        Args:
+            img_size (tuple(int)): The image size (h, w).
+            results (dict): The results dict.
+        """
+
+        assert isinstance(img_size, tuple)
+        h, w = img_size[:2]
+
+        key_masks = results["polys"]
+
+        x_valid_array = np.ones(w, dtype=np.int32)
+        y_valid_array = np.ones(h, dtype=np.int32)
+
+        selected_mask = key_masks[np.random.randint(0, len(key_masks))]
+        selected_mask = selected_mask.reshape((-1, 2)).astype(np.int32)
+        max_x_start = max(np.min(selected_mask[:, 0]) - 2, 0)
+        min_x_end = min(np.max(selected_mask[:, 0]) + 3, w - 1)
+        max_y_start = max(np.min(selected_mask[:, 1]) - 2, 0)
+        min_y_end = min(np.max(selected_mask[:, 1]) + 3, h - 1)
+
+        for mask in key_masks:
+            mask = mask.reshape((-1, 2)).astype(np.int32)
+            clip_x = np.clip(mask[:, 0], 0, w - 1)
+            clip_y = np.clip(mask[:, 1], 0, h - 1)
+            min_x, max_x = np.min(clip_x), np.max(clip_x)
+            min_y, max_y = np.min(clip_y), np.max(clip_y)
+
+            x_valid_array[min_x - 2 : max_x + 3] = 0
+            y_valid_array[min_y - 2 : max_y + 3] = 0
+
+        min_w = int(w * self.min_side_ratio)
+        min_h = int(h * self.min_side_ratio)
+
+        x1, x2 = self.sample_valid_start_end(
+            x_valid_array, min_w, max_x_start, min_x_end
+        )
+        y1, y2 = self.sample_valid_start_end(
+            y_valid_array, min_h, max_y_start, min_y_end
+        )
+
+        return np.array([x1, y1, x2, y2])
+
+    def crop_img(self, img, bbox):
+        assert img.ndim == 3
+        h, w, _ = img.shape
+        assert 0 <= bbox[1] < bbox[3] <= h
+        assert 0 <= bbox[0] < bbox[2] <= w
+        return img[bbox[1] : bbox[3], bbox[0] : bbox[2]]
+
+    def __call__(self, results):
+        image = results["image"]
+        polygons = results["polys"]
+        ignore_tags = results["ignore_tags"]
+        if len(polygons) < 1:
+            return results
+
+        if np.random.random_sample() < self.crop_ratio:
+            crop_box = self.sample_crop_box(image.shape, results)
+            img = self.crop_img(image, crop_box)
+            results["image"] = img
+            # crop and filter masks
+            x1, y1, x2, y2 = crop_box
+            w = max(x2 - x1, 1)
+            h = max(y2 - y1, 1)
+            polygons[:, :, 0::2] = polygons[:, :, 0::2] - x1
+            polygons[:, :, 1::2] = polygons[:, :, 1::2] - y1
+
+            valid_masks_list = []
+            valid_tags_list = []
+            for ind, polygon in enumerate(polygons):
+                if (
+                    (polygon[:, ::2] > -4).all()
+                    and (polygon[:, ::2] < w + 4).all()
+                    and (polygon[:, 1::2] > -4).all()
+                    and (polygon[:, 1::2] < h + 4).all()
+                ):
+                    polygon[:, ::2] = np.clip(polygon[:, ::2], 0, w)
+                    polygon[:, 1::2] = np.clip(polygon[:, 1::2], 0, h)
+                    valid_masks_list.append(polygon)
+                    valid_tags_list.append(ignore_tags[ind])
+
+            results["polys"] = np.array(valid_masks_list)
+            results["ignore_tags"] = valid_tags_list
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
+
+class RandomRotatePolyInstances:
+    def __init__(
+        self,
+        rotate_ratio=0.5,
+        max_angle=10,
+        pad_with_fixed_color=False,
+        pad_value=(0, 0, 0),
+        **kwargs,
+    ):
+        """Randomly rotate images and polygon masks.
+
+        Args:
+            rotate_ratio (float): The ratio of samples to operate rotation.
+            max_angle (int): The maximum rotation angle.
+            pad_with_fixed_color (bool): The flag for whether to pad rotated
+               image with fixed value. If set to False, the rotated image will
+               be padded onto cropped image.
+            pad_value (tuple(int)): The color value for padding rotated image.
+        """
+        self.rotate_ratio = rotate_ratio
+        self.max_angle = max_angle
+        self.pad_with_fixed_color = pad_with_fixed_color
+        self.pad_value = pad_value
+
+    def rotate(self, center, points, theta, center_shift=(0, 0)):
+        # rotate points.
+        (center_x, center_y) = center
+        center_y = -center_y
+        x, y = points[:, ::2], points[:, 1::2]
+        y = -y
+
+        theta = theta / 180 * math.pi
+        cos = math.cos(theta)
+        sin = math.sin(theta)
+
+        x = x - center_x
+        y = y - center_y
+
+        _x = center_x + x * cos - y * sin + center_shift[0]
+        _y = -(center_y + x * sin + y * cos) + center_shift[1]
+
+        points[:, ::2], points[:, 1::2] = _x, _y
+        return points
+
+    def cal_canvas_size(self, ori_size, degree):
+        assert isinstance(ori_size, tuple)
+        angle = degree * math.pi / 180.0
+        h, w = ori_size[:2]
+
+        cos = math.cos(angle)
+        sin = math.sin(angle)
+        canvas_h = int(w * math.fabs(sin) + h * math.fabs(cos))
+        canvas_w = int(w * math.fabs(cos) + h * math.fabs(sin))
+
+        canvas_size = (canvas_h, canvas_w)
+        return canvas_size
+
+    def sample_angle(self, max_angle):
+        angle = np.random.random_sample() * 2 * max_angle - max_angle
+        return angle
+
+    def rotate_img(self, img, angle, canvas_size):
+        h, w = img.shape[:2]
+        rotation_matrix = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
+        rotation_matrix[0, 2] += int((canvas_size[1] - w) / 2)
+        rotation_matrix[1, 2] += int((canvas_size[0] - h) / 2)
+
+        if self.pad_with_fixed_color:
+            target_img = cv2.warpAffine(
+                img,
+                rotation_matrix,
+                (canvas_size[1], canvas_size[0]),
+                flags=cv2.INTER_NEAREST,
+                borderValue=self.pad_value,
+            )
+        else:
+            mask = np.zeros_like(img)
+            (h_ind, w_ind) = (
+                np.random.randint(0, h * 7 // 8),
+                np.random.randint(0, w * 7 // 8),
+            )
+            img_cut = img[h_ind : (h_ind + h // 9), w_ind : (w_ind + w // 9)]
+            img_cut = cv2.resize(img_cut, (canvas_size[1], canvas_size[0]))
+
+            mask = cv2.warpAffine(
+                mask,
+                rotation_matrix,
+                (canvas_size[1], canvas_size[0]),
+                borderValue=[1, 1, 1],
+            )
+            target_img = cv2.warpAffine(
+                img,
+                rotation_matrix,
+                (canvas_size[1], canvas_size[0]),
+                borderValue=[0, 0, 0],
+            )
+            target_img = target_img + img_cut * mask
+
+        return target_img
+
+    def __call__(self, results):
+        if np.random.random_sample() < self.rotate_ratio:
+            image = results["image"]
+            polygons = results["polys"]
+            h, w = image.shape[:2]
+
+            angle = self.sample_angle(self.max_angle)
+            canvas_size = self.cal_canvas_size((h, w), angle)
+            center_shift = (
+                int((canvas_size[1] - w) / 2),
+                int((canvas_size[0] - h) / 2),
+            )
+            image = self.rotate_img(image, angle, canvas_size)
+            results["image"] = image
+            # rotate polygons
+            rotated_masks = []
+            for mask in polygons:
+                rotated_mask = self.rotate((w / 2, h / 2), mask, angle, center_shift)
+                rotated_masks.append(rotated_mask)
+            results["polys"] = np.array(rotated_masks)
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str
+
+
+class SquareResizePad:
+    def __init__(
+        self,
+        target_size,
+        pad_ratio=0.6,
+        pad_with_fixed_color=False,
+        pad_value=(0, 0, 0),
+        **kwargs,
+    ):
+        """Resize or pad images to be square shape.
+
+        Args:
+            target_size (int): The target size of square shaped image.
+            pad_with_fixed_color (bool): The flag for whether to pad rotated
+               image with fixed value. If set to False, the rescales image will
+               be padded onto cropped image.
+            pad_value (tuple(int)): The color value for padding rotated image.
+        """
+        assert isinstance(target_size, int)
+        assert isinstance(pad_ratio, float)
+        assert isinstance(pad_with_fixed_color, bool)
+        assert isinstance(pad_value, tuple)
+
+        self.target_size = target_size
+        self.pad_ratio = pad_ratio
+        self.pad_with_fixed_color = pad_with_fixed_color
+        self.pad_value = pad_value
+
+    def resize_img(self, img, keep_ratio=True):
+        h, w, _ = img.shape
+        if keep_ratio:
+            t_h = self.target_size if h >= w else int(h * self.target_size / w)
+            t_w = self.target_size if h <= w else int(w * self.target_size / h)
+        else:
+            t_h = t_w = self.target_size
+        img = cv2.resize(img, (t_w, t_h))
+        return img, (t_h, t_w)
+
+    def square_pad(self, img):
+        h, w = img.shape[:2]
+        if h == w:
+            return img, (0, 0)
+        pad_size = max(h, w)
+        if self.pad_with_fixed_color:
+            expand_img = np.ones((pad_size, pad_size, 3), dtype=np.uint8)
+            expand_img[:] = self.pad_value
+        else:
+            (h_ind, w_ind) = (
+                np.random.randint(0, h * 7 // 8),
+                np.random.randint(0, w * 7 // 8),
+            )
+            img_cut = img[h_ind : (h_ind + h // 9), w_ind : (w_ind + w // 9)]
+            expand_img = cv2.resize(img_cut, (pad_size, pad_size))
+        if h > w:
+            y0, x0 = 0, (h - w) // 2
+        else:
+            y0, x0 = (w - h) // 2, 0
+        expand_img[y0 : y0 + h, x0 : x0 + w] = img
+        offset = (x0, y0)
+
+        return expand_img, offset
+
+    def square_pad_mask(self, points, offset):
+        x0, y0 = offset
+        pad_points = points.copy()
+        pad_points[::2] = pad_points[::2] + x0
+        pad_points[1::2] = pad_points[1::2] + y0
+        return pad_points
+
+    def __call__(self, results):
+        image = results["image"]
+        polygons = results["polys"]
+        h, w = image.shape[:2]
+
+        if np.random.random_sample() < self.pad_ratio:
+            image, out_size = self.resize_img(image, keep_ratio=True)
+            image, offset = self.square_pad(image)
+        else:
+            image, out_size = self.resize_img(image, keep_ratio=False)
+            offset = (0, 0)
+        results["image"] = image
+        try:
+            polygons[:, :, 0::2] = polygons[:, :, 0::2] * out_size[1] / w + offset[0]
+            polygons[:, :, 1::2] = polygons[:, :, 1::2] * out_size[0] / h + offset[1]
+        except:
+            pass
+        results["polys"] = polygons
+
+        return results
+
+    def __repr__(self):
+        repr_str = self.__class__.__name__
+        return repr_str

ppocr/data/imaug/fce_targets.py

@@ -0,0 +1,697 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/fcenet_targets.py
+"""
+
+import cv2
+import numpy as np
+from numpy.fft import fft
+from numpy.linalg import norm
+import sys
+
+
+def vector_slope(vec):
+    assert len(vec) == 2
+    return abs(vec[1] / (vec[0] + 1e-8))
+
+
+class FCENetTargets:
+    """Generate the ground truth targets of FCENet: Fourier Contour Embedding
+    for Arbitrary-Shaped Text Detection.
+
+    [https://arxiv.org/abs/2104.10442]
+
+    Args:
+        fourier_degree (int): The maximum Fourier transform degree k.
+        resample_step (float): The step size for resampling the text center
+            line (TCL). It's better not to exceed half of the minimum width.
+        center_region_shrink_ratio (float): The shrink ratio of text center
+            region.
+        level_size_divisors (tuple(int)): The downsample ratio on each level.
+        level_proportion_range (tuple(tuple(int))): The range of text sizes
+            assigned to each level.
+    """
+
+    def __init__(
+        self,
+        fourier_degree=5,
+        resample_step=4.0,
+        center_region_shrink_ratio=0.3,
+        level_size_divisors=(8, 16, 32),
+        level_proportion_range=((0, 0.25), (0.2, 0.65), (0.55, 1.0)),
+        orientation_thr=2.0,
+        **kwargs,
+    ):
+        super().__init__()
+        assert isinstance(level_size_divisors, tuple)
+        assert isinstance(level_proportion_range, tuple)
+        assert len(level_size_divisors) == len(level_proportion_range)
+        self.fourier_degree = fourier_degree
+        self.resample_step = resample_step
+        self.center_region_shrink_ratio = center_region_shrink_ratio
+        self.level_size_divisors = level_size_divisors
+        self.level_proportion_range = level_proportion_range
+
+        self.orientation_thr = orientation_thr
+
+    def vector_angle(self, vec1, vec2):
+        if vec1.ndim > 1:
+            unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8).reshape((-1, 1))
+        else:
+            unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8)
+        if vec2.ndim > 1:
+            unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8).reshape((-1, 1))
+        else:
+            unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8)
+        return np.arccos(np.clip(np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0))
+
+    def resample_line(self, line, n):
+        """Resample n points on a line.
+
+        Args:
+            line (ndarray): The points composing a line.
+            n (int): The resampled points number.
+
+        Returns:
+            resampled_line (ndarray): The points composing the resampled line.
+        """
+
+        assert line.ndim == 2
+        assert line.shape[0] >= 2
+        assert line.shape[1] == 2
+        assert isinstance(n, int)
+        assert n > 0
+
+        length_list = [norm(line[i + 1] - line[i]) for i in range(len(line) - 1)]
+        total_length = sum(length_list)
+        length_cumsum = np.cumsum([0.0] + length_list)
+        delta_length = total_length / (float(n) + 1e-8)
+
+        current_edge_ind = 0
+        resampled_line = [line[0]]
+
+        for i in range(1, n):
+            current_line_len = i * delta_length
+
+            while (
+                current_edge_ind + 1 < len(length_cumsum)
+                and current_line_len >= length_cumsum[current_edge_ind + 1]
+            ):
+                current_edge_ind += 1
+
+            current_edge_end_shift = current_line_len - length_cumsum[current_edge_ind]
+
+            if current_edge_ind >= len(length_list):
+                break
+            end_shift_ratio = current_edge_end_shift / length_list[current_edge_ind]
+            current_point = (
+                line[current_edge_ind]
+                + (line[current_edge_ind + 1] - line[current_edge_ind])
+                * end_shift_ratio
+            )
+            resampled_line.append(current_point)
+        resampled_line.append(line[-1])
+        resampled_line = np.array(resampled_line)
+
+        return resampled_line
+
+    def reorder_poly_edge(self, points):
+        """Get the respective points composing head edge, tail edge, top
+        sideline and bottom sideline.
+
+        Args:
+            points (ndarray): The points composing a text polygon.
+
+        Returns:
+            head_edge (ndarray): The two points composing the head edge of text
+                polygon.
+            tail_edge (ndarray): The two points composing the tail edge of text
+                polygon.
+            top_sideline (ndarray): The points composing top curved sideline of
+                text polygon.
+            bot_sideline (ndarray): The points composing bottom curved sideline
+                of text polygon.
+        """
+
+        assert points.ndim == 2
+        assert points.shape[0] >= 4
+        assert points.shape[1] == 2
+
+        head_inds, tail_inds = self.find_head_tail(points, self.orientation_thr)
+        head_edge, tail_edge = points[head_inds], points[tail_inds]
+
+        pad_points = np.vstack([points, points])
+        if tail_inds[1] < 1:
+            tail_inds[1] = len(points)
+        sideline1 = pad_points[head_inds[1] : tail_inds[1]]
+        sideline2 = pad_points[tail_inds[1] : (head_inds[1] + len(points))]
+        sideline_mean_shift = np.mean(sideline1, axis=0) - np.mean(sideline2, axis=0)
+
+        if sideline_mean_shift[1] > 0:
+            top_sideline, bot_sideline = sideline2, sideline1
+        else:
+            top_sideline, bot_sideline = sideline1, sideline2
+
+        return head_edge, tail_edge, top_sideline, bot_sideline
+
+    def find_head_tail(self, points, orientation_thr):
+        """Find the head edge and tail edge of a text polygon.
+
+        Args:
+            points (ndarray): The points composing a text polygon.
+            orientation_thr (float): The threshold for distinguishing between
+                head edge and tail edge among the horizontal and vertical edges
+                of a quadrangle.
+
+        Returns:
+            head_inds (list): The indexes of two points composing head edge.
+            tail_inds (list): The indexes of two points composing tail edge.
+        """
+
+        assert points.ndim == 2
+        assert points.shape[0] >= 4
+        assert points.shape[1] == 2
+        assert isinstance(orientation_thr, float)
+
+        if len(points) > 4:
+            pad_points = np.vstack([points, points[0]])
+            edge_vec = pad_points[1:] - pad_points[:-1]
+
+            theta_sum = []
+            adjacent_vec_theta = []
+            for i, edge_vec1 in enumerate(edge_vec):
+                adjacent_ind = [x % len(edge_vec) for x in [i - 1, i + 1]]
+                adjacent_edge_vec = edge_vec[adjacent_ind]
+                temp_theta_sum = np.sum(self.vector_angle(edge_vec1, adjacent_edge_vec))
+                temp_adjacent_theta = self.vector_angle(
+                    adjacent_edge_vec[0], adjacent_edge_vec[1]
+                )
+                theta_sum.append(temp_theta_sum)
+                adjacent_vec_theta.append(temp_adjacent_theta)
+            theta_sum_score = np.array(theta_sum) / np.pi
+            adjacent_theta_score = np.array(adjacent_vec_theta) / np.pi
+            poly_center = np.mean(points, axis=0)
+            edge_dist = np.maximum(
+                norm(pad_points[1:] - poly_center, axis=-1),
+                norm(pad_points[:-1] - poly_center, axis=-1),
+            )
+            dist_score = edge_dist / np.max(edge_dist)
+            position_score = np.zeros(len(edge_vec))
+            score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score
+            score += 0.35 * dist_score
+            if len(points) % 2 == 0:
+                position_score[(len(score) // 2 - 1)] += 1
+                position_score[-1] += 1
+            score += 0.1 * position_score
+            pad_score = np.concatenate([score, score])
+            score_matrix = np.zeros((len(score), len(score) - 3))
+            x = np.arange(len(score) - 3) / float(len(score) - 4)
+            gaussian = (
+                1.0
+                / (np.sqrt(2.0 * np.pi) * 0.5)
+                * np.exp(-np.power((x - 0.5) / 0.5, 2.0) / 2)
+            )
+            gaussian = gaussian / np.max(gaussian)
+            for i in range(len(score)):
+                score_matrix[i, :] = (
+                    score[i]
+                    + pad_score[(i + 2) : (i + len(score) - 1)] * gaussian * 0.3
+                )
+
+            head_start, tail_increment = np.unravel_index(
+                score_matrix.argmax(), score_matrix.shape
+            )
+            tail_start = (head_start + tail_increment + 2) % len(points)
+            head_end = (head_start + 1) % len(points)
+            tail_end = (tail_start + 1) % len(points)
+
+            if head_end > tail_end:
+                head_start, tail_start = tail_start, head_start
+                head_end, tail_end = tail_end, head_end
+            head_inds = [head_start, head_end]
+            tail_inds = [tail_start, tail_end]
+        else:
+            if vector_slope(points[1] - points[0]) + vector_slope(
+                points[3] - points[2]
+            ) < vector_slope(points[2] - points[1]) + vector_slope(
+                points[0] - points[3]
+            ):
+                horizontal_edge_inds = [[0, 1], [2, 3]]
+                vertical_edge_inds = [[3, 0], [1, 2]]
+            else:
+                horizontal_edge_inds = [[3, 0], [1, 2]]
+                vertical_edge_inds = [[0, 1], [2, 3]]
+
+            vertical_len_sum = norm(
+                points[vertical_edge_inds[0][0]] - points[vertical_edge_inds[0][1]]
+            ) + norm(
+                points[vertical_edge_inds[1][0]] - points[vertical_edge_inds[1][1]]
+            )
+            horizontal_len_sum = norm(
+                points[horizontal_edge_inds[0][0]] - points[horizontal_edge_inds[0][1]]
+            ) + norm(
+                points[horizontal_edge_inds[1][0]] - points[horizontal_edge_inds[1][1]]
+            )
+
+            if vertical_len_sum > horizontal_len_sum * orientation_thr:
+                head_inds = horizontal_edge_inds[0]
+                tail_inds = horizontal_edge_inds[1]
+            else:
+                head_inds = vertical_edge_inds[0]
+                tail_inds = vertical_edge_inds[1]
+
+        return head_inds, tail_inds
+
+    def resample_sidelines(self, sideline1, sideline2, resample_step):
+        """Resample two sidelines to be of the same points number according to
+        step size.
+
+        Args:
+            sideline1 (ndarray): The points composing a sideline of a text
+                polygon.
+            sideline2 (ndarray): The points composing another sideline of a
+                text polygon.
+            resample_step (float): The resampled step size.
+
+        Returns:
+            resampled_line1 (ndarray): The resampled line 1.
+            resampled_line2 (ndarray): The resampled line 2.
+        """
+
+        assert sideline1.ndim == sideline2.ndim == 2
+        assert sideline1.shape[1] == sideline2.shape[1] == 2
+        assert sideline1.shape[0] >= 2
+        assert sideline2.shape[0] >= 2
+        assert isinstance(resample_step, float)
+
+        length1 = sum(
+            [norm(sideline1[i + 1] - sideline1[i]) for i in range(len(sideline1) - 1)]
+        )
+        length2 = sum(
+            [norm(sideline2[i + 1] - sideline2[i]) for i in range(len(sideline2) - 1)]
+        )
+
+        total_length = (length1 + length2) / 2
+        resample_point_num = max(int(float(total_length) / resample_step), 1)
+
+        resampled_line1 = self.resample_line(sideline1, resample_point_num)
+        resampled_line2 = self.resample_line(sideline2, resample_point_num)
+
+        return resampled_line1, resampled_line2
+
+    def generate_center_region_mask(self, img_size, text_polys):
+        """Generate text center region mask.
+
+        Args:
+            img_size (tuple): The image size of (height, width).
+            text_polys (list[list[ndarray]]): The list of text polygons.
+
+        Returns:
+            center_region_mask (ndarray): The text center region mask.
+        """
+
+        assert isinstance(img_size, tuple)
+        # assert check_argument.is_2dlist(text_polys)
+
+        h, w = img_size
+
+        center_region_mask = np.zeros((h, w), np.uint8)
+
+        center_region_boxes = []
+        for poly in text_polys:
+            # assert len(poly) == 1
+            polygon_points = poly.reshape(-1, 2)
+            _, _, top_line, bot_line = self.reorder_poly_edge(polygon_points)
+            resampled_top_line, resampled_bot_line = self.resample_sidelines(
+                top_line, bot_line, self.resample_step
+            )
+            resampled_bot_line = resampled_bot_line[::-1]
+            if len(resampled_top_line) != len(resampled_bot_line):
+                continue
+            center_line = (resampled_top_line + resampled_bot_line) / 2
+
+            line_head_shrink_len = (
+                norm(resampled_top_line[0] - resampled_bot_line[0]) / 4.0
+            )
+            line_tail_shrink_len = (
+                norm(resampled_top_line[-1] - resampled_bot_line[-1]) / 4.0
+            )
+            head_shrink_num = int(line_head_shrink_len // self.resample_step)
+            tail_shrink_num = int(line_tail_shrink_len // self.resample_step)
+            if len(center_line) > head_shrink_num + tail_shrink_num + 2:
+                center_line = center_line[
+                    head_shrink_num : len(center_line) - tail_shrink_num
+                ]
+                resampled_top_line = resampled_top_line[
+                    head_shrink_num : len(resampled_top_line) - tail_shrink_num
+                ]
+                resampled_bot_line = resampled_bot_line[
+                    head_shrink_num : len(resampled_bot_line) - tail_shrink_num
+                ]
+
+            for i in range(0, len(center_line) - 1):
+                tl = (
+                    center_line[i]
+                    + (resampled_top_line[i] - center_line[i])
+                    * self.center_region_shrink_ratio
+                )
+                tr = (
+                    center_line[i + 1]
+                    + (resampled_top_line[i + 1] - center_line[i + 1])
+                    * self.center_region_shrink_ratio
+                )
+                br = (
+                    center_line[i + 1]
+                    + (resampled_bot_line[i + 1] - center_line[i + 1])
+                    * self.center_region_shrink_ratio
+                )
+                bl = (
+                    center_line[i]
+                    + (resampled_bot_line[i] - center_line[i])
+                    * self.center_region_shrink_ratio
+                )
+                current_center_box = np.vstack([tl, tr, br, bl]).astype(np.int32)
+                center_region_boxes.append(current_center_box)
+
+        cv2.fillPoly(center_region_mask, center_region_boxes, 1)
+        return center_region_mask
+
+    def resample_polygon(self, polygon, n=400):
+        """Resample one polygon with n points on its boundary.
+
+        Args:
+            polygon (list[float]): The input polygon.
+            n (int): The number of resampled points.
+        Returns:
+            resampled_polygon (list[float]): The resampled polygon.
+        """
+        length = []
+
+        for i in range(len(polygon)):
+            p1 = polygon[i]
+            if i == len(polygon) - 1:
+                p2 = polygon[0]
+            else:
+                p2 = polygon[i + 1]
+            length.append(((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2) ** 0.5)
+
+        total_length = sum(length)
+        n_on_each_line = (np.array(length) / (total_length + 1e-8)) * n
+        n_on_each_line = n_on_each_line.astype(np.int32)
+        new_polygon = []
+
+        for i in range(len(polygon)):
+            num = n_on_each_line[i]
+            p1 = polygon[i]
+            if i == len(polygon) - 1:
+                p2 = polygon[0]
+            else:
+                p2 = polygon[i + 1]
+
+            if num == 0:
+                continue
+
+            dxdy = (p2 - p1) / num
+            for j in range(num):
+                point = p1 + dxdy * j
+                new_polygon.append(point)
+
+        return np.array(new_polygon)
+
+    def normalize_polygon(self, polygon):
+        """Normalize one polygon so that its start point is at right most.
+
+        Args:
+            polygon (list[float]): The origin polygon.
+        Returns:
+            new_polygon (lost[float]): The polygon with start point at right.
+        """
+        temp_polygon = polygon - polygon.mean(axis=0)
+        x = np.abs(temp_polygon[:, 0])
+        y = temp_polygon[:, 1]
+        index_x = np.argsort(x)
+        index_y = np.argmin(y[index_x[:8]])
+        index = index_x[index_y]
+        new_polygon = np.concatenate([polygon[index:], polygon[:index]])
+        return new_polygon
+
+    def poly2fourier(self, polygon, fourier_degree):
+        """Perform Fourier transformation to generate Fourier coefficients ck
+        from polygon.
+
+        Args:
+            polygon (ndarray): An input polygon.
+            fourier_degree (int): The maximum Fourier degree K.
+        Returns:
+            c (ndarray(complex)): Fourier coefficients.
+        """
+        points = polygon[:, 0] + polygon[:, 1] * 1j
+        c_fft = fft(points) / len(points)
+        c = np.hstack((c_fft[-fourier_degree:], c_fft[: fourier_degree + 1]))
+        return c
+
+    def clockwise(self, c, fourier_degree):
+        """Make sure the polygon reconstructed from Fourier coefficients c in
+        the clockwise direction.
+
+        Args:
+            polygon (list[float]): The origin polygon.
+        Returns:
+            new_polygon (lost[float]): The polygon in clockwise point order.
+        """
+        if np.abs(c[fourier_degree + 1]) > np.abs(c[fourier_degree - 1]):
+            return c
+        elif np.abs(c[fourier_degree + 1]) < np.abs(c[fourier_degree - 1]):
+            return c[::-1]
+        else:
+            if np.abs(c[fourier_degree + 2]) > np.abs(c[fourier_degree - 2]):
+                return c
+            else:
+                return c[::-1]
+
+    def cal_fourier_signature(self, polygon, fourier_degree):
+        """Calculate Fourier signature from input polygon.
+
+        Args:
+              polygon (ndarray): The input polygon.
+              fourier_degree (int): The maximum Fourier degree K.
+        Returns:
+              fourier_signature (ndarray): An array shaped (2k+1, 2) containing
+                  real part and image part of 2k+1 Fourier coefficients.
+        """
+        resampled_polygon = self.resample_polygon(polygon)
+        resampled_polygon = self.normalize_polygon(resampled_polygon)
+
+        fourier_coeff = self.poly2fourier(resampled_polygon, fourier_degree)
+        fourier_coeff = self.clockwise(fourier_coeff, fourier_degree)
+
+        real_part = np.real(fourier_coeff).reshape((-1, 1))
+        image_part = np.imag(fourier_coeff).reshape((-1, 1))
+        fourier_signature = np.hstack([real_part, image_part])
+
+        return fourier_signature
+
+    def generate_fourier_maps(self, img_size, text_polys):
+        """Generate Fourier coefficient maps.
+
+        Args:
+            img_size (tuple): The image size of (height, width).
+            text_polys (list[list[ndarray]]): The list of text polygons.
+
+        Returns:
+            fourier_real_map (ndarray): The Fourier coefficient real part maps.
+            fourier_image_map (ndarray): The Fourier coefficient image part
+                maps.
+        """
+
+        assert isinstance(img_size, tuple)
+
+        h, w = img_size
+        k = self.fourier_degree
+        real_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
+        imag_map = np.zeros((k * 2 + 1, h, w), dtype=np.float32)
+
+        for poly in text_polys:
+            mask = np.zeros((h, w), dtype=np.uint8)
+            polygon = np.array(poly).reshape((1, -1, 2))
+            cv2.fillPoly(mask, polygon.astype(np.int32), 1)
+            fourier_coeff = self.cal_fourier_signature(polygon[0], k)
+            for i in range(-k, k + 1):
+                if i != 0:
+                    real_map[i + k, :, :] = (
+                        mask * fourier_coeff[i + k, 0]
+                        + (1 - mask) * real_map[i + k, :, :]
+                    )
+                    imag_map[i + k, :, :] = (
+                        mask * fourier_coeff[i + k, 1]
+                        + (1 - mask) * imag_map[i + k, :, :]
+                    )
+                else:
+                    yx = np.argwhere(mask > 0.5)
+                    k_ind = np.ones((len(yx)), dtype=np.int64) * k
+                    y, x = yx[:, 0], yx[:, 1]
+                    real_map[k_ind, y, x] = fourier_coeff[k, 0] - x
+                    imag_map[k_ind, y, x] = fourier_coeff[k, 1] - y
+
+        return real_map, imag_map
+
+    def generate_text_region_mask(self, img_size, text_polys):
+        """Generate text center region mask and geometry attribute maps.
+
+        Args:
+            img_size (tuple): The image size (height, width).
+            text_polys (list[list[ndarray]]): The list of text polygons.
+
+        Returns:
+            text_region_mask (ndarray): The text region mask.
+        """
+
+        assert isinstance(img_size, tuple)
+
+        h, w = img_size
+        text_region_mask = np.zeros((h, w), dtype=np.uint8)
+
+        for poly in text_polys:
+            polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
+            cv2.fillPoly(text_region_mask, polygon, 1)
+
+        return text_region_mask
+
+    def generate_effective_mask(self, mask_size: tuple, polygons_ignore):
+        """Generate effective mask by setting the ineffective regions to 0 and
+        effective regions to 1.
+
+        Args:
+            mask_size (tuple): The mask size.
+            polygons_ignore (list[[ndarray]]: The list of ignored text
+                polygons.
+
+        Returns:
+            mask (ndarray): The effective mask of (height, width).
+        """
+
+        mask = np.ones(mask_size, dtype=np.uint8)
+
+        for poly in polygons_ignore:
+            instance = poly.reshape(-1, 2).astype(np.int32).reshape(1, -1, 2)
+            cv2.fillPoly(mask, instance, 0)
+
+        return mask
+
+    def generate_level_targets(self, img_size, text_polys, ignore_polys):
+        """Generate ground truth target on each level.
+
+        Args:
+            img_size (list[int]): Shape of input image.
+            text_polys (list[list[ndarray]]): A list of ground truth polygons.
+            ignore_polys (list[list[ndarray]]): A list of ignored polygons.
+        Returns:
+            level_maps (list(ndarray)): A list of ground target on each level.
+        """
+        h, w = img_size
+        lv_size_divs = self.level_size_divisors
+        lv_proportion_range = self.level_proportion_range
+        lv_text_polys = [[] for i in range(len(lv_size_divs))]
+        lv_ignore_polys = [[] for i in range(len(lv_size_divs))]
+        level_maps = []
+        for poly in text_polys:
+            polygon = np.array(poly, dtype=np.int32).reshape((1, -1, 2))
+            _, _, box_w, box_h = cv2.boundingRect(polygon)
+            proportion = max(box_h, box_w) / (h + 1e-8)
+
+            for ind, proportion_range in enumerate(lv_proportion_range):
+                if proportion_range[0] < proportion < proportion_range[1]:
+                    lv_text_polys[ind].append(poly / lv_size_divs[ind])
+
+        for ignore_poly in ignore_polys:
+            polygon = np.array(ignore_poly, dtype=np.int32).reshape((1, -1, 2))
+            _, _, box_w, box_h = cv2.boundingRect(polygon)
+            proportion = max(box_h, box_w) / (h + 1e-8)
+
+            for ind, proportion_range in enumerate(lv_proportion_range):
+                if proportion_range[0] < proportion < proportion_range[1]:
+                    lv_ignore_polys[ind].append(ignore_poly / lv_size_divs[ind])
+
+        for ind, size_divisor in enumerate(lv_size_divs):
+            current_level_maps = []
+            level_img_size = (h // size_divisor, w // size_divisor)
+
+            text_region = self.generate_text_region_mask(
+                level_img_size, lv_text_polys[ind]
+            )[None]
+            current_level_maps.append(text_region)
+
+            center_region = self.generate_center_region_mask(
+                level_img_size, lv_text_polys[ind]
+            )[None]
+            current_level_maps.append(center_region)
+
+            effective_mask = self.generate_effective_mask(
+                level_img_size, lv_ignore_polys[ind]
+            )[None]
+            current_level_maps.append(effective_mask)
+
+            fourier_real_map, fourier_image_maps = self.generate_fourier_maps(
+                level_img_size, lv_text_polys[ind]
+            )
+            current_level_maps.append(fourier_real_map)
+            current_level_maps.append(fourier_image_maps)
+
+            level_maps.append(np.concatenate(current_level_maps))
+
+        return level_maps
+
+    def generate_targets(self, results):
+        """Generate the ground truth targets for FCENet.
+
+        Args:
+            results (dict): The input result dictionary.
+
+        Returns:
+            results (dict): The output result dictionary.
+        """
+
+        assert isinstance(results, dict)
+        image = results["image"]
+        polygons = results["polys"]
+        ignore_tags = results["ignore_tags"]
+        h, w, _ = image.shape
+
+        polygon_masks = []
+        polygon_masks_ignore = []
+        for tag, polygon in zip(ignore_tags, polygons):
+            if tag is True:
+                polygon_masks_ignore.append(polygon)
+            else:
+                polygon_masks.append(polygon)
+
+        level_maps = self.generate_level_targets(
+            (h, w), polygon_masks, polygon_masks_ignore
+        )
+
+        mapping = {
+            "p3_maps": level_maps[0],
+            "p4_maps": level_maps[1],
+            "p5_maps": level_maps[2],
+        }
+        for key, value in mapping.items():
+            results[key] = value
+
+        return results
+
+    def __call__(self, results):
+        results = self.generate_targets(results)
+        return results

ppocr/data/imaug/iaa_augment.py

@@ -0,0 +1,213 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/iaa_augment.py
+"""
+import os
+
+# Prevent automatic updates in Albumentations for stability in augmentation behavior
+os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
+
+import numpy as np
+import albumentations as A
+from albumentations.core.transforms_interface import DualTransform
+from albumentations.augmentations.geometric import functional as fgeometric
+from packaging import version
+
+ALBU_VERSION = version.parse(A.__version__)
+IS_ALBU_NEW_VERSION = ALBU_VERSION >= version.parse("1.4.15")
+
+
+# Custom resize transformation mimicking Imgaug's behavior with scaling
+class ImgaugLikeResize(DualTransform):
+    def __init__(self, scale_range=(0.5, 3.0), interpolation=1, p=1.0):
+        super(ImgaugLikeResize, self).__init__(p)
+        self.scale_range = scale_range
+        self.interpolation = interpolation
+
+    # Resize the image based on a randomly chosen scale within the scale range
+    def apply(self, img, scale=1.0, **params):
+        height, width = img.shape[:2]
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+
+        if IS_ALBU_NEW_VERSION:
+            return fgeometric.resize(
+                img, (new_height, new_width), interpolation=self.interpolation
+            )
+        return fgeometric.resize(
+            img, new_height, new_width, interpolation=self.interpolation
+        )
+
+    # Apply the same scaling transformation to keypoints (e.g., polygon points)
+    def apply_to_keypoints(self, keypoints, scale=1.0, **params):
+        return np.array(
+            [(x * scale, y * scale) + tuple(rest) for x, y, *rest in keypoints]
+        )
+
+    # Get random scale parameter within the specified range
+    def get_params(self):
+        scale = np.random.uniform(self.scale_range[0], self.scale_range[1])
+        return {"scale": scale}
+
+
+# Builder class to translate custom augmenter arguments into Albumentations-compatible format
+class AugmenterBuilder(object):
+    def __init__(self):
+        # Map common Imgaug transformations to equivalent Albumentations transforms
+        self.imgaug_to_albu = {
+            "Fliplr": "HorizontalFlip",
+            "Flipud": "VerticalFlip",
+            "Affine": "Affine",
+            # Additional mappings can be added here if needed
+        }
+
+    # Recursive method to construct augmentation pipeline based on provided arguments
+    def build(self, args, root=True):
+        if args is None or len(args) == 0:
+            return None
+        elif isinstance(args, list):
+            # Build the full augmentation sequence if it's a root-level call
+            if root:
+                sequence = [self.build(value, root=False) for value in args]
+                return A.Compose(
+                    sequence,
+                    keypoint_params=A.KeypointParams(
+                        format="xy", remove_invisible=False
+                    ),
+                )
+            else:
+                # Build individual augmenters for nested arguments
+                augmenter_type = args[0]
+                augmenter_args = args[1] if len(args) > 1 else {}
+                augmenter_args_mapped = self.map_arguments(
+                    augmenter_type, augmenter_args
+                )
+                augmenter_type_mapped = self.imgaug_to_albu.get(
+                    augmenter_type, augmenter_type
+                )
+                if augmenter_type_mapped == "Resize":
+                    return ImgaugLikeResize(**augmenter_args_mapped)
+                else:
+                    cls = getattr(A, augmenter_type_mapped)
+                    return cls(
+                        **{
+                            k: self.to_tuple_if_list(v)
+                            for k, v in augmenter_args_mapped.items()
+                        }
+                    )
+        elif isinstance(args, dict):
+            # Process individual transformation specified as dictionary
+            augmenter_type = args["type"]
+            augmenter_args = args.get("args", {})
+            augmenter_args_mapped = self.map_arguments(augmenter_type, augmenter_args)
+            augmenter_type_mapped = self.imgaug_to_albu.get(
+                augmenter_type, augmenter_type
+            )
+            if augmenter_type_mapped == "Resize":
+                return ImgaugLikeResize(**augmenter_args_mapped)
+            else:
+                cls = getattr(A, augmenter_type_mapped)
+                return cls(
+                    **{
+                        k: self.to_tuple_if_list(v)
+                        for k, v in augmenter_args_mapped.items()
+                    }
+                )
+        else:
+            raise RuntimeError("Unknown augmenter arg: " + str(args))
+
+    # Map arguments to expected format for each augmenter type
+    def map_arguments(self, augmenter_type, augmenter_args):
+        augmenter_args = augmenter_args.copy()  # Avoid modifying the original arguments
+        if augmenter_type == "Resize":
+            # Ensure size is a valid 2-element list or tuple
+            size = augmenter_args.get("size")
+            if size:
+                if not isinstance(size, (list, tuple)) or len(size) != 2:
+                    raise ValueError(
+                        f"'size' must be a list or tuple of two numbers, but got {size}"
+                    )
+                min_scale, max_scale = size
+                return {
+                    "scale_range": (min_scale, max_scale),
+                    "interpolation": 1,  # Linear interpolation
+                    "p": 1.0,
+                }
+            else:
+                return {"scale_range": (1.0, 1.0), "interpolation": 1, "p": 1.0}
+        elif augmenter_type == "Affine":
+            # Map rotation to a tuple and ensure p=1.0 to apply transformation
+            rotate = augmenter_args.get("rotate", 0)
+            if isinstance(rotate, list):
+                rotate = tuple(rotate)
+            elif isinstance(rotate, (int, float)):
+                rotate = (float(rotate), float(rotate))
+            augmenter_args["rotate"] = rotate
+            augmenter_args["p"] = 1.0
+            return augmenter_args
+        else:
+            # For other augmenters, ensure 'p' probability is specified
+            p = augmenter_args.get("p", 1.0)
+            augmenter_args["p"] = p
+            return augmenter_args
+
+    # Convert lists to tuples for Albumentations compatibility
+    def to_tuple_if_list(self, obj):
+        if isinstance(obj, list):
+            return tuple(obj)
+        return obj
+
+
+# Wrapper class for image and polygon transformations using Imgaug-style augmentation
+class IaaAugment:
+    def __init__(self, augmenter_args=None, **kwargs):
+        if augmenter_args is None:
+            # Default augmenters if none are specified
+            augmenter_args = [
+                {"type": "Fliplr", "args": {"p": 0.5}},
+                {"type": "Affine", "args": {"rotate": [-10, 10]}},
+                {"type": "Resize", "args": {"size": [0.5, 3]}},
+            ]
+        self.augmenter = AugmenterBuilder().build(augmenter_args)
+
+    # Apply the augmentations to image and polygon data
+    def __call__(self, data):
+        image = data["image"]
+
+        if self.augmenter:
+            # Flatten polygons to individual keypoints for transformation
+            keypoints = []
+            keypoints_lengths = []
+            for poly in data["polys"]:
+                keypoints.extend([tuple(point) for point in poly])
+                keypoints_lengths.append(len(poly))
+
+            # Apply the augmentation pipeline to image and keypoints
+            transformed = self.augmenter(image=image, keypoints=keypoints)
+            data["image"] = transformed["image"]
+
+            # Extract transformed keypoints and reconstruct polygon structures
+            transformed_keypoints = transformed["keypoints"]
+
+            # Reassemble polygons from transformed keypoints
+            new_polys = []
+            idx = 0
+            for length in keypoints_lengths:
+                new_poly = transformed_keypoints[idx : idx + length]
+                new_polys.append(np.array([kp[:2] for kp in new_poly]))
+                idx += length
+            data["polys"] = np.array(new_polys)
+        return data

ppocr/data/imaug/latex_ocr_aug.py

@@ -0,0 +1,183 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+"""
+This code is refer from:
+https://github.com/lukas-blecher/LaTeX-OCR/blob/main/pix2tex/dataset/transforms.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import os
+
+os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
+
+import math
+import cv2
+import numpy as np
+import albumentations as A
+from PIL import Image
+
+
+class LatexTrainTransform:
+    def __init__(self, bitmap_prob=0.04, **kwargs):
+        # your init code
+        self.bitmap_prob = bitmap_prob
+        self.train_transform = A.Compose(
+            [
+                A.Compose(
+                    [
+                        A.ShiftScaleRotate(
+                            shift_limit=0,
+                            scale_limit=(-0.15, 0),
+                            rotate_limit=1,
+                            border_mode=0,
+                            interpolation=3,
+                            value=[255, 255, 255],
+                            p=1,
+                        ),
+                        A.GridDistortion(
+                            distort_limit=0.1,
+                            border_mode=0,
+                            interpolation=3,
+                            value=[255, 255, 255],
+                            p=0.5,
+                        ),
+                    ],
+                    p=0.15,
+                ),
+                A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.3),
+                A.GaussNoise(10, p=0.2),
+                A.RandomBrightnessContrast(0.05, (-0.2, 0), True, p=0.2),
+                A.ImageCompression(95, p=0.3),
+                A.ToGray(always_apply=True),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        if np.random.random() < self.bitmap_prob:
+            img[img != 255] = 0
+        img = self.train_transform(image=img)["image"]
+        data["image"] = img
+        return data
+
+
+class LatexTestTransform:
+    def __init__(self, **kwargs):
+        # your init code
+        self.test_transform = A.Compose(
+            [
+                A.ToGray(always_apply=True),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        img = self.test_transform(image=img)["image"]
+        data["image"] = img
+        return data
+
+
+class MinMaxResize:
+    def __init__(self, min_dimensions=[32, 32], max_dimensions=[672, 192], **kwargs):
+        # your init code
+        self.min_dimensions = min_dimensions
+        self.max_dimensions = max_dimensions
+        # pass
+
+    def pad_(self, img, divable=32):
+        threshold = 128
+        data = np.array(img.convert("LA"))
+        if data[..., -1].var() == 0:
+            data = (data[..., 0]).astype(np.uint8)
+        else:
+            data = (255 - data[..., -1]).astype(np.uint8)
+        data = (data - data.min()) / (data.max() - data.min()) * 255
+        if data.mean() > threshold:
+            # To invert the text to white
+            gray = 255 * (data < threshold).astype(np.uint8)
+        else:
+            gray = 255 * (data > threshold).astype(np.uint8)
+            data = 255 - data
+
+        coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+        a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+        rect = data[b : b + h, a : a + w]
+        im = Image.fromarray(rect).convert("L")
+        dims = []
+        for x in [w, h]:
+            div, mod = divmod(x, divable)
+            dims.append(divable * (div + (1 if mod > 0 else 0)))
+        padded = Image.new("L", dims, 255)
+        padded.paste(im, (0, 0, im.size[0], im.size[1]))
+        return padded
+
+    def minmax_size_(self, img, max_dimensions, min_dimensions):
+        if max_dimensions is not None:
+            ratios = [a / b for a, b in zip(img.size, max_dimensions)]
+            if any([r > 1 for r in ratios]):
+                size = np.array(img.size) // max(ratios)
+                img = img.resize(tuple(size.astype(int)), Image.BILINEAR)
+        if min_dimensions is not None:
+            # hypothesis: there is a dim in img smaller than min_dimensions, and return a proper dim >= min_dimensions
+            padded_size = [
+                max(img_dim, min_dim)
+                for img_dim, min_dim in zip(img.size, min_dimensions)
+            ]
+            if padded_size != list(img.size):  # assert hypothesis
+                padded_im = Image.new("L", padded_size, 255)
+                padded_im.paste(img, img.getbbox())
+                img = padded_im
+        return img
+
+    def __call__(self, data):
+        img = data["image"]
+        h, w = img.shape[:2]
+        if (
+            self.min_dimensions[0] <= w <= self.max_dimensions[0]
+            and self.min_dimensions[1] <= h <= self.max_dimensions[1]
+        ):
+            return data
+        else:
+            im = Image.fromarray(np.uint8(img))
+            im = self.minmax_size_(
+                self.pad_(im), self.max_dimensions, self.min_dimensions
+            )
+            im = np.array(im)
+            im = np.dstack((im, im, im))
+            data["image"] = im
+            return data
+
+
+class LatexImageFormat:
+    def __init__(self, **kwargs):
+        # your init code
+        pass
+
+    def __call__(self, data):
+        img = data["image"]
+        im_h, im_w = img.shape[:2]
+        divide_h = math.ceil(im_h / 16) * 16
+        divide_w = math.ceil(im_w / 16) * 16
+        img = img[:, :, 0]
+        img = np.pad(
+            img, ((0, divide_h - im_h), (0, divide_w - im_w)), constant_values=(1, 1)
+        )
+        img_expanded = img[:, :, np.newaxis].transpose(2, 0, 1)
+        data["image"] = img_expanded
+        return data

ppocr/data/imaug/make_border_map.py

@@ -0,0 +1,179 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_border_map.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import numpy as np
+import cv2
+
+np.seterr(divide="ignore", invalid="ignore")
+import pyclipper
+from shapely.geometry import Polygon
+import sys
+import warnings
+
+warnings.simplefilter("ignore")
+
+__all__ = ["MakeBorderMap"]
+
+
+class MakeBorderMap(object):
+    def __init__(self, shrink_ratio=0.4, thresh_min=0.3, thresh_max=0.7, **kwargs):
+        self.shrink_ratio = shrink_ratio
+        self.thresh_min = thresh_min
+        self.thresh_max = thresh_max
+        if "total_epoch" in kwargs and "epoch" in kwargs and kwargs["epoch"] != "None":
+            self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs["epoch"] / float(
+                kwargs["total_epoch"]
+            )
+
+    def __call__(self, data):
+        img = data["image"]
+        text_polys = data["polys"]
+        ignore_tags = data["ignore_tags"]
+
+        canvas = np.zeros(img.shape[:2], dtype=np.float32)
+        mask = np.zeros(img.shape[:2], dtype=np.float32)
+
+        for i in range(len(text_polys)):
+            if ignore_tags[i]:
+                continue
+            self.draw_border_map(text_polys[i], canvas, mask=mask)
+        canvas = canvas * (self.thresh_max - self.thresh_min) + self.thresh_min
+
+        data["threshold_map"] = canvas
+        data["threshold_mask"] = mask
+        return data
+
+    def draw_border_map(self, polygon, canvas, mask):
+        polygon = np.array(polygon)
+        assert polygon.ndim == 2
+        assert polygon.shape[1] == 2
+
+        polygon_shape = Polygon(polygon)
+        if polygon_shape.area <= 0:
+            return
+        distance = (
+            polygon_shape.area
+            * (1 - np.power(self.shrink_ratio, 2))
+            / polygon_shape.length
+        )
+        subject = [tuple(l) for l in polygon]
+        padding = pyclipper.PyclipperOffset()
+        padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+
+        padded_polygon = np.array(padding.Execute(distance)[0])
+        cv2.fillPoly(mask, [padded_polygon.astype(np.int32)], 1.0)
+
+        xmin = padded_polygon[:, 0].min()
+        xmax = padded_polygon[:, 0].max()
+        ymin = padded_polygon[:, 1].min()
+        ymax = padded_polygon[:, 1].max()
+        width = xmax - xmin + 1
+        height = ymax - ymin + 1
+
+        polygon[:, 0] = polygon[:, 0] - xmin
+        polygon[:, 1] = polygon[:, 1] - ymin
+
+        xs = np.broadcast_to(
+            np.linspace(0, width - 1, num=width).reshape(1, width), (height, width)
+        )
+        ys = np.broadcast_to(
+            np.linspace(0, height - 1, num=height).reshape(height, 1), (height, width)
+        )
+
+        distance_map = np.zeros((polygon.shape[0], height, width), dtype=np.float32)
+        for i in range(polygon.shape[0]):
+            j = (i + 1) % polygon.shape[0]
+            absolute_distance = self._distance(xs, ys, polygon[i], polygon[j])
+            distance_map[i] = np.clip(absolute_distance / distance, 0, 1)
+        distance_map = distance_map.min(axis=0)
+
+        xmin_valid = min(max(0, xmin), canvas.shape[1] - 1)
+        xmax_valid = min(max(0, xmax), canvas.shape[1] - 1)
+        ymin_valid = min(max(0, ymin), canvas.shape[0] - 1)
+        ymax_valid = min(max(0, ymax), canvas.shape[0] - 1)
+        canvas[ymin_valid : ymax_valid + 1, xmin_valid : xmax_valid + 1] = np.fmax(
+            1
+            - distance_map[
+                ymin_valid - ymin : ymax_valid - ymax + height,
+                xmin_valid - xmin : xmax_valid - xmax + width,
+            ],
+            canvas[ymin_valid : ymax_valid + 1, xmin_valid : xmax_valid + 1],
+        )
+
+    def _distance(self, xs, ys, point_1, point_2):
+        """
+        compute the distance from point to a line
+        ys: coordinates in the first axis
+        xs: coordinates in the second axis
+        point_1, point_2: (x, y), the end of the line
+        """
+        height, width = xs.shape[:2]
+        square_distance_1 = np.square(xs - point_1[0]) + np.square(ys - point_1[1])
+        square_distance_2 = np.square(xs - point_2[0]) + np.square(ys - point_2[1])
+        square_distance = np.square(point_1[0] - point_2[0]) + np.square(
+            point_1[1] - point_2[1]
+        )
+
+        cosin = (square_distance - square_distance_1 - square_distance_2) / (
+            2 * np.sqrt(square_distance_1 * square_distance_2)
+        )
+        square_sin = 1 - np.square(cosin)
+        square_sin = np.nan_to_num(square_sin)
+        result = np.sqrt(
+            square_distance_1 * square_distance_2 * square_sin / square_distance
+        )
+
+        result[cosin < 0] = np.sqrt(np.fmin(square_distance_1, square_distance_2))[
+            cosin < 0
+        ]
+        # self.extend_line(point_1, point_2, result)
+        return result
+
+    def extend_line(self, point_1, point_2, result, shrink_ratio):
+        ex_point_1 = (
+            int(round(point_1[0] + (point_1[0] - point_2[0]) * (1 + shrink_ratio))),
+            int(round(point_1[1] + (point_1[1] - point_2[1]) * (1 + shrink_ratio))),
+        )
+        cv2.line(
+            result,
+            tuple(ex_point_1),
+            tuple(point_1),
+            4096.0,
+            1,
+            lineType=cv2.LINE_AA,
+            shift=0,
+        )
+        ex_point_2 = (
+            int(round(point_2[0] + (point_2[0] - point_1[0]) * (1 + shrink_ratio))),
+            int(round(point_2[1] + (point_2[1] - point_1[1]) * (1 + shrink_ratio))),
+        )
+        cv2.line(
+            result,
+            tuple(ex_point_2),
+            tuple(point_2),
+            4096.0,
+            1,
+            lineType=cv2.LINE_AA,
+            shift=0,
+        )
+        return ex_point_1, ex_point_2

ppocr/data/imaug/make_pse_gt.py

@@ -0,0 +1,104 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+__all__ = ["MakePseGt"]
+
+
+class MakePseGt(object):
+    def __init__(self, kernel_num=7, size=640, min_shrink_ratio=0.4, **kwargs):
+        self.kernel_num = kernel_num
+        self.min_shrink_ratio = min_shrink_ratio
+        self.size = size
+
+    def __call__(self, data):
+        image = data["image"]
+        text_polys = data["polys"]
+        ignore_tags = data["ignore_tags"]
+
+        h, w, _ = image.shape
+        short_edge = min(h, w)
+        if short_edge < self.size:
+            # keep short_size >= self.size
+            scale = self.size / short_edge
+            image = cv2.resize(image, dsize=None, fx=scale, fy=scale)
+            text_polys *= scale
+
+        gt_kernels = []
+        for i in range(1, self.kernel_num + 1):
+            # s1->sn, from big to small
+            rate = 1.0 - (1.0 - self.min_shrink_ratio) / (self.kernel_num - 1) * i
+            text_kernel, ignore_tags = self.generate_kernel(
+                image.shape[0:2], rate, text_polys, ignore_tags
+            )
+            gt_kernels.append(text_kernel)
+
+        training_mask = np.ones(image.shape[0:2], dtype="uint8")
+        for i in range(text_polys.shape[0]):
+            if ignore_tags[i]:
+                cv2.fillPoly(
+                    training_mask, text_polys[i].astype(np.int32)[np.newaxis, :, :], 0
+                )
+
+        gt_kernels = np.array(gt_kernels)
+        gt_kernels[gt_kernels > 0] = 1
+
+        data["image"] = image
+        data["polys"] = text_polys
+        data["gt_kernels"] = gt_kernels[0:]
+        data["gt_text"] = gt_kernels[0]
+        data["mask"] = training_mask.astype("float32")
+        return data
+
+    def generate_kernel(self, img_size, shrink_ratio, text_polys, ignore_tags=None):
+        """
+        Refer to part of the code:
+        https://github.com/open-mmlab/mmocr/blob/main/mmocr/datasets/pipelines/textdet_targets/base_textdet_targets.py
+        """
+
+        h, w = img_size
+        text_kernel = np.zeros((h, w), dtype=np.float32)
+        for i, poly in enumerate(text_polys):
+            polygon = Polygon(poly)
+            distance = (
+                polygon.area
+                * (1 - shrink_ratio * shrink_ratio)
+                / (polygon.length + 1e-6)
+            )
+            subject = [tuple(l) for l in poly]
+            pco = pyclipper.PyclipperOffset()
+            pco.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+            shrunk = np.array(pco.Execute(-distance))
+
+            if len(shrunk) == 0 or shrunk.size == 0:
+                if ignore_tags is not None:
+                    ignore_tags[i] = True
+                continue
+            try:
+                shrunk = np.array(shrunk[0]).reshape(-1, 2)
+            except:
+                if ignore_tags is not None:
+                    ignore_tags[i] = True
+                continue
+            cv2.fillPoly(text_kernel, [shrunk.astype(np.int32)], i + 1)
+        return text_kernel, ignore_tags

ppocr/data/imaug/make_shrink_map.py

@@ -0,0 +1,125 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/make_shrink_map.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import numpy as np
+import cv2
+from shapely.geometry import Polygon
+import pyclipper
+
+__all__ = ["MakeShrinkMap"]
+
+
+class MakeShrinkMap(object):
+    r"""
+    Making binary mask from detection data with ICDAR format.
+    Typically following the process of class `MakeICDARData`.
+    """
+
+    def __init__(self, min_text_size=8, shrink_ratio=0.4, **kwargs):
+        self.min_text_size = min_text_size
+        self.shrink_ratio = shrink_ratio
+        if "total_epoch" in kwargs and "epoch" in kwargs and kwargs["epoch"] != "None":
+            self.shrink_ratio = self.shrink_ratio + 0.2 * kwargs["epoch"] / float(
+                kwargs["total_epoch"]
+            )
+
+    def __call__(self, data):
+        image = data["image"]
+        text_polys = data["polys"]
+        ignore_tags = data["ignore_tags"]
+
+        h, w = image.shape[:2]
+        text_polys, ignore_tags = self.validate_polygons(text_polys, ignore_tags, h, w)
+        gt = np.zeros((h, w), dtype=np.float32)
+        mask = np.ones((h, w), dtype=np.float32)
+        for i in range(len(text_polys)):
+            polygon = text_polys[i]
+            height = max(polygon[:, 1]) - min(polygon[:, 1])
+            width = max(polygon[:, 0]) - min(polygon[:, 0])
+            if ignore_tags[i] or min(height, width) < self.min_text_size:
+                cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
+                ignore_tags[i] = True
+            else:
+                polygon_shape = Polygon(polygon)
+                subject = [tuple(l) for l in polygon]
+                padding = pyclipper.PyclipperOffset()
+                padding.AddPath(subject, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+                shrunk = []
+
+                # Increase the shrink ratio every time we get multiple polygon returned back
+                possible_ratios = np.arange(self.shrink_ratio, 1, self.shrink_ratio)
+                np.append(possible_ratios, 1)
+                # print(possible_ratios)
+                for ratio in possible_ratios:
+                    # print(f"Change shrink ratio to {ratio}")
+                    distance = (
+                        polygon_shape.area
+                        * (1 - np.power(ratio, 2))
+                        / polygon_shape.length
+                    )
+                    shrunk = padding.Execute(-distance)
+                    if len(shrunk) == 1:
+                        break
+
+                if shrunk == []:
+                    cv2.fillPoly(mask, polygon.astype(np.int32)[np.newaxis, :, :], 0)
+                    ignore_tags[i] = True
+                    continue
+
+                for each_shrink in shrunk:
+                    shrink = np.array(each_shrink).reshape(-1, 2)
+                    cv2.fillPoly(gt, [shrink.astype(np.int32)], 1)
+
+        data["shrink_map"] = gt
+        data["shrink_mask"] = mask
+        return data
+
+    def validate_polygons(self, polygons, ignore_tags, h, w):
+        """
+        polygons (numpy.array, required): of shape (num_instances, num_points, 2)
+        """
+        if len(polygons) == 0:
+            return polygons, ignore_tags
+        assert len(polygons) == len(ignore_tags)
+        for polygon in polygons:
+            polygon[:, 0] = np.clip(polygon[:, 0], 0, w - 1)
+            polygon[:, 1] = np.clip(polygon[:, 1], 0, h - 1)
+
+        for i in range(len(polygons)):
+            area = self.polygon_area(polygons[i])
+            if abs(area) < 1:
+                ignore_tags[i] = True
+            if area > 0:
+                polygons[i] = polygons[i][::-1, :]
+        return polygons, ignore_tags
+
+    def polygon_area(self, polygon):
+        """
+        compute polygon area
+        """
+        area = 0
+        q = polygon[-1]
+        for p in polygon:
+            area += p[0] * q[1] - p[1] * q[0]
+            q = p
+        return area / 2.0

ppocr/data/imaug/operators.py

@@ -0,0 +1,523 @@
+"""
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
+#
+# 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
+#
+#     http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import sys
+import cv2
+import numpy as np
+import math
+from PIL import Image
+
+
+class DecodeImage(object):
+    """decode image"""
+
+    def __init__(
+        self, img_mode="RGB", channel_first=False, ignore_orientation=False, **kwargs
+    ):
+        self.img_mode = img_mode
+        self.channel_first = channel_first
+        self.ignore_orientation = ignore_orientation
+
+    def __call__(self, data):
+        img = data["image"]
+        assert type(img) is bytes and len(img) > 0, "invalid input 'img' in DecodeImage"
+        img = np.frombuffer(img, dtype="uint8")
+        if self.ignore_orientation:
+            img = cv2.imdecode(img, cv2.IMREAD_IGNORE_ORIENTATION | cv2.IMREAD_COLOR)
+        else:
+            img = cv2.imdecode(img, 1)
+        if img is None:
+            return None
+        if self.img_mode == "GRAY":
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        elif self.img_mode == "RGB":
+            assert img.shape[2] == 3, "invalid shape of image[%s]" % (img.shape)
+            img = img[:, :, ::-1]
+
+        if self.channel_first:
+            img = img.transpose((2, 0, 1))
+
+        data["image"] = img
+        return data
+
+
+class NormalizeImage(object):
+    """normalize image such as subtract mean, divide std"""
+
+    def __init__(self, scale=None, mean=None, std=None, order="chw", **kwargs):
+        if isinstance(scale, str):
+            scale = eval(scale)
+        self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
+        mean = mean if mean is not None else [0.485, 0.456, 0.406]
+        std = std if std is not None else [0.229, 0.224, 0.225]
+
+        shape = (3, 1, 1) if order == "chw" else (1, 1, 3)
+        self.mean = np.array(mean).reshape(shape).astype("float32")
+        self.std = np.array(std).reshape(shape).astype("float32")
+
+    def __call__(self, data):
+        img = data["image"]
+        from PIL import Image
+
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        assert isinstance(img, np.ndarray), "invalid input 'img' in NormalizeImage"
+        data["image"] = (img.astype("float32") * self.scale - self.mean) / self.std
+        return data
+
+
+class ToCHWImage(object):
+    """convert hwc image to chw image"""
+
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        img = data["image"]
+        from PIL import Image
+
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        data["image"] = img.transpose((2, 0, 1))
+        return data
+
+
+class Fasttext(object):
+    def __init__(self, path="None", **kwargs):
+        import fasttext
+
+        self.fast_model = fasttext.load_model(path)
+
+    def __call__(self, data):
+        label = data["label"]
+        fast_label = self.fast_model[label]
+        data["fast_label"] = fast_label
+        return data
+
+
+class KeepKeys(object):
+    def __init__(self, keep_keys, **kwargs):
+        self.keep_keys = keep_keys
+
+    def __call__(self, data):
+        data_list = []
+        for key in self.keep_keys:
+            data_list.append(data[key])
+        return data_list
+
+
+class Pad(object):
+    def __init__(self, size=None, size_div=32, **kwargs):
+        if size is not None and not isinstance(size, (int, list, tuple)):
+            raise TypeError(
+                "Type of target_size is invalid. Now is {}".format(type(size))
+            )
+        if isinstance(size, int):
+            size = [size, size]
+        self.size = size
+        self.size_div = size_div
+
+    def __call__(self, data):
+        img = data["image"]
+        img_h, img_w = img.shape[0], img.shape[1]
+        if self.size:
+            resize_h2, resize_w2 = self.size
+            assert (
+                img_h < resize_h2 and img_w < resize_w2
+            ), "(h, w) of target size should be greater than (img_h, img_w)"
+        else:
+            resize_h2 = max(
+                int(math.ceil(img.shape[0] / self.size_div) * self.size_div),
+                self.size_div,
+            )
+            resize_w2 = max(
+                int(math.ceil(img.shape[1] / self.size_div) * self.size_div),
+                self.size_div,
+            )
+        img = cv2.copyMakeBorder(
+            img,
+            0,
+            resize_h2 - img_h,
+            0,
+            resize_w2 - img_w,
+            cv2.BORDER_CONSTANT,
+            value=0,
+        )
+        data["image"] = img
+        return data
+
+
+class Resize(object):
+    def __init__(self, size=(640, 640), **kwargs):
+        self.size = size
+
+    def resize_image(self, img):
+        resize_h, resize_w = self.size
+        ori_h, ori_w = img.shape[:2]  # (h, w, c)
+        ratio_h = float(resize_h) / ori_h
+        ratio_w = float(resize_w) / ori_w
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        return img, [ratio_h, ratio_w]
+
+    def __call__(self, data):
+        img = data["image"]
+        if "polys" in data:
+            text_polys = data["polys"]
+
+        img_resize, [ratio_h, ratio_w] = self.resize_image(img)
+        if "polys" in data:
+            new_boxes = []
+            for box in text_polys:
+                new_box = []
+                for cord in box:
+                    new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
+                new_boxes.append(new_box)
+            data["polys"] = np.array(new_boxes, dtype=np.float32)
+        data["image"] = img_resize
+        return data
+
+
+class DetResizeForTest(object):
+    def __init__(self, **kwargs):
+        super(DetResizeForTest, self).__init__()
+        self.resize_type = 0
+        self.keep_ratio = False
+        if "image_shape" in kwargs:
+            self.image_shape = kwargs["image_shape"]
+            self.resize_type = 1
+            if "keep_ratio" in kwargs:
+                self.keep_ratio = kwargs["keep_ratio"]
+        elif "limit_side_len" in kwargs:
+            self.limit_side_len = kwargs["limit_side_len"]
+            self.limit_type = kwargs.get("limit_type", "min")
+        elif "resize_long" in kwargs:
+            self.resize_type = 2
+            self.resize_long = kwargs.get("resize_long", 960)
+        else:
+            self.limit_side_len = 736
+            self.limit_type = "min"
+
+    def __call__(self, data):
+        img = data["image"]
+        src_h, src_w, _ = img.shape
+        if sum([src_h, src_w]) < 64:
+            img = self.image_padding(img)
+
+        if self.resize_type == 0:
+            # img, shape = self.resize_image_type0(img)
+            img, [ratio_h, ratio_w] = self.resize_image_type0(img)
+        elif self.resize_type == 2:
+            img, [ratio_h, ratio_w] = self.resize_image_type2(img)
+        else:
+            # img, shape = self.resize_image_type1(img)
+            img, [ratio_h, ratio_w] = self.resize_image_type1(img)
+        data["image"] = img
+        data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
+        return data
+
+    def image_padding(self, im, value=0):
+        h, w, c = im.shape
+        im_pad = np.zeros((max(32, h), max(32, w), c), np.uint8) + value
+        im_pad[:h, :w, :] = im
+        return im_pad
+
+    def resize_image_type1(self, img):
+        resize_h, resize_w = self.image_shape
+        ori_h, ori_w = img.shape[:2]  # (h, w, c)
+        if self.keep_ratio is True:
+            resize_w = ori_w * resize_h / ori_h
+            N = math.ceil(resize_w / 32)
+            resize_w = N * 32
+        ratio_h = float(resize_h) / ori_h
+        ratio_w = float(resize_w) / ori_w
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        # return img, np.array([ori_h, ori_w])
+        return img, [ratio_h, ratio_w]
+
+    def resize_image_type0(self, img):
+        """
+        resize image to a size multiple of 32 which is required by the network
+        args:
+            img(array): array with shape [h, w, c]
+        return(tuple):
+            img, (ratio_h, ratio_w)
+        """
+        limit_side_len = self.limit_side_len
+        h, w, c = img.shape
+
+        # limit the max side
+        if self.limit_type == "max":
+            if max(h, w) > limit_side_len:
+                if h > w:
+                    ratio = float(limit_side_len) / h
+                else:
+                    ratio = float(limit_side_len) / w
+            else:
+                ratio = 1.0
+        elif self.limit_type == "min":
+            if min(h, w) < limit_side_len:
+                if h < w:
+                    ratio = float(limit_side_len) / h
+                else:
+                    ratio = float(limit_side_len) / w
+            else:
+                ratio = 1.0
+        elif self.limit_type == "resize_long":
+            ratio = float(limit_side_len) / max(h, w)
+        else:
+            raise Exception("not support limit type, image ")
+        resize_h = int(h * ratio)
+        resize_w = int(w * ratio)
+
+        resize_h = max(int(round(resize_h / 32) * 32), 32)
+        resize_w = max(int(round(resize_w / 32) * 32), 32)
+
+        try:
+            if int(resize_w) <= 0 or int(resize_h) <= 0:
+                return None, (None, None)
+            img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        except:
+            print(img.shape, resize_w, resize_h)
+            sys.exit(0)
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+        return img, [ratio_h, ratio_w]
+
+    def resize_image_type2(self, img):
+        h, w, _ = img.shape
+
+        resize_w = w
+        resize_h = h
+
+        if resize_h > resize_w:
+            ratio = float(self.resize_long) / resize_h
+        else:
+            ratio = float(self.resize_long) / resize_w
+
+        resize_h = int(resize_h * ratio)
+        resize_w = int(resize_w * ratio)
+
+        max_stride = 128
+        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+
+        return img, [ratio_h, ratio_w]
+
+
+class E2EResizeForTest(object):
+    def __init__(self, **kwargs):
+        super(E2EResizeForTest, self).__init__()
+        self.max_side_len = kwargs["max_side_len"]
+        self.valid_set = kwargs["valid_set"]
+
+    def __call__(self, data):
+        img = data["image"]
+        src_h, src_w, _ = img.shape
+        if self.valid_set == "totaltext":
+            im_resized, [ratio_h, ratio_w] = self.resize_image_for_totaltext(
+                img, max_side_len=self.max_side_len
+            )
+        else:
+            im_resized, (ratio_h, ratio_w) = self.resize_image(
+                img, max_side_len=self.max_side_len
+            )
+        data["image"] = im_resized
+        data["shape"] = np.array([src_h, src_w, ratio_h, ratio_w])
+        return data
+
+    def resize_image_for_totaltext(self, im, max_side_len=512):
+        h, w, _ = im.shape
+        resize_w = w
+        resize_h = h
+        ratio = 1.25
+        if h * ratio > max_side_len:
+            ratio = float(max_side_len) / resize_h
+        resize_h = int(resize_h * ratio)
+        resize_w = int(resize_w * ratio)
+
+        max_stride = 128
+        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+        im = cv2.resize(im, (int(resize_w), int(resize_h)))
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+        return im, (ratio_h, ratio_w)
+
+    def resize_image(self, im, max_side_len=512):
+        """
+        resize image to a size multiple of max_stride which is required by the network
+        :param im: the resized image
+        :param max_side_len: limit of max image size to avoid out of memory in gpu
+        :return: the resized image and the resize ratio
+        """
+        h, w, _ = im.shape
+
+        resize_w = w
+        resize_h = h
+
+        # Fix the longer side
+        if resize_h > resize_w:
+            ratio = float(max_side_len) / resize_h
+        else:
+            ratio = float(max_side_len) / resize_w
+
+        resize_h = int(resize_h * ratio)
+        resize_w = int(resize_w * ratio)
+
+        max_stride = 128
+        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+        im = cv2.resize(im, (int(resize_w), int(resize_h)))
+        ratio_h = resize_h / float(h)
+        ratio_w = resize_w / float(w)
+
+        return im, (ratio_h, ratio_w)
+
+
+class KieResize(object):
+    def __init__(self, **kwargs):
+        super(KieResize, self).__init__()
+        self.max_side, self.min_side = kwargs["img_scale"][0], kwargs["img_scale"][1]
+
+    def __call__(self, data):
+        img = data["image"]
+        points = data["points"]
+        src_h, src_w, _ = img.shape
+        (
+            im_resized,
+            scale_factor,
+            [ratio_h, ratio_w],
+            [new_h, new_w],
+        ) = self.resize_image(img)
+        resize_points = self.resize_boxes(img, points, scale_factor)
+        data["ori_image"] = img
+        data["ori_boxes"] = points
+        data["points"] = resize_points
+        data["image"] = im_resized
+        data["shape"] = np.array([new_h, new_w])
+        return data
+
+    def resize_image(self, img):
+        norm_img = np.zeros([1024, 1024, 3], dtype="float32")
+        scale = [512, 1024]
+        h, w = img.shape[:2]
+        max_long_edge = max(scale)
+        max_short_edge = min(scale)
+        scale_factor = min(max_long_edge / max(h, w), max_short_edge / min(h, w))
+        resize_w, resize_h = int(w * float(scale_factor) + 0.5), int(
+            h * float(scale_factor) + 0.5
+        )
+        max_stride = 32
+        resize_h = (resize_h + max_stride - 1) // max_stride * max_stride
+        resize_w = (resize_w + max_stride - 1) // max_stride * max_stride
+        im = cv2.resize(img, (resize_w, resize_h))
+        new_h, new_w = im.shape[:2]
+        w_scale = new_w / w
+        h_scale = new_h / h
+        scale_factor = np.array([w_scale, h_scale, w_scale, h_scale], dtype=np.float32)
+        norm_img[:new_h, :new_w, :] = im
+        return norm_img, scale_factor, [h_scale, w_scale], [new_h, new_w]
+
+    def resize_boxes(self, im, points, scale_factor):
+        points = points * scale_factor
+        img_shape = im.shape[:2]
+        points[:, 0::2] = np.clip(points[:, 0::2], 0, img_shape[1])
+        points[:, 1::2] = np.clip(points[:, 1::2], 0, img_shape[0])
+        return points
+
+
+class SRResize(object):
+    def __init__(
+        self,
+        imgH=32,
+        imgW=128,
+        down_sample_scale=4,
+        keep_ratio=False,
+        min_ratio=1,
+        mask=False,
+        infer_mode=False,
+        **kwargs,
+    ):
+        self.imgH = imgH
+        self.imgW = imgW
+        self.keep_ratio = keep_ratio
+        self.min_ratio = min_ratio
+        self.down_sample_scale = down_sample_scale
+        self.mask = mask
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        imgH = self.imgH
+        imgW = self.imgW
+        images_lr = data["image_lr"]
+        transform2 = ResizeNormalize(
+            (imgW // self.down_sample_scale, imgH // self.down_sample_scale)
+        )
+        images_lr = transform2(images_lr)
+        data["img_lr"] = images_lr
+        if self.infer_mode:
+            return data
+
+        images_HR = data["image_hr"]
+        label_strs = data["label"]
+        transform = ResizeNormalize((imgW, imgH))
+        images_HR = transform(images_HR)
+        data["img_hr"] = images_HR
+        return data
+
+
+class ResizeNormalize(object):
+    def __init__(self, size, interpolation=Image.BICUBIC):
+        self.size = size
+        self.interpolation = interpolation
+
+    def __call__(self, img):
+        img = img.resize(self.size, self.interpolation)
+        img_numpy = np.array(img).astype("float32")
+        img_numpy = img_numpy.transpose((2, 0, 1)) / 255
+        return img_numpy
+
+
+class GrayImageChannelFormat(object):
+    """
+    format gray scale image's channel: (3,h,w) -> (1,h,w)
+    Args:
+        inverse: inverse gray image
+    """
+
+    def __init__(self, inverse=False, **kwargs):
+        self.inverse = inverse
+
+    def __call__(self, data):
+        img = data["image"]
+        img_single_channel = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        img_expanded = np.expand_dims(img_single_channel, 0)
+
+        if self.inverse:
+            data["image"] = np.abs(img_expanded - 1)
+        else:
+            data["image"] = img_expanded
+
+        data["src_image"] = img
+        return data

ppocr/data/imaug/randaugment.py

@@ -0,0 +1,134 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from PIL import Image, ImageEnhance, ImageOps
+import numpy as np
+import random
+
+
+class RawRandAugment(object):
+    def __init__(self, num_layers=2, magnitude=5, fillcolor=(128, 128, 128), **kwargs):
+        self.num_layers = num_layers
+        self.magnitude = magnitude
+        self.max_level = 10
+
+        abso_level = self.magnitude / self.max_level
+        self.level_map = {
+            "shearX": 0.3 * abso_level,
+            "shearY": 0.3 * abso_level,
+            "translateX": 150.0 / 331 * abso_level,
+            "translateY": 150.0 / 331 * abso_level,
+            "rotate": 30 * abso_level,
+            "color": 0.9 * abso_level,
+            "posterize": int(4.0 * abso_level),
+            "solarize": 256.0 * abso_level,
+            "contrast": 0.9 * abso_level,
+            "sharpness": 0.9 * abso_level,
+            "brightness": 0.9 * abso_level,
+            "autocontrast": 0,
+            "equalize": 0,
+            "invert": 0,
+        }
+
+        # from https://stackoverflow.com/questions/5252170/
+        # specify-image-filling-color-when-rotating-in-python-with-pil-and-setting-expand
+        def rotate_with_fill(img, magnitude):
+            rot = img.convert("RGBA").rotate(magnitude)
+            return Image.composite(
+                rot, Image.new("RGBA", rot.size, (128,) * 4), rot
+            ).convert(img.mode)
+
+        rnd_ch_op = random.choice
+
+        self.func = {
+            "shearX": lambda img, magnitude: img.transform(
+                img.size,
+                Image.AFFINE,
+                (1, magnitude * rnd_ch_op([-1, 1]), 0, 0, 1, 0),
+                Image.BICUBIC,
+                fillcolor=fillcolor,
+            ),
+            "shearY": lambda img, magnitude: img.transform(
+                img.size,
+                Image.AFFINE,
+                (1, 0, 0, magnitude * rnd_ch_op([-1, 1]), 1, 0),
+                Image.BICUBIC,
+                fillcolor=fillcolor,
+            ),
+            "translateX": lambda img, magnitude: img.transform(
+                img.size,
+                Image.AFFINE,
+                (1, 0, magnitude * img.size[0] * rnd_ch_op([-1, 1]), 0, 1, 0),
+                fillcolor=fillcolor,
+            ),
+            "translateY": lambda img, magnitude: img.transform(
+                img.size,
+                Image.AFFINE,
+                (1, 0, 0, 0, 1, magnitude * img.size[1] * rnd_ch_op([-1, 1])),
+                fillcolor=fillcolor,
+            ),
+            "rotate": lambda img, magnitude: rotate_with_fill(img, magnitude),
+            "color": lambda img, magnitude: ImageEnhance.Color(img).enhance(
+                1 + magnitude * rnd_ch_op([-1, 1])
+            ),
+            "posterize": lambda img, magnitude: ImageOps.posterize(img, magnitude),
+            "solarize": lambda img, magnitude: ImageOps.solarize(img, magnitude),
+            "contrast": lambda img, magnitude: ImageEnhance.Contrast(img).enhance(
+                1 + magnitude * rnd_ch_op([-1, 1])
+            ),
+            "sharpness": lambda img, magnitude: ImageEnhance.Sharpness(img).enhance(
+                1 + magnitude * rnd_ch_op([-1, 1])
+            ),
+            "brightness": lambda img, magnitude: ImageEnhance.Brightness(img).enhance(
+                1 + magnitude * rnd_ch_op([-1, 1])
+            ),
+            "autocontrast": lambda img, magnitude: ImageOps.autocontrast(img),
+            "equalize": lambda img, magnitude: ImageOps.equalize(img),
+            "invert": lambda img, magnitude: ImageOps.invert(img),
+        }
+
+    def __call__(self, img):
+        avaiable_op_names = list(self.level_map.keys())
+        for layer_num in range(self.num_layers):
+            op_name = np.random.choice(avaiable_op_names)
+            img = self.func[op_name](img, self.level_map[op_name])
+        return img
+
+
+class RandAugment(RawRandAugment):
+    """RandAugment wrapper to auto fit different img types"""
+
+    def __init__(self, prob=0.5, *args, **kwargs):
+        self.prob = prob
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, data):
+        if np.random.rand() > self.prob:
+            return data
+        img = data["image"]
+        if not isinstance(img, Image.Image):
+            img = np.ascontiguousarray(img)
+            img = Image.fromarray(img)
+
+        img = super().__call__(img)
+
+        if isinstance(img, Image.Image):
+            img = np.asarray(img)
+        data["image"] = img
+        return data

ppocr/data/imaug/random_crop_data.py

@@ -0,0 +1,238 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/data_loader/modules/random_crop_data.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import numpy as np
+import cv2
+import random
+
+
+def is_poly_in_rect(poly, x, y, w, h):
+    poly = np.array(poly)
+    if poly[:, 0].min() < x or poly[:, 0].max() > x + w:
+        return False
+    if poly[:, 1].min() < y or poly[:, 1].max() > y + h:
+        return False
+    return True
+
+
+def is_poly_outside_rect(poly, x, y, w, h):
+    poly = np.array(poly)
+    if poly[:, 0].max() < x or poly[:, 0].min() > x + w:
+        return True
+    if poly[:, 1].max() < y or poly[:, 1].min() > y + h:
+        return True
+    return False
+
+
+def split_regions(axis):
+    regions = []
+    min_axis = 0
+    for i in range(1, axis.shape[0]):
+        if axis[i] != axis[i - 1] + 1:
+            region = axis[min_axis:i]
+            min_axis = i
+            regions.append(region)
+    return regions
+
+
+def random_select(axis, max_size):
+    xx = np.random.choice(axis, size=2)
+    xmin = np.min(xx)
+    xmax = np.max(xx)
+    xmin = np.clip(xmin, 0, max_size - 1)
+    xmax = np.clip(xmax, 0, max_size - 1)
+    return xmin, xmax
+
+
+def region_wise_random_select(regions, max_size):
+    selected_index = list(np.random.choice(len(regions), 2))
+    selected_values = []
+    for index in selected_index:
+        axis = regions[index]
+        xx = int(np.random.choice(axis, size=1))
+        selected_values.append(xx)
+    xmin = min(selected_values)
+    xmax = max(selected_values)
+    return xmin, xmax
+
+
+def crop_area(im, text_polys, min_crop_side_ratio, max_tries):
+    h, w, _ = im.shape
+    h_array = np.zeros(h, dtype=np.int32)
+    w_array = np.zeros(w, dtype=np.int32)
+    for points in text_polys:
+        points = np.round(points, decimals=0).astype(np.int32)
+        minx = np.min(points[:, 0])
+        maxx = np.max(points[:, 0])
+        w_array[minx:maxx] = 1
+        miny = np.min(points[:, 1])
+        maxy = np.max(points[:, 1])
+        h_array[miny:maxy] = 1
+    # ensure the cropped area not across a text
+    h_axis = np.where(h_array == 0)[0]
+    w_axis = np.where(w_array == 0)[0]
+
+    if len(h_axis) == 0 or len(w_axis) == 0:
+        return 0, 0, w, h
+
+    h_regions = split_regions(h_axis)
+    w_regions = split_regions(w_axis)
+
+    for i in range(max_tries):
+        if len(w_regions) > 1:
+            xmin, xmax = region_wise_random_select(w_regions, w)
+        else:
+            xmin, xmax = random_select(w_axis, w)
+        if len(h_regions) > 1:
+            ymin, ymax = region_wise_random_select(h_regions, h)
+        else:
+            ymin, ymax = random_select(h_axis, h)
+
+        if (
+            xmax - xmin < min_crop_side_ratio * w
+            or ymax - ymin < min_crop_side_ratio * h
+        ):
+            # area too small
+            continue
+        num_poly_in_rect = 0
+        for poly in text_polys:
+            if not is_poly_outside_rect(poly, xmin, ymin, xmax - xmin, ymax - ymin):
+                num_poly_in_rect += 1
+                break
+
+        if num_poly_in_rect > 0:
+            return xmin, ymin, xmax - xmin, ymax - ymin
+
+    return 0, 0, w, h
+
+
+class EastRandomCropData(object):
+    def __init__(
+        self,
+        size=(640, 640),
+        max_tries=10,
+        min_crop_side_ratio=0.1,
+        keep_ratio=True,
+        **kwargs,
+    ):
+        self.size = size
+        self.max_tries = max_tries
+        self.min_crop_side_ratio = min_crop_side_ratio
+        self.keep_ratio = keep_ratio
+
+    def __call__(self, data):
+        img = data["image"]
+        text_polys = data["polys"]
+        ignore_tags = data["ignore_tags"]
+        texts = data["texts"]
+        all_care_polys = [text_polys[i] for i, tag in enumerate(ignore_tags) if not tag]
+        # 计算crop区域
+        crop_x, crop_y, crop_w, crop_h = crop_area(
+            img, all_care_polys, self.min_crop_side_ratio, self.max_tries
+        )
+        # crop 图片 保持比例填充
+        scale_w = self.size[0] / crop_w
+        scale_h = self.size[1] / crop_h
+        scale = min(scale_w, scale_h)
+        h = int(crop_h * scale)
+        w = int(crop_w * scale)
+        if self.keep_ratio:
+            padimg = np.zeros((self.size[1], self.size[0], img.shape[2]), img.dtype)
+            padimg[:h, :w] = cv2.resize(
+                img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w], (w, h)
+            )
+            img = padimg
+        else:
+            img = cv2.resize(
+                img[crop_y : crop_y + crop_h, crop_x : crop_x + crop_w],
+                tuple(self.size),
+            )
+        # crop 文本框
+        text_polys_crop = []
+        ignore_tags_crop = []
+        texts_crop = []
+        for poly, text, tag in zip(text_polys, texts, ignore_tags):
+            poly = ((poly - (crop_x, crop_y)) * scale).tolist()
+            if not is_poly_outside_rect(poly, 0, 0, w, h):
+                text_polys_crop.append(poly)
+                ignore_tags_crop.append(tag)
+                texts_crop.append(text)
+        data["image"] = img
+        data["polys"] = np.array(text_polys_crop)
+        data["ignore_tags"] = ignore_tags_crop
+        data["texts"] = texts_crop
+        return data
+
+
+class RandomCropImgMask(object):
+    def __init__(self, size, main_key, crop_keys, p=3 / 8, **kwargs):
+        self.size = size
+        self.main_key = main_key
+        self.crop_keys = crop_keys
+        self.p = p
+
+    def __call__(self, data):
+        image = data["image"]
+
+        h, w = image.shape[0:2]
+        th, tw = self.size
+        if w == tw and h == th:
+            return data
+
+        mask = data[self.main_key]
+        if np.max(mask) > 0 and random.random() > self.p:
+            # make sure to crop the text region
+            tl = np.min(np.where(mask > 0), axis=1) - (th, tw)
+            tl[tl < 0] = 0
+            br = np.max(np.where(mask > 0), axis=1) - (th, tw)
+            br[br < 0] = 0
+
+            br[0] = min(br[0], h - th)
+            br[1] = min(br[1], w - tw)
+
+            i = random.randint(tl[0], br[0]) if tl[0] < br[0] else 0
+            j = random.randint(tl[1], br[1]) if tl[1] < br[1] else 0
+        else:
+            i = random.randint(0, h - th) if h - th > 0 else 0
+            j = random.randint(0, w - tw) if w - tw > 0 else 0
+
+        # return i, j, th, tw
+        for k in data:
+            if k in self.crop_keys:
+                if len(data[k].shape) == 3:
+                    if np.argmin(data[k].shape) == 0:
+                        img = data[k][:, i : i + th, j : j + tw]
+                        if img.shape[1] != img.shape[2]:
+                            a = 1
+                    elif np.argmin(data[k].shape) == 2:
+                        img = data[k][i : i + th, j : j + tw, :]
+                        if img.shape[1] != img.shape[0]:
+                            a = 1
+                    else:
+                        img = data[k]
+                else:
+                    img = data[k][i : i + th, j : j + tw]
+                    if img.shape[0] != img.shape[1]:
+                        a = 1
+                data[k] = img
+        return data

ppocr/data/imaug/rec_img_aug.py

@@ -0,0 +1,932 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import math
+import cv2
+import numpy as np
+import random
+import copy
+from PIL import Image
+import PIL
+from .text_image_aug import tia_perspective, tia_stretch, tia_distort
+from .abinet_aug import (
+    CVGeometry,
+    CVDeterioration,
+    CVColorJitter,
+    SVTRGeometry,
+    SVTRDeterioration,
+    ParseQDeterioration,
+)
+from paddle.vision.transforms import Compose
+
+
+class RecAug(object):
+    def __init__(
+        self,
+        tia_prob=0.4,
+        crop_prob=0.4,
+        reverse_prob=0.4,
+        noise_prob=0.4,
+        jitter_prob=0.4,
+        blur_prob=0.4,
+        hsv_aug_prob=0.4,
+        **kwargs,
+    ):
+        self.tia_prob = tia_prob
+        self.bda = BaseDataAugmentation(
+            crop_prob, reverse_prob, noise_prob, jitter_prob, blur_prob, hsv_aug_prob
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        h, w, _ = img.shape
+
+        # tia
+        if random.random() <= self.tia_prob:
+            if h >= 20 and w >= 20:
+                img = tia_distort(img, random.randint(3, 6))
+                img = tia_stretch(img, random.randint(3, 6))
+                img = tia_perspective(img)
+
+        # bda
+        data["image"] = img
+        data = self.bda(data)
+        return data
+
+
+class BaseDataAugmentation(object):
+    def __init__(
+        self,
+        crop_prob=0.4,
+        reverse_prob=0.4,
+        noise_prob=0.4,
+        jitter_prob=0.4,
+        blur_prob=0.4,
+        hsv_aug_prob=0.4,
+        **kwargs,
+    ):
+        self.crop_prob = crop_prob
+        self.reverse_prob = reverse_prob
+        self.noise_prob = noise_prob
+        self.jitter_prob = jitter_prob
+        self.blur_prob = blur_prob
+        self.hsv_aug_prob = hsv_aug_prob
+        # for GaussianBlur
+        self.fil = cv2.getGaussianKernel(ksize=5, sigma=1, ktype=cv2.CV_32F)
+
+    def __call__(self, data):
+        img = data["image"]
+        h, w, _ = img.shape
+
+        if random.random() <= self.crop_prob and h >= 20 and w >= 20:
+            img = get_crop(img)
+
+        if random.random() <= self.blur_prob:
+            # GaussianBlur
+            img = cv2.sepFilter2D(img, -1, self.fil, self.fil)
+
+        if random.random() <= self.hsv_aug_prob:
+            img = hsv_aug(img)
+
+        if random.random() <= self.jitter_prob:
+            img = jitter(img)
+
+        if random.random() <= self.noise_prob:
+            img = add_gasuss_noise(img)
+
+        if random.random() <= self.reverse_prob:
+            img = 255 - img
+
+        data["image"] = img
+        return data
+
+
+class ABINetRecAug(object):
+    def __init__(
+        self, geometry_p=0.5, deterioration_p=0.25, colorjitter_p=0.25, **kwargs
+    ):
+        self.transforms = Compose(
+            [
+                CVGeometry(
+                    degrees=45,
+                    translate=(0.0, 0.0),
+                    scale=(0.5, 2.0),
+                    shear=(45, 15),
+                    distortion=0.5,
+                    p=geometry_p,
+                ),
+                CVDeterioration(var=20, degrees=6, factor=4, p=deterioration_p),
+                CVColorJitter(
+                    brightness=0.5,
+                    contrast=0.5,
+                    saturation=0.5,
+                    hue=0.1,
+                    p=colorjitter_p,
+                ),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        img = self.transforms(img)
+        data["image"] = img
+        return data
+
+
+class RecConAug(object):
+    def __init__(
+        self,
+        prob=0.5,
+        image_shape=(32, 320, 3),
+        max_text_length=25,
+        ext_data_num=1,
+        **kwargs,
+    ):
+        self.ext_data_num = ext_data_num
+        self.prob = prob
+        self.max_text_length = max_text_length
+        self.image_shape = image_shape
+        self.max_wh_ratio = self.image_shape[1] / self.image_shape[0]
+
+    def merge_ext_data(self, data, ext_data):
+        ori_w = round(
+            data["image"].shape[1] / data["image"].shape[0] * self.image_shape[0]
+        )
+        ext_w = round(
+            ext_data["image"].shape[1]
+            / ext_data["image"].shape[0]
+            * self.image_shape[0]
+        )
+        data["image"] = cv2.resize(data["image"], (ori_w, self.image_shape[0]))
+        ext_data["image"] = cv2.resize(ext_data["image"], (ext_w, self.image_shape[0]))
+        data["image"] = np.concatenate([data["image"], ext_data["image"]], axis=1)
+        data["label"] += ext_data["label"]
+        return data
+
+    def __call__(self, data):
+        rnd_num = random.random()
+        if rnd_num > self.prob:
+            return data
+        for idx, ext_data in enumerate(data["ext_data"]):
+            if len(data["label"]) + len(ext_data["label"]) > self.max_text_length:
+                break
+            concat_ratio = (
+                data["image"].shape[1] / data["image"].shape[0]
+                + ext_data["image"].shape[1] / ext_data["image"].shape[0]
+            )
+            if concat_ratio > self.max_wh_ratio:
+                break
+            data = self.merge_ext_data(data, ext_data)
+        data.pop("ext_data")
+        return data
+
+
+class SVTRRecAug(object):
+    def __init__(
+        self,
+        aug_type=0,
+        geometry_p=0.5,
+        deterioration_p=0.25,
+        colorjitter_p=0.25,
+        **kwargs,
+    ):
+        self.transforms = Compose(
+            [
+                SVTRGeometry(
+                    aug_type=aug_type,
+                    degrees=45,
+                    translate=(0.0, 0.0),
+                    scale=(0.5, 2.0),
+                    shear=(45, 15),
+                    distortion=0.5,
+                    p=geometry_p,
+                ),
+                SVTRDeterioration(var=20, degrees=6, factor=4, p=deterioration_p),
+                CVColorJitter(
+                    brightness=0.5,
+                    contrast=0.5,
+                    saturation=0.5,
+                    hue=0.1,
+                    p=colorjitter_p,
+                ),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        img = self.transforms(img)
+        data["image"] = img
+        return data
+
+
+class ParseQRecAug(object):
+    def __init__(
+        self,
+        aug_type=0,
+        geometry_p=0.5,
+        deterioration_p=0.25,
+        colorjitter_p=0.25,
+        **kwargs,
+    ):
+        self.transforms = Compose(
+            [
+                SVTRGeometry(
+                    aug_type=aug_type,
+                    degrees=45,
+                    translate=(0.0, 0.0),
+                    scale=(0.5, 2.0),
+                    shear=(45, 15),
+                    distortion=0.5,
+                    p=geometry_p,
+                ),
+                ParseQDeterioration(
+                    var=20, degrees=6, lam=20, radius=2.0, factor=4, p=deterioration_p
+                ),
+                CVColorJitter(
+                    brightness=0.5,
+                    contrast=0.5,
+                    saturation=0.5,
+                    hue=0.1,
+                    p=colorjitter_p,
+                ),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        img = self.transforms(img)
+        data["image"] = img
+        return data
+
+
+class ClsResizeImg(object):
+    def __init__(self, image_shape, **kwargs):
+        self.image_shape = image_shape
+
+    def __call__(self, data):
+        img = data["image"]
+        norm_img, _ = resize_norm_img(img, self.image_shape)
+        data["image"] = norm_img
+        return data
+
+
+class RecResizeImg(object):
+    def __init__(
+        self,
+        image_shape,
+        infer_mode=False,
+        eval_mode=False,
+        character_dict_path="./ppocr/utils/ppocr_keys_v1.txt",
+        padding=True,
+        **kwargs,
+    ):
+        self.image_shape = image_shape
+        self.infer_mode = infer_mode
+        self.eval_mode = eval_mode
+        self.character_dict_path = character_dict_path
+        self.padding = padding
+
+    def __call__(self, data):
+        img = data["image"]
+        if self.eval_mode or (self.infer_mode and self.character_dict_path is not None):
+            norm_img, valid_ratio = resize_norm_img_chinese(img, self.image_shape)
+        else:
+            norm_img, valid_ratio = resize_norm_img(img, self.image_shape, self.padding)
+        data["image"] = norm_img
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class VLRecResizeImg(object):
+    def __init__(
+        self,
+        image_shape,
+        infer_mode=False,
+        character_dict_path="./ppocr/utils/ppocr_keys_v1.txt",
+        padding=True,
+        **kwargs,
+    ):
+        self.image_shape = image_shape
+        self.infer_mode = infer_mode
+        self.character_dict_path = character_dict_path
+        self.padding = padding
+
+    def __call__(self, data):
+        img = data["image"]
+
+        imgC, imgH, imgW = self.image_shape
+        resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
+        resized_w = imgW
+        resized_image = resized_image.astype("float32")
+        if self.image_shape[0] == 1:
+            resized_image = resized_image / 255
+            norm_img = resized_image[np.newaxis, :]
+        else:
+            norm_img = resized_image.transpose((2, 0, 1)) / 255
+        valid_ratio = min(1.0, float(resized_w / imgW))
+
+        data["image"] = norm_img
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class RFLRecResizeImg(object):
+    def __init__(self, image_shape, padding=True, interpolation=1, **kwargs):
+        self.image_shape = image_shape
+        self.padding = padding
+
+        self.interpolation = interpolation
+        if self.interpolation == 0:
+            self.interpolation = cv2.INTER_NEAREST
+        elif self.interpolation == 1:
+            self.interpolation = cv2.INTER_LINEAR
+        elif self.interpolation == 2:
+            self.interpolation = cv2.INTER_CUBIC
+        elif self.interpolation == 3:
+            self.interpolation = cv2.INTER_AREA
+        else:
+            raise Exception("Unsupported interpolation type !!!")
+
+    def __call__(self, data):
+        img = data["image"]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        norm_img, valid_ratio = resize_norm_img(
+            img, self.image_shape, self.padding, self.interpolation
+        )
+        data["image"] = norm_img
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class SRNRecResizeImg(object):
+    def __init__(self, image_shape, num_heads, max_text_length, **kwargs):
+        self.image_shape = image_shape
+        self.num_heads = num_heads
+        self.max_text_length = max_text_length
+
+    def __call__(self, data):
+        img = data["image"]
+        norm_img = resize_norm_img_srn(img, self.image_shape)
+        data["image"] = norm_img
+        [
+            encoder_word_pos,
+            gsrm_word_pos,
+            gsrm_slf_attn_bias1,
+            gsrm_slf_attn_bias2,
+        ] = srn_other_inputs(self.image_shape, self.num_heads, self.max_text_length)
+
+        data["encoder_word_pos"] = encoder_word_pos
+        data["gsrm_word_pos"] = gsrm_word_pos
+        data["gsrm_slf_attn_bias1"] = gsrm_slf_attn_bias1
+        data["gsrm_slf_attn_bias2"] = gsrm_slf_attn_bias2
+        return data
+
+
+class SARRecResizeImg(object):
+    def __init__(self, image_shape, width_downsample_ratio=0.25, **kwargs):
+        self.image_shape = image_shape
+        self.width_downsample_ratio = width_downsample_ratio
+
+    def __call__(self, data):
+        img = data["image"]
+        norm_img, resize_shape, pad_shape, valid_ratio = resize_norm_img_sar(
+            img, self.image_shape, self.width_downsample_ratio
+        )
+        data["image"] = norm_img
+        data["resized_shape"] = resize_shape
+        data["pad_shape"] = pad_shape
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class PRENResizeImg(object):
+    def __init__(self, image_shape, **kwargs):
+        """
+        According to original paper's realization, it's a hard resize method here.
+        So maybe you should optimize it to fit for your task better.
+        """
+        self.dst_h, self.dst_w = image_shape
+
+    def __call__(self, data):
+        img = data["image"]
+        resized_img = cv2.resize(
+            img, (self.dst_w, self.dst_h), interpolation=cv2.INTER_LINEAR
+        )
+        resized_img = resized_img.transpose((2, 0, 1)) / 255
+        resized_img -= 0.5
+        resized_img /= 0.5
+        data["image"] = resized_img.astype(np.float32)
+        return data
+
+
+class SPINRecResizeImg(object):
+    def __init__(
+        self,
+        image_shape,
+        interpolation=2,
+        mean=(127.5, 127.5, 127.5),
+        std=(127.5, 127.5, 127.5),
+        **kwargs,
+    ):
+        self.image_shape = image_shape
+
+        self.mean = np.array(mean, dtype=np.float32)
+        self.std = np.array(std, dtype=np.float32)
+        self.interpolation = interpolation
+
+    def __call__(self, data):
+        img = data["image"]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        # different interpolation type corresponding the OpenCV
+        if self.interpolation == 0:
+            interpolation = cv2.INTER_NEAREST
+        elif self.interpolation == 1:
+            interpolation = cv2.INTER_LINEAR
+        elif self.interpolation == 2:
+            interpolation = cv2.INTER_CUBIC
+        elif self.interpolation == 3:
+            interpolation = cv2.INTER_AREA
+        else:
+            raise Exception("Unsupported interpolation type !!!")
+        # Deal with the image error during image loading
+        if img is None:
+            return None
+
+        img = cv2.resize(img, tuple(self.image_shape), interpolation)
+        img = np.array(img, np.float32)
+        img = np.expand_dims(img, -1)
+        img = img.transpose((2, 0, 1))
+        # normalize the image
+        img = img.copy().astype(np.float32)
+        mean = np.float64(self.mean.reshape(1, -1))
+        stdinv = 1 / np.float64(self.std.reshape(1, -1))
+        img -= mean
+        img *= stdinv
+        data["image"] = img
+        return data
+
+
+class GrayRecResizeImg(object):
+    def __init__(
+        self,
+        image_shape,
+        resize_type,
+        inter_type="Image.Resampling.LANCZOS",
+        scale=True,
+        padding=False,
+        **kwargs,
+    ):
+        self.image_shape = image_shape
+        self.resize_type = resize_type
+        self.padding = padding
+        self.inter_type = eval(inter_type)
+        self.scale = scale
+
+    def __call__(self, data):
+        img = data["image"]
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        image_shape = self.image_shape
+        if self.padding:
+            imgC, imgH, imgW = image_shape
+            # todo: change to 0 and modified image shape
+            h = img.shape[0]
+            w = img.shape[1]
+            ratio = w / float(h)
+            if math.ceil(imgH * ratio) > imgW:
+                resized_w = imgW
+            else:
+                resized_w = int(math.ceil(imgH * ratio))
+            resized_image = cv2.resize(img, (resized_w, imgH))
+            norm_img = np.expand_dims(resized_image, -1)
+            norm_img = norm_img.transpose((2, 0, 1))
+            resized_image = norm_img.astype(np.float32) / 128.0 - 1.0
+            padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
+            padding_im[:, :, 0:resized_w] = resized_image
+            data["image"] = padding_im
+            return data
+        if self.resize_type == "PIL":
+            image_pil = Image.fromarray(np.uint8(img))
+            img = image_pil.resize(self.image_shape, self.inter_type)
+            img = np.array(img)
+        if self.resize_type == "OpenCV":
+            img = cv2.resize(img, self.image_shape)
+        norm_img = np.expand_dims(img, -1)
+        norm_img = norm_img.transpose((2, 0, 1))
+        if self.scale:
+            data["image"] = norm_img.astype(np.float32) / 128.0 - 1.0
+        else:
+            data["image"] = norm_img.astype(np.float32) / 255.0
+        return data
+
+
+class ABINetRecResizeImg(object):
+    def __init__(self, image_shape, **kwargs):
+        self.image_shape = image_shape
+
+    def __call__(self, data):
+        img = data["image"]
+        norm_img, valid_ratio = resize_norm_img_abinet(img, self.image_shape)
+        data["image"] = norm_img
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class SVTRRecResizeImg(object):
+    def __init__(self, image_shape, padding=True, **kwargs):
+        self.image_shape = image_shape
+        self.padding = padding
+
+    def __call__(self, data):
+        img = data["image"]
+
+        norm_img, valid_ratio = resize_norm_img(img, self.image_shape, self.padding)
+        data["image"] = norm_img
+        data["valid_ratio"] = valid_ratio
+        return data
+
+
+class RobustScannerRecResizeImg(object):
+    def __init__(
+        self, image_shape, max_text_length, width_downsample_ratio=0.25, **kwargs
+    ):
+        self.image_shape = image_shape
+        self.width_downsample_ratio = width_downsample_ratio
+        self.max_text_length = max_text_length
+
+    def __call__(self, data):
+        img = data["image"]
+        norm_img, resize_shape, pad_shape, valid_ratio = resize_norm_img_sar(
+            img, self.image_shape, self.width_downsample_ratio
+        )
+        word_positons = np.array(range(0, self.max_text_length)).astype("int64")
+        data["image"] = norm_img
+        data["resized_shape"] = resize_shape
+        data["pad_shape"] = pad_shape
+        data["valid_ratio"] = valid_ratio
+        data["word_positons"] = word_positons
+        return data
+
+
+def resize_norm_img_sar(img, image_shape, width_downsample_ratio=0.25):
+    imgC, imgH, imgW_min, imgW_max = image_shape
+    h = img.shape[0]
+    w = img.shape[1]
+    valid_ratio = 1.0
+    # make sure new_width is an integral multiple of width_divisor.
+    width_divisor = int(1 / width_downsample_ratio)
+    # resize
+    ratio = w / float(h)
+    resize_w = math.ceil(imgH * ratio)
+    if resize_w % width_divisor != 0:
+        resize_w = round(resize_w / width_divisor) * width_divisor
+    if imgW_min is not None:
+        resize_w = max(imgW_min, resize_w)
+    if imgW_max is not None:
+        valid_ratio = min(1.0, 1.0 * resize_w / imgW_max)
+        resize_w = min(imgW_max, resize_w)
+    resized_image = cv2.resize(img, (resize_w, imgH))
+    resized_image = resized_image.astype("float32")
+    # norm
+    if image_shape[0] == 1:
+        resized_image = resized_image / 255
+        resized_image = resized_image[np.newaxis, :]
+    else:
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+    resized_image -= 0.5
+    resized_image /= 0.5
+    resize_shape = resized_image.shape
+    padding_im = -1.0 * np.ones((imgC, imgH, imgW_max), dtype=np.float32)
+    padding_im[:, :, 0:resize_w] = resized_image
+    pad_shape = padding_im.shape
+
+    return padding_im, resize_shape, pad_shape, valid_ratio
+
+
+def resize_norm_img(img, image_shape, padding=True, interpolation=cv2.INTER_LINEAR):
+    imgC, imgH, imgW = image_shape
+    h = img.shape[0]
+    w = img.shape[1]
+    if not padding:
+        resized_image = cv2.resize(img, (imgW, imgH), interpolation=interpolation)
+        resized_w = imgW
+    else:
+        ratio = w / float(h)
+        if math.ceil(imgH * ratio) > imgW:
+            resized_w = imgW
+        else:
+            resized_w = int(math.ceil(imgH * ratio))
+        resized_image = cv2.resize(img, (resized_w, imgH))
+    resized_image = resized_image.astype("float32")
+    if image_shape[0] == 1:
+        resized_image = resized_image / 255
+        resized_image = resized_image[np.newaxis, :]
+    else:
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+    resized_image -= 0.5
+    resized_image /= 0.5
+    padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
+    padding_im[:, :, 0:resized_w] = resized_image
+    valid_ratio = min(1.0, float(resized_w / imgW))
+    return padding_im, valid_ratio
+
+
+def resize_norm_img_chinese(img, image_shape):
+    imgC, imgH, imgW = image_shape
+    # todo: change to 0 and modified image shape
+    max_wh_ratio = imgW * 1.0 / imgH
+    h, w = img.shape[0], img.shape[1]
+    ratio = w * 1.0 / h
+    max_wh_ratio = max(max_wh_ratio, ratio)
+    imgW = int(imgH * max_wh_ratio)
+    if math.ceil(imgH * ratio) > imgW:
+        resized_w = imgW
+    else:
+        resized_w = int(math.ceil(imgH * ratio))
+    resized_image = cv2.resize(img, (resized_w, imgH))
+    resized_image = resized_image.astype("float32")
+    if image_shape[0] == 1:
+        resized_image = resized_image / 255
+        resized_image = resized_image[np.newaxis, :]
+    else:
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+    resized_image -= 0.5
+    resized_image /= 0.5
+    padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
+    padding_im[:, :, 0:resized_w] = resized_image
+    valid_ratio = min(1.0, float(resized_w / imgW))
+    return padding_im, valid_ratio
+
+
+def resize_norm_img_srn(img, image_shape):
+    imgC, imgH, imgW = image_shape
+
+    img_black = np.zeros((imgH, imgW))
+    im_hei = img.shape[0]
+    im_wid = img.shape[1]
+
+    if im_wid <= im_hei * 1:
+        img_new = cv2.resize(img, (imgH * 1, imgH))
+    elif im_wid <= im_hei * 2:
+        img_new = cv2.resize(img, (imgH * 2, imgH))
+    elif im_wid <= im_hei * 3:
+        img_new = cv2.resize(img, (imgH * 3, imgH))
+    else:
+        img_new = cv2.resize(img, (imgW, imgH))
+
+    img_np = np.asarray(img_new)
+    img_np = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
+    img_black[:, 0 : img_np.shape[1]] = img_np
+    img_black = img_black[:, :, np.newaxis]
+
+    row, col, c = img_black.shape
+    c = 1
+
+    return np.reshape(img_black, (c, row, col)).astype(np.float32)
+
+
+def resize_norm_img_abinet(img, image_shape):
+    imgC, imgH, imgW = image_shape
+
+    resized_image = cv2.resize(img, (imgW, imgH), interpolation=cv2.INTER_LINEAR)
+    resized_w = imgW
+    resized_image = resized_image.astype("float32")
+    resized_image = resized_image / 255.0
+
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    resized_image = (resized_image - mean[None, None, ...]) / std[None, None, ...]
+    resized_image = resized_image.transpose((2, 0, 1))
+    resized_image = resized_image.astype("float32")
+
+    valid_ratio = min(1.0, float(resized_w / imgW))
+    return resized_image, valid_ratio
+
+
+def srn_other_inputs(image_shape, num_heads, max_text_length):
+    imgC, imgH, imgW = image_shape
+    feature_dim = int((imgH / 8) * (imgW / 8))
+
+    encoder_word_pos = (
+        np.array(range(0, feature_dim)).reshape((feature_dim, 1)).astype("int64")
+    )
+    gsrm_word_pos = (
+        np.array(range(0, max_text_length))
+        .reshape((max_text_length, 1))
+        .astype("int64")
+    )
+
+    gsrm_attn_bias_data = np.ones((1, max_text_length, max_text_length))
+    gsrm_slf_attn_bias1 = np.triu(gsrm_attn_bias_data, 1).reshape(
+        [1, max_text_length, max_text_length]
+    )
+    gsrm_slf_attn_bias1 = np.tile(gsrm_slf_attn_bias1, [num_heads, 1, 1]) * [-1e9]
+
+    gsrm_slf_attn_bias2 = np.tril(gsrm_attn_bias_data, -1).reshape(
+        [1, max_text_length, max_text_length]
+    )
+    gsrm_slf_attn_bias2 = np.tile(gsrm_slf_attn_bias2, [num_heads, 1, 1]) * [-1e9]
+
+    return [encoder_word_pos, gsrm_word_pos, gsrm_slf_attn_bias1, gsrm_slf_attn_bias2]
+
+
+def flag():
+    """
+    flag
+    """
+    return 1 if random.random() > 0.5000001 else -1
+
+
+def hsv_aug(img):
+    """
+    cvtColor
+    """
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+    delta = 0.001 * random.random() * flag()
+    hsv[:, :, 2] = hsv[:, :, 2] * (1 + delta)
+    new_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+    return new_img
+
+
+def blur(img):
+    """
+    blur
+    """
+    h, w, _ = img.shape
+    if h > 10 and w > 10:
+        return cv2.GaussianBlur(img, (5, 5), 1)
+    else:
+        return img
+
+
+def jitter(img):
+    """
+    jitter
+    """
+    w, h, _ = img.shape
+    if h > 10 and w > 10:
+        thres = min(w, h)
+        s = int(random.random() * thres * 0.01)
+        src_img = img.copy()
+        for i in range(s):
+            img[i:, i:, :] = src_img[: w - i, : h - i, :]
+        return img
+    else:
+        return img
+
+
+def add_gasuss_noise(image, mean=0, var=0.1):
+    """
+    Gasuss noise
+    """
+
+    noise = np.random.normal(mean, var**0.5, image.shape)
+    out = image + 0.5 * noise
+    out = np.clip(out, 0, 255)
+    out = np.uint8(out)
+    return out
+
+
+def get_crop(image):
+    """
+    random crop
+    """
+    h, w, _ = image.shape
+    top_min = 1
+    top_max = 8
+    top_crop = int(random.randint(top_min, top_max))
+    top_crop = min(top_crop, h - 1)
+    crop_img = image.copy()
+    ratio = random.randint(0, 1)
+    if ratio:
+        crop_img = crop_img[top_crop:h, :, :]
+    else:
+        crop_img = crop_img[0 : h - top_crop, :, :]
+    return crop_img
+
+
+def rad(x):
+    """
+    rad
+    """
+    return x * np.pi / 180
+
+
+def get_warpR(config):
+    """
+    get_warpR
+    """
+    anglex, angley, anglez, fov, w, h, r = (
+        config.anglex,
+        config.angley,
+        config.anglez,
+        config.fov,
+        config.w,
+        config.h,
+        config.r,
+    )
+    if w > 69 and w < 112:
+        anglex = anglex * 1.5
+
+    z = np.sqrt(w**2 + h**2) / 2 / np.tan(rad(fov / 2))
+    # Homogeneous coordinate transformation matrix
+    rx = np.array(
+        [
+            [1, 0, 0, 0],
+            [0, np.cos(rad(anglex)), -np.sin(rad(anglex)), 0],
+            [
+                0,
+                -np.sin(rad(anglex)),
+                np.cos(rad(anglex)),
+                0,
+            ],
+            [0, 0, 0, 1],
+        ],
+        np.float32,
+    )
+    ry = np.array(
+        [
+            [np.cos(rad(angley)), 0, np.sin(rad(angley)), 0],
+            [0, 1, 0, 0],
+            [
+                -np.sin(rad(angley)),
+                0,
+                np.cos(rad(angley)),
+                0,
+            ],
+            [0, 0, 0, 1],
+        ],
+        np.float32,
+    )
+    rz = np.array(
+        [
+            [np.cos(rad(anglez)), np.sin(rad(anglez)), 0, 0],
+            [-np.sin(rad(anglez)), np.cos(rad(anglez)), 0, 0],
+            [0, 0, 1, 0],
+            [0, 0, 0, 1],
+        ],
+        np.float32,
+    )
+    r = rx.dot(ry).dot(rz)
+    # generate 4 points
+    pcenter = np.array([h / 2, w / 2, 0, 0], np.float32)
+    p1 = np.array([0, 0, 0, 0], np.float32) - pcenter
+    p2 = np.array([w, 0, 0, 0], np.float32) - pcenter
+    p3 = np.array([0, h, 0, 0], np.float32) - pcenter
+    p4 = np.array([w, h, 0, 0], np.float32) - pcenter
+    dst1 = r.dot(p1)
+    dst2 = r.dot(p2)
+    dst3 = r.dot(p3)
+    dst4 = r.dot(p4)
+    list_dst = np.array([dst1, dst2, dst3, dst4])
+    org = np.array([[0, 0], [w, 0], [0, h], [w, h]], np.float32)
+    dst = np.zeros((4, 2), np.float32)
+    # Project onto the image plane
+    dst[:, 0] = list_dst[:, 0] * z / (z - list_dst[:, 2]) + pcenter[0]
+    dst[:, 1] = list_dst[:, 1] * z / (z - list_dst[:, 2]) + pcenter[1]
+
+    warpR = cv2.getPerspectiveTransform(org, dst)
+
+    dst1, dst2, dst3, dst4 = dst
+    r1 = int(min(dst1[1], dst2[1]))
+    r2 = int(max(dst3[1], dst4[1]))
+    c1 = int(min(dst1[0], dst3[0]))
+    c2 = int(max(dst2[0], dst4[0]))
+
+    try:
+        ratio = min(1.0 * h / (r2 - r1), 1.0 * w / (c2 - c1))
+
+        dx = -c1
+        dy = -r1
+        T1 = np.float32([[1.0, 0, dx], [0, 1.0, dy], [0, 0, 1.0 / ratio]])
+        ret = T1.dot(warpR)
+    except:
+        ratio = 1.0
+        T1 = np.float32([[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]])
+        ret = T1
+    return ret, (-r1, -c1), ratio, dst
+
+
+def get_warpAffine(config):
+    """
+    get_warpAffine
+    """
+    anglez = config.anglez
+    rz = np.array(
+        [
+            [np.cos(rad(anglez)), np.sin(rad(anglez)), 0],
+            [-np.sin(rad(anglez)), np.cos(rad(anglez)), 0],
+        ],
+        np.float32,
+    )
+    return rz

ppocr/data/imaug/sast_process.py

@@ -0,0 +1,810 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This part code is referred from:
+https://github.com/songdejia/EAST/blob/master/data_utils.py
+"""
+import math
+import cv2
+import numpy as np
+import json
+import sys
+import os
+
+__all__ = ["SASTProcessTrain"]
+
+
+class SASTProcessTrain(object):
+    def __init__(
+        self,
+        image_shape=[512, 512],
+        min_crop_size=24,
+        min_crop_side_ratio=0.3,
+        min_text_size=10,
+        max_text_size=512,
+        **kwargs,
+    ):
+        self.input_size = image_shape[1]
+        self.min_crop_size = min_crop_size
+        self.min_crop_side_ratio = min_crop_side_ratio
+        self.min_text_size = min_text_size
+        self.max_text_size = max_text_size
+
+    def quad_area(self, poly):
+        """
+        compute area of a polygon
+        :param poly:
+        :return:
+        """
+        edge = [
+            (poly[1][0] - poly[0][0]) * (poly[1][1] + poly[0][1]),
+            (poly[2][0] - poly[1][0]) * (poly[2][1] + poly[1][1]),
+            (poly[3][0] - poly[2][0]) * (poly[3][1] + poly[2][1]),
+            (poly[0][0] - poly[3][0]) * (poly[0][1] + poly[3][1]),
+        ]
+        return np.sum(edge) / 2.0
+
+    def gen_quad_from_poly(self, poly):
+        """
+        Generate min area quad from poly.
+        """
+        point_num = poly.shape[0]
+        min_area_quad = np.zeros((4, 2), dtype=np.float32)
+        if True:
+            rect = cv2.minAreaRect(
+                poly.astype(np.int32)
+            )  # (center (x,y), (width, height), angle of rotation)
+            center_point = rect[0]
+            box = np.array(cv2.boxPoints(rect))
+
+            first_point_idx = 0
+            min_dist = 1e4
+            for i in range(4):
+                dist = (
+                    np.linalg.norm(box[(i + 0) % 4] - poly[0])
+                    + np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1])
+                    + np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2])
+                    + np.linalg.norm(box[(i + 3) % 4] - poly[-1])
+                )
+                if dist < min_dist:
+                    min_dist = dist
+                    first_point_idx = i
+            for i in range(4):
+                min_area_quad[i] = box[(first_point_idx + i) % 4]
+
+        return min_area_quad
+
+    def check_and_validate_polys(self, polys, tags, xxx_todo_changeme):
+        """
+        check so that the text poly is in the same direction,
+        and also filter some invalid polygons
+        :param polys:
+        :param tags:
+        :return:
+        """
+        (h, w) = xxx_todo_changeme
+        if polys.shape[0] == 0:
+            return polys, np.array([]), np.array([])
+        polys[:, :, 0] = np.clip(polys[:, :, 0], 0, w - 1)
+        polys[:, :, 1] = np.clip(polys[:, :, 1], 0, h - 1)
+
+        validated_polys = []
+        validated_tags = []
+        hv_tags = []
+        for poly, tag in zip(polys, tags):
+            quad = self.gen_quad_from_poly(poly)
+            p_area = self.quad_area(quad)
+            if abs(p_area) < 1:
+                print("invalid poly")
+                continue
+            if p_area > 0:
+                if tag == False:
+                    print("poly in wrong direction")
+                    tag = True  # reversed cases should be ignore
+                poly = poly[(0, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1), :]
+                quad = quad[(0, 3, 2, 1), :]
+
+            len_w = np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(
+                quad[3] - quad[2]
+            )
+            len_h = np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(
+                quad[1] - quad[2]
+            )
+            hv_tag = 1
+
+            if len_w * 2.0 < len_h:
+                hv_tag = 0
+
+            validated_polys.append(poly)
+            validated_tags.append(tag)
+            hv_tags.append(hv_tag)
+        return np.array(validated_polys), np.array(validated_tags), np.array(hv_tags)
+
+    def crop_area(self, im, polys, tags, hv_tags, crop_background=False, max_tries=25):
+        """
+        make random crop from the input image
+        :param im:
+        :param polys:
+        :param tags:
+        :param crop_background:
+        :param max_tries: 50 -> 25
+        :return:
+        """
+        h, w, _ = im.shape
+        pad_h = h // 10
+        pad_w = w // 10
+        h_array = np.zeros((h + pad_h * 2), dtype=np.int32)
+        w_array = np.zeros((w + pad_w * 2), dtype=np.int32)
+        for poly in polys:
+            poly = np.round(poly, decimals=0).astype(np.int32)
+            minx = np.min(poly[:, 0])
+            maxx = np.max(poly[:, 0])
+            w_array[minx + pad_w : maxx + pad_w] = 1
+            miny = np.min(poly[:, 1])
+            maxy = np.max(poly[:, 1])
+            h_array[miny + pad_h : maxy + pad_h] = 1
+        # ensure the cropped area not across a text
+        h_axis = np.where(h_array == 0)[0]
+        w_axis = np.where(w_array == 0)[0]
+        if len(h_axis) == 0 or len(w_axis) == 0:
+            return im, polys, tags, hv_tags
+        for i in range(max_tries):
+            xx = np.random.choice(w_axis, size=2)
+            xmin = np.min(xx) - pad_w
+            xmax = np.max(xx) - pad_w
+            xmin = np.clip(xmin, 0, w - 1)
+            xmax = np.clip(xmax, 0, w - 1)
+            yy = np.random.choice(h_axis, size=2)
+            ymin = np.min(yy) - pad_h
+            ymax = np.max(yy) - pad_h
+            ymin = np.clip(ymin, 0, h - 1)
+            ymax = np.clip(ymax, 0, h - 1)
+            # if xmax - xmin < ARGS.min_crop_side_ratio * w or \
+            #   ymax - ymin < ARGS.min_crop_side_ratio * h:
+            if xmax - xmin < self.min_crop_size or ymax - ymin < self.min_crop_size:
+                # area too small
+                continue
+            if polys.shape[0] != 0:
+                poly_axis_in_area = (
+                    (polys[:, :, 0] >= xmin)
+                    & (polys[:, :, 0] <= xmax)
+                    & (polys[:, :, 1] >= ymin)
+                    & (polys[:, :, 1] <= ymax)
+                )
+                selected_polys = np.where(np.sum(poly_axis_in_area, axis=1) == 4)[0]
+            else:
+                selected_polys = []
+            if len(selected_polys) == 0:
+                # no text in this area
+                if crop_background:
+                    return (
+                        im[ymin : ymax + 1, xmin : xmax + 1, :],
+                        polys[selected_polys],
+                        tags[selected_polys],
+                        hv_tags[selected_polys],
+                    )
+                else:
+                    continue
+            im = im[ymin : ymax + 1, xmin : xmax + 1, :]
+            polys = polys[selected_polys]
+            tags = tags[selected_polys]
+            hv_tags = hv_tags[selected_polys]
+            polys[:, :, 0] -= xmin
+            polys[:, :, 1] -= ymin
+            return im, polys, tags, hv_tags
+
+        return im, polys, tags, hv_tags
+
+    def generate_direction_map(self, poly_quads, direction_map):
+        """ """
+        width_list = []
+        height_list = []
+        for quad in poly_quads:
+            quad_w = (
+                np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3])
+            ) / 2.0
+            quad_h = (
+                np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[2] - quad[1])
+            ) / 2.0
+            width_list.append(quad_w)
+            height_list.append(quad_h)
+        norm_width = max(sum(width_list) / (len(width_list) + 1e-6), 1.0)
+        average_height = max(sum(height_list) / (len(height_list) + 1e-6), 1.0)
+
+        for quad in poly_quads:
+            direct_vector_full = ((quad[1] + quad[2]) - (quad[0] + quad[3])) / 2.0
+            direct_vector = (
+                direct_vector_full
+                / (np.linalg.norm(direct_vector_full) + 1e-6)
+                * norm_width
+            )
+            direction_label = tuple(
+                map(
+                    float,
+                    [direct_vector[0], direct_vector[1], 1.0 / (average_height + 1e-6)],
+                )
+            )
+            cv2.fillPoly(
+                direction_map,
+                quad.round().astype(np.int32)[np.newaxis, :, :],
+                direction_label,
+            )
+        return direction_map
+
+    def calculate_average_height(self, poly_quads):
+        """ """
+        height_list = []
+        for quad in poly_quads:
+            quad_h = (
+                np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[2] - quad[1])
+            ) / 2.0
+            height_list.append(quad_h)
+        average_height = max(sum(height_list) / len(height_list), 1.0)
+        return average_height
+
+    def generate_tcl_label(
+        self, hw, polys, tags, ds_ratio, tcl_ratio=0.3, shrink_ratio_of_width=0.15
+    ):
+        """
+        Generate polygon.
+        """
+        h, w = hw
+        h, w = int(h * ds_ratio), int(w * ds_ratio)
+        polys = polys * ds_ratio
+
+        score_map = np.zeros(
+            (
+                h,
+                w,
+            ),
+            dtype=np.float32,
+        )
+        tbo_map = np.zeros((h, w, 5), dtype=np.float32)
+        training_mask = np.ones(
+            (
+                h,
+                w,
+            ),
+            dtype=np.float32,
+        )
+        direction_map = np.ones((h, w, 3)) * np.array([0, 0, 1]).reshape(
+            [1, 1, 3]
+        ).astype(np.float32)
+
+        for poly_idx, poly_tag in enumerate(zip(polys, tags)):
+            poly = poly_tag[0]
+            tag = poly_tag[1]
+
+            # generate min_area_quad
+            min_area_quad, center_point = self.gen_min_area_quad_from_poly(poly)
+            min_area_quad_h = 0.5 * (
+                np.linalg.norm(min_area_quad[0] - min_area_quad[3])
+                + np.linalg.norm(min_area_quad[1] - min_area_quad[2])
+            )
+            min_area_quad_w = 0.5 * (
+                np.linalg.norm(min_area_quad[0] - min_area_quad[1])
+                + np.linalg.norm(min_area_quad[2] - min_area_quad[3])
+            )
+
+            if (
+                min(min_area_quad_h, min_area_quad_w) < self.min_text_size * ds_ratio
+                or min(min_area_quad_h, min_area_quad_w) > self.max_text_size * ds_ratio
+            ):
+                continue
+
+            if tag:
+                # continue
+                cv2.fillPoly(
+                    training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0.15
+                )
+            else:
+                tcl_poly = self.poly2tcl(poly, tcl_ratio)
+                tcl_quads = self.poly2quads(tcl_poly)
+                poly_quads = self.poly2quads(poly)
+                # stcl map
+                stcl_quads, quad_index = self.shrink_poly_along_width(
+                    tcl_quads,
+                    shrink_ratio_of_width=shrink_ratio_of_width,
+                    expand_height_ratio=1.0 / tcl_ratio,
+                )
+                # generate tcl map
+                cv2.fillPoly(score_map, np.round(stcl_quads).astype(np.int32), 1.0)
+
+                # generate tbo map
+                for idx, quad in enumerate(stcl_quads):
+                    quad_mask = np.zeros((h, w), dtype=np.float32)
+                    quad_mask = cv2.fillPoly(
+                        quad_mask,
+                        np.round(quad[np.newaxis, :, :]).astype(np.int32),
+                        1.0,
+                    )
+                    tbo_map = self.gen_quad_tbo(
+                        poly_quads[quad_index[idx]], quad_mask, tbo_map
+                    )
+        return score_map, tbo_map, training_mask
+
+    def generate_tvo_and_tco(self, hw, polys, tags, tcl_ratio=0.3, ds_ratio=0.25):
+        """
+        Generate tcl map, tvo map and tbo map.
+        """
+        h, w = hw
+        h, w = int(h * ds_ratio), int(w * ds_ratio)
+        polys = polys * ds_ratio
+        poly_mask = np.zeros((h, w), dtype=np.float32)
+
+        tvo_map = np.ones((9, h, w), dtype=np.float32)
+        tvo_map[0:-1:2] = np.tile(np.arange(0, w), (h, 1))
+        tvo_map[1:-1:2] = np.tile(np.arange(0, w), (h, 1)).T
+        poly_tv_xy_map = np.zeros((8, h, w), dtype=np.float32)
+
+        # tco map
+        tco_map = np.ones((3, h, w), dtype=np.float32)
+        tco_map[0] = np.tile(np.arange(0, w), (h, 1))
+        tco_map[1] = np.tile(np.arange(0, w), (h, 1)).T
+        poly_tc_xy_map = np.zeros((2, h, w), dtype=np.float32)
+
+        poly_short_edge_map = np.ones((h, w), dtype=np.float32)
+
+        for poly, poly_tag in zip(polys, tags):
+            if poly_tag == True:
+                continue
+
+            # adjust point order for vertical poly
+            poly = self.adjust_point(poly)
+
+            # generate min_area_quad
+            min_area_quad, center_point = self.gen_min_area_quad_from_poly(poly)
+            min_area_quad_h = 0.5 * (
+                np.linalg.norm(min_area_quad[0] - min_area_quad[3])
+                + np.linalg.norm(min_area_quad[1] - min_area_quad[2])
+            )
+            min_area_quad_w = 0.5 * (
+                np.linalg.norm(min_area_quad[0] - min_area_quad[1])
+                + np.linalg.norm(min_area_quad[2] - min_area_quad[3])
+            )
+
+            # generate tcl map and text, 128 * 128
+            tcl_poly = self.poly2tcl(poly, tcl_ratio)
+
+            # generate poly_tv_xy_map
+            for idx in range(4):
+                cv2.fillPoly(
+                    poly_tv_xy_map[2 * idx],
+                    np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
+                    float(min(max(min_area_quad[idx, 0], 0), w)),
+                )
+                cv2.fillPoly(
+                    poly_tv_xy_map[2 * idx + 1],
+                    np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
+                    float(min(max(min_area_quad[idx, 1], 0), h)),
+                )
+
+            # generate poly_tc_xy_map
+            for idx in range(2):
+                cv2.fillPoly(
+                    poly_tc_xy_map[idx],
+                    np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
+                    float(center_point[idx]),
+                )
+
+            # generate poly_short_edge_map
+            cv2.fillPoly(
+                poly_short_edge_map,
+                np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32),
+                float(max(min(min_area_quad_h, min_area_quad_w), 1.0)),
+            )
+
+            # generate poly_mask and training_mask
+            cv2.fillPoly(
+                poly_mask, np.round(tcl_poly[np.newaxis, :, :]).astype(np.int32), 1
+            )
+
+        tvo_map *= poly_mask
+        tvo_map[:8] -= poly_tv_xy_map
+        tvo_map[-1] /= poly_short_edge_map
+        tvo_map = tvo_map.transpose((1, 2, 0))
+
+        tco_map *= poly_mask
+        tco_map[:2] -= poly_tc_xy_map
+        tco_map[-1] /= poly_short_edge_map
+        tco_map = tco_map.transpose((1, 2, 0))
+
+        return tvo_map, tco_map
+
+    def adjust_point(self, poly):
+        """
+        adjust point order.
+        """
+        point_num = poly.shape[0]
+        if point_num == 4:
+            len_1 = np.linalg.norm(poly[0] - poly[1])
+            len_2 = np.linalg.norm(poly[1] - poly[2])
+            len_3 = np.linalg.norm(poly[2] - poly[3])
+            len_4 = np.linalg.norm(poly[3] - poly[0])
+
+            if (len_1 + len_3) * 1.5 < (len_2 + len_4):
+                poly = poly[[1, 2, 3, 0], :]
+
+        elif point_num > 4:
+            vector_1 = poly[0] - poly[1]
+            vector_2 = poly[1] - poly[2]
+            cos_theta = np.dot(vector_1, vector_2) / (
+                np.linalg.norm(vector_1) * np.linalg.norm(vector_2) + 1e-6
+            )
+            theta = np.arccos(np.round(cos_theta, decimals=4))
+
+            if abs(theta) > (70 / 180 * math.pi):
+                index = list(range(1, point_num)) + [0]
+                poly = poly[np.array(index), :]
+        return poly
+
+    def gen_min_area_quad_from_poly(self, poly):
+        """
+        Generate min area quad from poly.
+        """
+        point_num = poly.shape[0]
+        min_area_quad = np.zeros((4, 2), dtype=np.float32)
+        if point_num == 4:
+            min_area_quad = poly
+            center_point = np.sum(poly, axis=0) / 4
+        else:
+            rect = cv2.minAreaRect(
+                poly.astype(np.int32)
+            )  # (center (x,y), (width, height), angle of rotation)
+            center_point = rect[0]
+            box = np.array(cv2.boxPoints(rect))
+
+            first_point_idx = 0
+            min_dist = 1e4
+            for i in range(4):
+                dist = (
+                    np.linalg.norm(box[(i + 0) % 4] - poly[0])
+                    + np.linalg.norm(box[(i + 1) % 4] - poly[point_num // 2 - 1])
+                    + np.linalg.norm(box[(i + 2) % 4] - poly[point_num // 2])
+                    + np.linalg.norm(box[(i + 3) % 4] - poly[-1])
+                )
+                if dist < min_dist:
+                    min_dist = dist
+                    first_point_idx = i
+
+            for i in range(4):
+                min_area_quad[i] = box[(first_point_idx + i) % 4]
+
+        return min_area_quad, center_point
+
+    def shrink_quad_along_width(self, quad, begin_width_ratio=0.0, end_width_ratio=1.0):
+        """
+        Generate shrink_quad_along_width.
+        """
+        ratio_pair = np.array(
+            [[begin_width_ratio], [end_width_ratio]], dtype=np.float32
+        )
+        p0_1 = quad[0] + (quad[1] - quad[0]) * ratio_pair
+        p3_2 = quad[3] + (quad[2] - quad[3]) * ratio_pair
+        return np.array([p0_1[0], p0_1[1], p3_2[1], p3_2[0]])
+
+    def shrink_poly_along_width(
+        self, quads, shrink_ratio_of_width, expand_height_ratio=1.0
+    ):
+        """
+        shrink poly with given length.
+        """
+        upper_edge_list = []
+
+        def get_cut_info(edge_len_list, cut_len):
+            for idx, edge_len in enumerate(edge_len_list):
+                cut_len -= edge_len
+                if cut_len <= 0.000001:
+                    ratio = (cut_len + edge_len_list[idx]) / edge_len_list[idx]
+                    return idx, ratio
+
+        for quad in quads:
+            upper_edge_len = np.linalg.norm(quad[0] - quad[1])
+            upper_edge_list.append(upper_edge_len)
+
+        # length of left edge and right edge.
+        left_length = np.linalg.norm(quads[0][0] - quads[0][3]) * expand_height_ratio
+        right_length = np.linalg.norm(quads[-1][1] - quads[-1][2]) * expand_height_ratio
+
+        shrink_length = (
+            min(left_length, right_length, sum(upper_edge_list)) * shrink_ratio_of_width
+        )
+        # shrinking length
+        upper_len_left = shrink_length
+        upper_len_right = sum(upper_edge_list) - shrink_length
+
+        left_idx, left_ratio = get_cut_info(upper_edge_list, upper_len_left)
+        left_quad = self.shrink_quad_along_width(
+            quads[left_idx], begin_width_ratio=left_ratio, end_width_ratio=1
+        )
+        right_idx, right_ratio = get_cut_info(upper_edge_list, upper_len_right)
+        right_quad = self.shrink_quad_along_width(
+            quads[right_idx], begin_width_ratio=0, end_width_ratio=right_ratio
+        )
+
+        out_quad_list = []
+        if left_idx == right_idx:
+            out_quad_list.append(
+                [left_quad[0], right_quad[1], right_quad[2], left_quad[3]]
+            )
+        else:
+            out_quad_list.append(left_quad)
+            for idx in range(left_idx + 1, right_idx):
+                out_quad_list.append(quads[idx])
+            out_quad_list.append(right_quad)
+
+        return np.array(out_quad_list), list(range(left_idx, right_idx + 1))
+
+    def vector_angle(self, A, B):
+        """
+        Calculate the angle between vector AB and x-axis positive direction.
+        """
+        AB = np.array([B[1] - A[1], B[0] - A[0]])
+        return np.arctan2(*AB)
+
+    def theta_line_cross_point(self, theta, point):
+        """
+        Calculate the line through given point and angle in ax + by + c =0 form.
+        """
+        x, y = point
+        cos = np.cos(theta)
+        sin = np.sin(theta)
+        return [sin, -cos, cos * y - sin * x]
+
+    def line_cross_two_point(self, A, B):
+        """
+        Calculate the line through given point A and B in ax + by + c =0 form.
+        """
+        angle = self.vector_angle(A, B)
+        return self.theta_line_cross_point(angle, A)
+
+    def average_angle(self, poly):
+        """
+        Calculate the average angle between left and right edge in given poly.
+        """
+        p0, p1, p2, p3 = poly
+        angle30 = self.vector_angle(p3, p0)
+        angle21 = self.vector_angle(p2, p1)
+        return (angle30 + angle21) / 2
+
+    def line_cross_point(self, line1, line2):
+        """
+        line1 and line2 in  0=ax+by+c form, compute the cross point of line1 and line2
+        """
+        a1, b1, c1 = line1
+        a2, b2, c2 = line2
+        d = a1 * b2 - a2 * b1
+
+        if d == 0:
+            # print("line1", line1)
+            # print("line2", line2)
+            print("Cross point does not exist")
+            return np.array([0, 0], dtype=np.float32)
+        else:
+            x = (b1 * c2 - b2 * c1) / d
+            y = (a2 * c1 - a1 * c2) / d
+
+        return np.array([x, y], dtype=np.float32)
+
+    def quad2tcl(self, poly, ratio):
+        """
+        Generate center line by poly clock-wise point. (4, 2)
+        """
+        ratio_pair = np.array([[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
+        p0_3 = poly[0] + (poly[3] - poly[0]) * ratio_pair
+        p1_2 = poly[1] + (poly[2] - poly[1]) * ratio_pair
+        return np.array([p0_3[0], p1_2[0], p1_2[1], p0_3[1]])
+
+    def poly2tcl(self, poly, ratio):
+        """
+        Generate center line by poly clock-wise point.
+        """
+        ratio_pair = np.array([[0.5 - ratio / 2], [0.5 + ratio / 2]], dtype=np.float32)
+        tcl_poly = np.zeros_like(poly)
+        point_num = poly.shape[0]
+
+        for idx in range(point_num // 2):
+            point_pair = (
+                poly[idx] + (poly[point_num - 1 - idx] - poly[idx]) * ratio_pair
+            )
+            tcl_poly[idx] = point_pair[0]
+            tcl_poly[point_num - 1 - idx] = point_pair[1]
+        return tcl_poly
+
+    def gen_quad_tbo(self, quad, tcl_mask, tbo_map):
+        """
+        Generate tbo_map for give quad.
+        """
+        # upper and lower line function: ax + by + c = 0;
+        up_line = self.line_cross_two_point(quad[0], quad[1])
+        lower_line = self.line_cross_two_point(quad[3], quad[2])
+
+        quad_h = 0.5 * (
+            np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[1] - quad[2])
+        )
+        quad_w = 0.5 * (
+            np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3])
+        )
+
+        # average angle of left and right line.
+        angle = self.average_angle(quad)
+
+        xy_in_poly = np.argwhere(tcl_mask == 1)
+        for y, x in xy_in_poly:
+            point = (x, y)
+            line = self.theta_line_cross_point(angle, point)
+            cross_point_upper = self.line_cross_point(up_line, line)
+            cross_point_lower = self.line_cross_point(lower_line, line)
+            ##FIX, offset reverse
+            upper_offset_x, upper_offset_y = cross_point_upper - point
+            lower_offset_x, lower_offset_y = cross_point_lower - point
+            tbo_map[y, x, 0] = upper_offset_y
+            tbo_map[y, x, 1] = upper_offset_x
+            tbo_map[y, x, 2] = lower_offset_y
+            tbo_map[y, x, 3] = lower_offset_x
+            tbo_map[y, x, 4] = 1.0 / max(min(quad_h, quad_w), 1.0) * 2
+        return tbo_map
+
+    def poly2quads(self, poly):
+        """
+        Split poly into quads.
+        """
+        quad_list = []
+        point_num = poly.shape[0]
+
+        # point pair
+        point_pair_list = []
+        for idx in range(point_num // 2):
+            point_pair = [poly[idx], poly[point_num - 1 - idx]]
+            point_pair_list.append(point_pair)
+
+        quad_num = point_num // 2 - 1
+        for idx in range(quad_num):
+            # reshape and adjust to clock-wise
+            quad_list.append(
+                (np.array(point_pair_list)[[idx, idx + 1]]).reshape(4, 2)[[0, 2, 3, 1]]
+            )
+
+        return np.array(quad_list)
+
+    def __call__(self, data):
+        im = data["image"]
+        text_polys = data["polys"]
+        text_tags = data["ignore_tags"]
+        if im is None:
+            return None
+        if text_polys.shape[0] == 0:
+            return None
+
+        h, w, _ = im.shape
+        text_polys, text_tags, hv_tags = self.check_and_validate_polys(
+            text_polys, text_tags, (h, w)
+        )
+
+        if text_polys.shape[0] == 0:
+            return None
+
+        # set aspect ratio and keep area fix
+        asp_scales = np.arange(1.0, 1.55, 0.1)
+        asp_scale = np.random.choice(asp_scales)
+
+        if np.random.rand() < 0.5:
+            asp_scale = 1.0 / asp_scale
+        asp_scale = math.sqrt(asp_scale)
+
+        asp_wx = asp_scale
+        asp_hy = 1.0 / asp_scale
+        im = cv2.resize(im, dsize=None, fx=asp_wx, fy=asp_hy)
+        text_polys[:, :, 0] *= asp_wx
+        text_polys[:, :, 1] *= asp_hy
+
+        h, w, _ = im.shape
+        if max(h, w) > 2048:
+            rd_scale = 2048.0 / max(h, w)
+            im = cv2.resize(im, dsize=None, fx=rd_scale, fy=rd_scale)
+            text_polys *= rd_scale
+        h, w, _ = im.shape
+        if min(h, w) < 16:
+            return None
+
+        # no background
+        im, text_polys, text_tags, hv_tags = self.crop_area(
+            im, text_polys, text_tags, hv_tags, crop_background=False
+        )
+
+        if text_polys.shape[0] == 0:
+            return None
+        # continue for all ignore case
+        if np.sum((text_tags * 1.0)) >= text_tags.size:
+            return None
+        new_h, new_w, _ = im.shape
+        if (new_h is None) or (new_w is None):
+            return None
+        # resize image
+        std_ratio = float(self.input_size) / max(new_w, new_h)
+        rand_scales = np.array(
+            [0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0, 1.0, 1.0, 1.0, 1.0]
+        )
+        rz_scale = std_ratio * np.random.choice(rand_scales)
+        im = cv2.resize(im, dsize=None, fx=rz_scale, fy=rz_scale)
+        text_polys[:, :, 0] *= rz_scale
+        text_polys[:, :, 1] *= rz_scale
+
+        # add gaussian blur
+        if np.random.rand() < 0.1 * 0.5:
+            ks = np.random.permutation(5)[0] + 1
+            ks = int(ks / 2) * 2 + 1
+            im = cv2.GaussianBlur(im, ksize=(ks, ks), sigmaX=0, sigmaY=0)
+        # add brighter
+        if np.random.rand() < 0.1 * 0.5:
+            im = im * (1.0 + np.random.rand() * 0.5)
+            im = np.clip(im, 0.0, 255.0)
+        # add darker
+        if np.random.rand() < 0.1 * 0.5:
+            im = im * (1.0 - np.random.rand() * 0.5)
+            im = np.clip(im, 0.0, 255.0)
+
+        # Padding the im to [input_size, input_size]
+        new_h, new_w, _ = im.shape
+        if min(new_w, new_h) < self.input_size * 0.5:
+            return None
+
+        im_padded = np.ones((self.input_size, self.input_size, 3), dtype=np.float32)
+        im_padded[:, :, 2] = 0.485 * 255
+        im_padded[:, :, 1] = 0.456 * 255
+        im_padded[:, :, 0] = 0.406 * 255
+
+        # Random the start position
+        del_h = self.input_size - new_h
+        del_w = self.input_size - new_w
+        sh, sw = 0, 0
+        if del_h > 1:
+            sh = int(np.random.rand() * del_h)
+        if del_w > 1:
+            sw = int(np.random.rand() * del_w)
+
+        # Padding
+        im_padded[sh : sh + new_h, sw : sw + new_w, :] = im.copy()
+        text_polys[:, :, 0] += sw
+        text_polys[:, :, 1] += sh
+
+        score_map, border_map, training_mask = self.generate_tcl_label(
+            (self.input_size, self.input_size), text_polys, text_tags, 0.25
+        )
+
+        # SAST head
+        tvo_map, tco_map = self.generate_tvo_and_tco(
+            (self.input_size, self.input_size),
+            text_polys,
+            text_tags,
+            tcl_ratio=0.3,
+            ds_ratio=0.25,
+        )
+        # print("test--------tvo_map shape:", tvo_map.shape)
+
+        im_padded[:, :, 2] -= 0.485 * 255
+        im_padded[:, :, 1] -= 0.456 * 255
+        im_padded[:, :, 0] -= 0.406 * 255
+        im_padded[:, :, 2] /= 255.0 * 0.229
+        im_padded[:, :, 1] /= 255.0 * 0.224
+        im_padded[:, :, 0] /= 255.0 * 0.225
+        im_padded = im_padded.transpose((2, 0, 1))
+
+        data["image"] = im_padded[::-1, :, :]
+        data["score_map"] = score_map[np.newaxis, :, :]
+        data["border_map"] = border_map.transpose((2, 0, 1))
+        data["training_mask"] = training_mask[np.newaxis, :, :]
+        data["tvo_map"] = tvo_map.transpose((2, 0, 1))
+        data["tco_map"] = tco_map.transpose((2, 0, 1))
+        return data

ppocr/data/imaug/ssl_img_aug.py

@@ -0,0 +1,55 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import math
+import cv2
+import numpy as np
+import random
+from PIL import Image
+
+from .rec_img_aug import resize_norm_img
+
+
+class SSLRotateResize(object):
+    def __init__(
+        self, image_shape, padding=False, select_all=True, mode="train", **kwargs
+    ):
+        self.image_shape = image_shape
+        self.padding = padding
+        self.select_all = select_all
+        self.mode = mode
+
+    def __call__(self, data):
+        img = data["image"]
+
+        data["image_r90"] = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
+        data["image_r180"] = cv2.rotate(data["image_r90"], cv2.ROTATE_90_CLOCKWISE)
+        data["image_r270"] = cv2.rotate(data["image_r180"], cv2.ROTATE_90_CLOCKWISE)
+
+        images = []
+        for key in ["image", "image_r90", "image_r180", "image_r270"]:
+            images.append(
+                resize_norm_img(
+                    data.pop(key), image_shape=self.image_shape, padding=self.padding
+                )[0]
+            )
+        data["image"] = np.stack(images, axis=0)
+        data["label"] = np.array(list(range(4)))
+        if not self.select_all:
+            data["image"] = data["image"][0::2]  # just choose 0 and 180
+            data["label"] = data["label"][0:2]  # label needs to be continuous
+        if self.mode == "test":
+            data["image"] = data["image"][0]
+            data["label"] = data["label"][0]
+        return data

ppocr/data/imaug/table_ops.py

@@ -0,0 +1,232 @@
+"""
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
+#
+# 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
+#
+#     http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import sys
+import cv2
+import numpy as np
+
+
+class GenTableMask(object):
+    """gen table mask"""
+
+    def __init__(self, shrink_h_max, shrink_w_max, mask_type=0, **kwargs):
+        self.shrink_h_max = 5
+        self.shrink_w_max = 5
+        self.mask_type = mask_type
+
+    def projection(self, erosion, h, w, spilt_threshold=0):
+        # 水平投影
+        projection_map = np.ones_like(erosion)
+        project_val_array = [0 for _ in range(0, h)]
+
+        for j in range(0, h):
+            for i in range(0, w):
+                if erosion[j, i] == 255:
+                    project_val_array[j] += 1
+        # 根据数组，获取切割点
+        start_idx = 0  # 记录进入字符区的索引
+        end_idx = 0  # 记录进入空白区域的索引
+        in_text = False  # 是否遍历到了字符区内
+        box_list = []
+        for i in range(len(project_val_array)):
+            if (
+                in_text == False and project_val_array[i] > spilt_threshold
+            ):  # 进入字符区了
+                in_text = True
+                start_idx = i
+            elif (
+                project_val_array[i] <= spilt_threshold and in_text == True
+            ):  # 进入空白区了
+                end_idx = i
+                in_text = False
+                if end_idx - start_idx <= 2:
+                    continue
+                box_list.append((start_idx, end_idx + 1))
+
+        if in_text:
+            box_list.append((start_idx, h - 1))
+        # 绘制投影直方图
+        for j in range(0, h):
+            for i in range(0, project_val_array[j]):
+                projection_map[j, i] = 0
+        return box_list, projection_map
+
+    def projection_cx(self, box_img):
+        box_gray_img = cv2.cvtColor(box_img, cv2.COLOR_BGR2GRAY)
+        h, w = box_gray_img.shape
+        # 灰度图片进行二值化处理
+        ret, thresh1 = cv2.threshold(box_gray_img, 200, 255, cv2.THRESH_BINARY_INV)
+        # 纵向腐蚀
+        if h < w:
+            kernel = np.ones((2, 1), np.uint8)
+            erode = cv2.erode(thresh1, kernel, iterations=1)
+        else:
+            erode = thresh1
+        # 水平膨胀
+        kernel = np.ones((1, 5), np.uint8)
+        erosion = cv2.dilate(erode, kernel, iterations=1)
+        # 水平投影
+        projection_map = np.ones_like(erosion)
+        project_val_array = [0 for _ in range(0, h)]
+
+        for j in range(0, h):
+            for i in range(0, w):
+                if erosion[j, i] == 255:
+                    project_val_array[j] += 1
+        # 根据数组，获取切割点
+        start_idx = 0  # 记录进入字符区的索引
+        end_idx = 0  # 记录进入空白区域的索引
+        in_text = False  # 是否遍历到了字符区内
+        box_list = []
+        spilt_threshold = 0
+        for i in range(len(project_val_array)):
+            if (
+                in_text == False and project_val_array[i] > spilt_threshold
+            ):  # 进入字符区了
+                in_text = True
+                start_idx = i
+            elif (
+                project_val_array[i] <= spilt_threshold and in_text == True
+            ):  # 进入空白区了
+                end_idx = i
+                in_text = False
+                if end_idx - start_idx <= 2:
+                    continue
+                box_list.append((start_idx, end_idx + 1))
+
+        if in_text:
+            box_list.append((start_idx, h - 1))
+        # 绘制投影直方图
+        for j in range(0, h):
+            for i in range(0, project_val_array[j]):
+                projection_map[j, i] = 0
+        split_bbox_list = []
+        if len(box_list) > 1:
+            for i, (h_start, h_end) in enumerate(box_list):
+                if i == 0:
+                    h_start = 0
+                if i == len(box_list):
+                    h_end = h
+                word_img = erosion[h_start : h_end + 1, :]
+                word_h, word_w = word_img.shape
+                w_split_list, w_projection_map = self.projection(
+                    word_img.T, word_w, word_h
+                )
+                w_start, w_end = w_split_list[0][0], w_split_list[-1][1]
+                if h_start > 0:
+                    h_start -= 1
+                h_end += 1
+                word_img = box_img[h_start : h_end + 1 :, w_start : w_end + 1, :]
+                split_bbox_list.append([w_start, h_start, w_end, h_end])
+        else:
+            split_bbox_list.append([0, 0, w, h])
+        return split_bbox_list
+
+    def shrink_bbox(self, bbox):
+        left, top, right, bottom = bbox
+        sh_h = min(max(int((bottom - top) * 0.1), 1), self.shrink_h_max)
+        sh_w = min(max(int((right - left) * 0.1), 1), self.shrink_w_max)
+        left_new = left + sh_w
+        right_new = right - sh_w
+        top_new = top + sh_h
+        bottom_new = bottom - sh_h
+        if left_new >= right_new:
+            left_new = left
+            right_new = right
+        if top_new >= bottom_new:
+            top_new = top
+            bottom_new = bottom
+        return [left_new, top_new, right_new, bottom_new]
+
+    def __call__(self, data):
+        img = data["image"]
+        cells = data["cells"]
+        height, width = img.shape[0:2]
+        if self.mask_type == 1:
+            mask_img = np.zeros((height, width), dtype=np.float32)
+        else:
+            mask_img = np.zeros((height, width, 3), dtype=np.float32)
+        cell_num = len(cells)
+        for cno in range(cell_num):
+            if "bbox" in cells[cno]:
+                bbox = cells[cno]["bbox"]
+                left, top, right, bottom = bbox
+                box_img = img[top:bottom, left:right, :].copy()
+                split_bbox_list = self.projection_cx(box_img)
+                for sno in range(len(split_bbox_list)):
+                    split_bbox_list[sno][0] += left
+                    split_bbox_list[sno][1] += top
+                    split_bbox_list[sno][2] += left
+                    split_bbox_list[sno][3] += top
+
+                for sno in range(len(split_bbox_list)):
+                    left, top, right, bottom = split_bbox_list[sno]
+                    left, top, right, bottom = self.shrink_bbox(
+                        [left, top, right, bottom]
+                    )
+                    if self.mask_type == 1:
+                        mask_img[top:bottom, left:right] = 1.0
+                        data["mask_img"] = mask_img
+                    else:
+                        mask_img[top:bottom, left:right, :] = (255, 255, 255)
+                        data["image"] = mask_img
+        return data
+
+
+class ResizeTableImage(object):
+    def __init__(self, max_len, resize_bboxes=False, infer_mode=False, **kwargs):
+        super(ResizeTableImage, self).__init__()
+        self.max_len = max_len
+        self.resize_bboxes = resize_bboxes
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        img = data["image"]
+        height, width = img.shape[0:2]
+        ratio = self.max_len / (max(height, width) * 1.0)
+        resize_h = int(height * ratio)
+        resize_w = int(width * ratio)
+        resize_img = cv2.resize(img, (resize_w, resize_h))
+        if self.resize_bboxes and not self.infer_mode:
+            data["bboxes"] = data["bboxes"] * ratio
+        data["image"] = resize_img
+        data["src_img"] = img
+        data["shape"] = np.array([height, width, ratio, ratio])
+        data["max_len"] = self.max_len
+        return data
+
+
+class PaddingTableImage(object):
+    def __init__(self, size, **kwargs):
+        super(PaddingTableImage, self).__init__()
+        self.size = size
+
+    def __call__(self, data):
+        img = data["image"]
+        pad_h, pad_w = self.size
+        padding_img = np.zeros((pad_h, pad_w, 3), dtype=np.float32)
+        height, width = img.shape[0:2]
+        padding_img[0:height, 0:width, :] = img.copy()
+        data["image"] = padding_img
+        shape = data["shape"].tolist()
+        shape.extend([pad_h, pad_w])
+        data["shape"] = np.array(shape)
+        return data

ppocr/data/imaug/text_image_aug/__init__.py

@@ -0,0 +1,17 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 .augment import tia_perspective, tia_distort, tia_stretch
+
+__all__ = ["tia_distort", "tia_stretch", "tia_perspective"]

ppocr/data/imaug/text_image_aug/augment.py

@@ -0,0 +1,123 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/RubanSeven/Text-Image-Augmentation-python/blob/master/augment.py
+"""
+
+import numpy as np
+from .warp_mls import WarpMLS
+
+
+def tia_distort(src, segment=4):
+    img_h, img_w = src.shape[:2]
+
+    cut = img_w // segment
+    thresh = cut // 3
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([np.random.randint(thresh), np.random.randint(thresh)])
+    dst_pts.append([img_w - np.random.randint(thresh), np.random.randint(thresh)])
+    dst_pts.append(
+        [img_w - np.random.randint(thresh), img_h - np.random.randint(thresh)]
+    )
+    dst_pts.append([np.random.randint(thresh), img_h - np.random.randint(thresh)])
+
+    half_thresh = thresh * 0.5
+
+    for cut_idx in np.arange(1, segment, 1):
+        src_pts.append([cut * cut_idx, 0])
+        src_pts.append([cut * cut_idx, img_h])
+        dst_pts.append(
+            [
+                cut * cut_idx + np.random.randint(thresh) - half_thresh,
+                np.random.randint(thresh) - half_thresh,
+            ]
+        )
+        dst_pts.append(
+            [
+                cut * cut_idx + np.random.randint(thresh) - half_thresh,
+                img_h + np.random.randint(thresh) - half_thresh,
+            ]
+        )
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst
+
+
+def tia_stretch(src, segment=4):
+    img_h, img_w = src.shape[:2]
+
+    cut = img_w // segment
+    thresh = cut * 4 // 5
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([0, 0])
+    dst_pts.append([img_w, 0])
+    dst_pts.append([img_w, img_h])
+    dst_pts.append([0, img_h])
+
+    half_thresh = thresh * 0.5
+
+    for cut_idx in np.arange(1, segment, 1):
+        move = np.random.randint(thresh) - half_thresh
+        src_pts.append([cut * cut_idx, 0])
+        src_pts.append([cut * cut_idx, img_h])
+        dst_pts.append([cut * cut_idx + move, 0])
+        dst_pts.append([cut * cut_idx + move, img_h])
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst
+
+
+def tia_perspective(src):
+    img_h, img_w = src.shape[:2]
+
+    thresh = img_h // 2
+
+    src_pts = list()
+    dst_pts = list()
+
+    src_pts.append([0, 0])
+    src_pts.append([img_w, 0])
+    src_pts.append([img_w, img_h])
+    src_pts.append([0, img_h])
+
+    dst_pts.append([0, np.random.randint(thresh)])
+    dst_pts.append([img_w, np.random.randint(thresh)])
+    dst_pts.append([img_w, img_h - np.random.randint(thresh)])
+    dst_pts.append([0, img_h - np.random.randint(thresh)])
+
+    trans = WarpMLS(src, src_pts, dst_pts, img_w, img_h)
+    dst = trans.generate()
+
+    return dst

ppocr/data/imaug/text_image_aug/warp_mls.py

@@ -0,0 +1,187 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/RubanSeven/Text-Image-Augmentation-python/blob/master/warp_mls.py
+"""
+
+import numpy as np
+
+
+class WarpMLS:
+    def __init__(self, src, src_pts, dst_pts, dst_w, dst_h, trans_ratio=1.0):
+        self.src = src
+        self.src_pts = src_pts
+        self.dst_pts = dst_pts
+        self.pt_count = len(self.dst_pts)
+        self.dst_w = dst_w
+        self.dst_h = dst_h
+        self.trans_ratio = trans_ratio
+        self.grid_size = 100
+        self.rdx = np.zeros((self.dst_h, self.dst_w))
+        self.rdy = np.zeros((self.dst_h, self.dst_w))
+
+    @staticmethod
+    def __bilinear_interp(x, y, v11, v12, v21, v22):
+        return (v11 * (1 - y) + v12 * y) * (1 - x) + (v21 * (1 - y) + v22 * y) * x
+
+    def generate(self):
+        self.calc_delta()
+        return self.gen_img()
+
+    def calc_delta(self):
+        w = np.zeros(self.pt_count, dtype=np.float32)
+
+        if self.pt_count < 2:
+            return
+
+        i = 0
+        while 1:
+            if self.dst_w <= i < self.dst_w + self.grid_size - 1:
+                i = self.dst_w - 1
+            elif i >= self.dst_w:
+                break
+
+            j = 0
+            while 1:
+                if self.dst_h <= j < self.dst_h + self.grid_size - 1:
+                    j = self.dst_h - 1
+                elif j >= self.dst_h:
+                    break
+
+                sw = 0
+                swp = np.zeros(2, dtype=np.float32)
+                swq = np.zeros(2, dtype=np.float32)
+                new_pt = np.zeros(2, dtype=np.float32)
+                cur_pt = np.array([i, j], dtype=np.float32)
+
+                k = 0
+                for k in range(self.pt_count):
+                    if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                        break
+
+                    w[k] = 1.0 / (
+                        (i - self.dst_pts[k][0]) * (i - self.dst_pts[k][0])
+                        + (j - self.dst_pts[k][1]) * (j - self.dst_pts[k][1])
+                    )
+
+                    sw += w[k]
+                    swp = swp + w[k] * np.array(self.dst_pts[k])
+                    swq = swq + w[k] * np.array(self.src_pts[k])
+
+                if k == self.pt_count - 1:
+                    pstar = 1 / sw * swp
+                    qstar = 1 / sw * swq
+
+                    miu_s = 0
+                    for k in range(self.pt_count):
+                        if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                            continue
+                        pt_i = self.dst_pts[k] - pstar
+                        miu_s += w[k] * np.sum(pt_i * pt_i)
+
+                    cur_pt -= pstar
+                    cur_pt_j = np.array([-cur_pt[1], cur_pt[0]])
+
+                    for k in range(self.pt_count):
+                        if i == self.dst_pts[k][0] and j == self.dst_pts[k][1]:
+                            continue
+
+                        pt_i = self.dst_pts[k] - pstar
+                        pt_j = np.array([-pt_i[1], pt_i[0]])
+
+                        tmp_pt = np.zeros(2, dtype=np.float32)
+                        tmp_pt[0] = (
+                            np.sum(pt_i * cur_pt) * self.src_pts[k][0]
+                            - np.sum(pt_j * cur_pt) * self.src_pts[k][1]
+                        )
+                        tmp_pt[1] = (
+                            -np.sum(pt_i * cur_pt_j) * self.src_pts[k][0]
+                            + np.sum(pt_j * cur_pt_j) * self.src_pts[k][1]
+                        )
+                        tmp_pt *= w[k] / miu_s
+                        new_pt += tmp_pt
+
+                    new_pt += qstar
+                else:
+                    new_pt = self.src_pts[k]
+
+                self.rdx[j, i] = new_pt[0] - i
+                self.rdy[j, i] = new_pt[1] - j
+
+                j += self.grid_size
+            i += self.grid_size
+
+    def gen_img(self):
+        src_h, src_w = self.src.shape[:2]
+        dst = np.zeros_like(self.src, dtype=np.float32)
+
+        for i in np.arange(0, self.dst_h, self.grid_size):
+            for j in np.arange(0, self.dst_w, self.grid_size):
+                ni = i + self.grid_size
+                nj = j + self.grid_size
+                w = h = self.grid_size
+                if ni >= self.dst_h:
+                    ni = self.dst_h - 1
+                    h = ni - i + 1
+                if nj >= self.dst_w:
+                    nj = self.dst_w - 1
+                    w = nj - j + 1
+
+                di = np.reshape(np.arange(h), (-1, 1))
+                dj = np.reshape(np.arange(w), (1, -1))
+                delta_x = self.__bilinear_interp(
+                    di / h,
+                    dj / w,
+                    self.rdx[i, j],
+                    self.rdx[i, nj],
+                    self.rdx[ni, j],
+                    self.rdx[ni, nj],
+                )
+                delta_y = self.__bilinear_interp(
+                    di / h,
+                    dj / w,
+                    self.rdy[i, j],
+                    self.rdy[i, nj],
+                    self.rdy[ni, j],
+                    self.rdy[ni, nj],
+                )
+                nx = j + dj + delta_x * self.trans_ratio
+                ny = i + di + delta_y * self.trans_ratio
+                nx = np.clip(nx, 0, src_w - 1)
+                ny = np.clip(ny, 0, src_h - 1)
+                nxi = np.array(np.floor(nx), dtype=np.int32)
+                nyi = np.array(np.floor(ny), dtype=np.int32)
+                nxi1 = np.array(np.ceil(nx), dtype=np.int32)
+                nyi1 = np.array(np.ceil(ny), dtype=np.int32)
+
+                if len(self.src.shape) == 3:
+                    x = np.tile(np.expand_dims(ny - nyi, axis=-1), (1, 1, 3))
+                    y = np.tile(np.expand_dims(nx - nxi, axis=-1), (1, 1, 3))
+                else:
+                    x = ny - nyi
+                    y = nx - nxi
+                dst[i : i + h, j : j + w] = self.__bilinear_interp(
+                    x,
+                    y,
+                    self.src[nyi, nxi],
+                    self.src[nyi, nxi1],
+                    self.src[nyi1, nxi],
+                    self.src[nyi1, nxi1],
+                )
+
+        dst = np.clip(dst, 0, 255)
+        dst = np.array(dst, dtype=np.uint8)
+
+        return dst

ppocr/data/imaug/unimernet_aug.py

@@ -0,0 +1,807 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import os
+
+os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
+
+import cv2
+import math
+import numpy as np
+from io import BytesIO
+import albumentations as A
+from PIL import Image, ImageOps, ImageDraw
+from scipy.ndimage import zoom as scizoom
+
+
+class Erosion(A.ImageOnlyTransform):
+    def __init__(self, scale, always_apply=False, p=0.5):
+        super().__init__(always_apply=always_apply, p=p)
+        if type(scale) is tuple or type(scale) is list:
+            assert len(scale) == 2
+            self.scale = scale
+        else:
+            self.scale = (scale, scale)
+
+    def apply(self, img, **params):
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2))
+        )
+        img = cv2.erode(img, kernel, iterations=1)
+        return img
+
+
+class Dilation(A.ImageOnlyTransform):
+    def __init__(self, scale, always_apply=False, p=0.5):
+        super().__init__(always_apply=always_apply, p=p)
+        if type(scale) is tuple or type(scale) is list:
+            assert len(scale) == 2
+            self.scale = scale
+        else:
+            self.scale = (scale, scale)
+
+    def apply(self, img, **params):
+        kernel = cv2.getStructuringElement(
+            cv2.MORPH_ELLIPSE, tuple(np.random.randint(self.scale[0], self.scale[1], 2))
+        )
+        img = cv2.dilate(img, kernel, iterations=1)
+        return img
+
+
+class Bitmap(A.ImageOnlyTransform):
+
+    def __init__(self, value=0, lower=200, always_apply=False, p=0.5):
+        super().__init__(always_apply=always_apply, p=p)
+        self.lower = lower
+        self.value = value
+
+    def apply(self, img, **params):
+        img = img.copy()
+        img[img < self.lower] = self.value
+        return img
+
+
+def clipped_zoom(img, zoom_factor):
+    h = img.shape[1]
+    ch = int(np.ceil(h / float(zoom_factor)))
+    top = (h - ch) // 2
+    img = scizoom(
+        img[top : top + ch, top : top + ch], (zoom_factor, zoom_factor, 1), order=1
+    )
+    trim_top = (img.shape[0] - h) // 2
+
+    return img[trim_top : trim_top + h, trim_top : trim_top + h]
+
+
+def disk(radius, alias_blur=0.1, dtype=np.float32):
+    if radius <= 8:
+        coords = np.arange(-8, 8 + 1)
+        ksize = (3, 3)
+    else:
+        coords = np.arange(-radius, radius + 1)
+        ksize = (5, 5)
+    x, y = np.meshgrid(coords, coords)
+    aliased_disk = np.asarray((x**2 + y**2) <= radius**2, dtype=dtype)
+    aliased_disk /= np.sum(aliased_disk)
+
+    return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)
+
+
+def plasma_fractal(mapsize=256, wibbledecay=3, rng=None):
+    """
+    Generate a heightmap using diamond-square algorithm.
+    Return square 2d array, side length 'mapsize', of floats in range 0-255.
+    'mapsize' must be a power of two.
+    """
+    assert mapsize & (mapsize - 1) == 0
+    maparray = np.empty((mapsize, mapsize), dtype=np.float_)
+    maparray[0, 0] = 0
+    stepsize = mapsize
+    wibble = 100
+    if rng is None:
+        rng = np.random.default_rng()
+
+    def wibbledmean(array):
+        return array / 4 + wibble * rng.uniform(-wibble, wibble, array.shape)
+
+    def fillsquares():
+        """For each square of points stepsize apart,
+        calculate middle value as mean of points + wibble"""
+        cornerref = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        squareaccum = cornerref + np.roll(cornerref, shift=-1, axis=0)
+        squareaccum += np.roll(squareaccum, shift=-1, axis=1)
+        maparray[
+            stepsize // 2 : mapsize : stepsize, stepsize // 2 : mapsize : stepsize
+        ] = wibbledmean(squareaccum)
+
+    def filldiamonds():
+        """For each diamond of points stepsize apart,
+        calculate middle value as mean of points + wibble"""
+        drgrid = maparray[
+            stepsize // 2 : mapsize : stepsize, stepsize // 2 : mapsize : stepsize
+        ]
+        ulgrid = maparray[0:mapsize:stepsize, 0:mapsize:stepsize]
+        ldrsum = drgrid + np.roll(drgrid, 1, axis=0)
+        lulsum = ulgrid + np.roll(ulgrid, -1, axis=1)
+        ltsum = ldrsum + lulsum
+        maparray[0:mapsize:stepsize, stepsize // 2 : mapsize : stepsize] = wibbledmean(
+            ltsum
+        )
+        tdrsum = drgrid + np.roll(drgrid, 1, axis=1)
+        tulsum = ulgrid + np.roll(ulgrid, -1, axis=0)
+        ttsum = tdrsum + tulsum
+        maparray[stepsize // 2 : mapsize : stepsize, 0:mapsize:stepsize] = wibbledmean(
+            ttsum
+        )
+
+    while stepsize >= 2:
+        fillsquares()
+        filldiamonds()
+        stepsize //= 2
+        wibble /= wibbledecay
+
+    maparray -= maparray.min()
+    return maparray / maparray.max()
+
+
+class Fog(A.ImageOnlyTransform):
+    def __init__(self, mag=-1, always_apply=False, p=1.0):
+        super().__init__(always_apply=always_apply, p=p)
+        self.rng = np.random.default_rng()
+        self.mag = mag
+
+    def apply(self, img, **params):
+        img = Image.fromarray(img.astype(np.uint8))
+        w, h = img.size
+        c = [(1.5, 2), (2.0, 2), (2.5, 1.7)]
+        if self.mag < 0 or self.mag >= len(c):
+            index = self.rng.integers(0, len(c))
+        else:
+            index = self.mag
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.asarray(img) / 255.0
+        max_val = img.max()
+        max_size = 2 ** math.ceil(math.log2(max(w, h)) + 1)
+        fog = (
+            c[0]
+            * plasma_fractal(mapsize=max_size, wibbledecay=c[1], rng=self.rng)[:h, :w][
+                ..., np.newaxis
+            ]
+        )
+        if isgray:
+            fog = np.squeeze(fog)
+        else:
+            fog = np.repeat(fog, 3, axis=2)
+
+        img += fog
+        img = np.clip(img * max_val / (max_val + c[0]), 0, 1) * 255
+        return img.astype(np.uint8)
+
+
+class Frost(A.ImageOnlyTransform):
+    def __init__(self, mag=-1, always_apply=False, p=1.0):
+        super().__init__(always_apply=always_apply, p=p)
+        self.rng = np.random.default_rng()
+        self.mag = mag
+
+    def apply(self, img, **params):
+        img = Image.fromarray(img.astype(np.uint8))
+        w, h = img.size
+        c = [(0.78, 0.22), (0.64, 0.36), (0.5, 0.5)]
+        if self.mag < 0 or self.mag >= len(c):
+            index = self.rng.integers(0, len(c))
+        else:
+            index = self.mag
+        c = c[index]
+
+        file_dir = os.path.dirname(__file__)
+        filename = [
+            os.path.join(file_dir, "frost_img", "frost1.jpg"),
+            os.path.join(file_dir, "frost_img", "frost2.png"),
+            os.path.join(file_dir, "frost_img", "frost3.png"),
+            os.path.join(file_dir, "frost_img", "frost4.jpg"),
+            os.path.join(file_dir, "frost_img", "frost5.jpg"),
+            os.path.join(file_dir, "frost_img", "frost6.jpg"),
+        ]
+        index = self.rng.integers(0, len(filename))
+        filename = filename[index]
+        frost = Image.open(filename).convert("RGB")
+
+        f_w, f_h = frost.size
+        if w / h > f_w / f_h:
+            f_h = round(f_h * w / f_w)
+            f_w = w
+        else:
+            f_w = round(f_w * h / f_h)
+            f_h = h
+        frost = np.asarray(frost.resize((f_w, f_h)))
+
+        # randomly crop
+        y_start, x_start = self.rng.integers(0, f_h - h + 1), self.rng.integers(
+            0, f_w - w + 1
+        )
+        frost = frost[y_start : y_start + h, x_start : x_start + w]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.asarray(img)
+
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+
+        img = np.clip(np.round(c[0] * img + c[1] * frost), 0, 255)
+        img = img.astype(np.uint8)
+        if isgray:
+            img = np.squeeze(img)
+        return img
+
+
+class Snow(A.ImageOnlyTransform):
+    def __init__(self, mag=-1, always_apply=False, p=1.0):
+        super().__init__(always_apply=always_apply, p=p)
+        self.rng = np.random.default_rng()
+        self.mag = mag
+
+    def apply(self, img, **params):
+        from wand.image import Image as WandImage
+
+        img = Image.fromarray(img.astype(np.uint8))
+        w, h = img.size
+        c = [
+            (0.1, 0.3, 3, 0.5, 10, 4, 0.8),
+            (0.2, 0.3, 2, 0.5, 12, 4, 0.7),
+            (0.55, 0.3, 4, 0.9, 12, 8, 0.7),
+        ]
+        if self.mag < 0 or self.mag >= len(c):
+            index = self.rng.integers(0, len(c))
+        else:
+            index = self.mag
+        c = c[index]
+
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        img = np.asarray(img, dtype=np.float32) / 255.0
+        if isgray:
+            img = np.expand_dims(img, axis=2)
+            img = np.repeat(img, 3, axis=2)
+
+        snow_layer = self.rng.normal(size=img.shape[:2], loc=c[0], scale=c[1])
+
+        snow_layer[snow_layer < c[3]] = 0
+
+        snow_layer = Image.fromarray(
+            (np.clip(snow_layer.squeeze(), 0, 1) * 255).astype(np.uint8), mode="L"
+        )
+        output = BytesIO()
+        snow_layer.save(output, format="PNG")
+        snow_layer = WandImage(blob=output.getvalue())
+
+        snow_layer.motion_blur(
+            radius=c[4], sigma=c[5], angle=self.rng.uniform(-135, -45)
+        )
+
+        snow_layer = (
+            cv2.imdecode(
+                np.frombuffer(snow_layer.make_blob(), np.uint8), cv2.IMREAD_UNCHANGED
+            )
+            / 255.0
+        )
+
+        snow_layer = snow_layer[..., np.newaxis]
+
+        img = c[6] * img
+        gray_img = (1 - c[6]) * np.maximum(
+            img, cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).reshape(h, w, 1) * 1.5 + 0.5
+        )
+        img += gray_img
+        img = np.clip(img + snow_layer + np.rot90(snow_layer, k=2), 0, 1) * 255
+        img = img.astype(np.uint8)
+        if isgray:
+            img = np.squeeze(img)
+        return img
+
+
+class Rain(A.ImageOnlyTransform):
+    def __init__(self, mag=-1, always_apply=False, p=1.0):
+        super().__init__(always_apply=always_apply, p=p)
+        self.rng = np.random.default_rng()
+        self.mag = mag
+
+    def apply(self, img, **params):
+        img = Image.fromarray(img.astype(np.uint8))
+        img = img.copy()
+        w, h = img.size
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+        line_width = self.rng.integers(1, 2)
+
+        c = [50, 70, 90]
+        if self.mag < 0 or self.mag >= len(c):
+            index = 0
+        else:
+            index = self.mag
+        c = c[index]
+
+        n_rains = self.rng.integers(c, c + 20)
+        slant = self.rng.integers(-60, 60)
+        fillcolor = 200 if isgray else (200, 200, 200)
+
+        draw = ImageDraw.Draw(img)
+        max_length = min(w, h, 10)
+        for i in range(1, n_rains):
+            length = self.rng.integers(5, max_length)
+            x1 = self.rng.integers(0, w - length)
+            y1 = self.rng.integers(0, h - length)
+            x2 = x1 + length * math.sin(slant * math.pi / 180.0)
+            y2 = y1 + length * math.cos(slant * math.pi / 180.0)
+            x2 = int(x2)
+            y2 = int(y2)
+            draw.line([(x1, y1), (x2, y2)], width=line_width, fill=fillcolor)
+        img = np.asarray(img).astype(np.uint8)
+        return img
+
+
+class Shadow(A.ImageOnlyTransform):
+    def __init__(self, mag=-1, always_apply=False, p=1.0):
+        super().__init__(always_apply=always_apply, p=p)
+        self.rng = np.random.default_rng()
+        self.mag = mag
+
+    def apply(self, img, **params):
+        img = Image.fromarray(img.astype(np.uint8))
+        w, h = img.size
+        n_channels = len(img.getbands())
+        isgray = n_channels == 1
+
+        c = [64, 96, 128]
+        if self.mag < 0 or self.mag >= len(c):
+            index = 0
+        else:
+            index = self.mag
+        c = c[index]
+
+        img = img.convert("RGBA")
+        overlay = Image.new("RGBA", img.size, (255, 255, 255, 0))
+        draw = ImageDraw.Draw(overlay)
+        transparency = self.rng.integers(c, c + 32)
+        x1 = self.rng.integers(0, w // 2)
+        y1 = 0
+
+        x2 = self.rng.integers(w // 2, w)
+        y2 = 0
+
+        x3 = self.rng.integers(w // 2, w)
+        y3 = h - 1
+
+        x4 = self.rng.integers(0, w // 2)
+        y4 = h - 1
+
+        draw.polygon(
+            [(x1, y1), (x2, y2), (x3, y3), (x4, y4)], fill=(0, 0, 0, transparency)
+        )
+
+        img = Image.alpha_composite(img, overlay)
+        img = img.convert("RGB")
+        if isgray:
+            img = ImageOps.grayscale(img)
+        img = np.asarray(img).astype(np.uint8)
+        return img
+
+
+class UniMERNetTrainTransform:
+    def __init__(self, bitmap_prob=0.04, **kwargs):
+        self.bitmap_prob = bitmap_prob
+        self.train_transform = A.Compose(
+            [
+                A.Compose(
+                    [
+                        Bitmap(p=0.05),
+                        A.OneOf([Fog(), Frost(), Snow(), Rain(), Shadow()], p=0.2),
+                        A.OneOf([Erosion((2, 3)), Dilation((2, 3))], p=0.2),
+                        A.ShiftScaleRotate(
+                            shift_limit=0,
+                            scale_limit=(-0.15, 0),
+                            rotate_limit=1,
+                            border_mode=0,
+                            interpolation=3,
+                            value=[255, 255, 255],
+                            p=1,
+                        ),
+                        A.GridDistortion(
+                            distort_limit=0.1,
+                            border_mode=0,
+                            interpolation=3,
+                            value=[255, 255, 255],
+                            p=0.5,
+                        ),
+                    ],
+                    p=0.15,
+                ),
+                A.RGBShift(r_shift_limit=15, g_shift_limit=15, b_shift_limit=15, p=0.3),
+                A.GaussNoise(10, p=0.2),
+                A.RandomBrightnessContrast(0.05, (-0.2, 0), True, p=0.2),
+                A.ImageCompression(95, p=0.3),
+                A.ToGray(always_apply=True),
+                A.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        if np.random.random() < self.bitmap_prob:
+            img[img != 255] = 0
+        img = self.train_transform(image=img)["image"]
+        data["image"] = img
+        return data
+
+
+class UniMERNetTestTransform:
+    def __init__(self, **kwargs):
+        self.test_transform = A.Compose(
+            [
+                A.ToGray(always_apply=True),
+                A.Normalize((0.7931, 0.7931, 0.7931), (0.1738, 0.1738, 0.1738)),
+            ]
+        )
+
+    def __call__(self, data):
+        img = data["image"]
+        img = self.test_transform(image=img)["image"]
+        data["image"] = img
+        return data
+
+
+class GoTImgDecode:
+    def __init__(self, input_size, random_padding=False, **kwargs):
+        self.input_size = input_size
+        self.random_padding = random_padding
+
+    def crop_margin(self, img):
+        data = np.array(img.convert("L"))
+        data = data.astype(np.uint8)
+        max_val = data.max()
+        min_val = data.min()
+        if max_val == min_val:
+            return img
+        data = (data - min_val) / (max_val - min_val) * 255
+        gray = 255 * (data < 200).astype(np.uint8)
+        coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+        a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+        return img.crop((a, b, w + a, h + b))
+
+    def get_dimensions(self, img):
+        if hasattr(img, "getbands"):
+            channels = len(img.getbands())
+        else:
+            channels = img.channels
+        width, height = img.size
+        return [channels, height, width]
+
+    def _compute_resized_output_size(self, image_size, size, max_size=None):
+        if len(size) == 1:  # specified size only for the smallest edge
+            h, w = image_size
+            short, long = (w, h) if w <= h else (h, w)
+            requested_new_short = size if isinstance(size, int) else size[0]
+
+            new_short, new_long = requested_new_short, int(
+                requested_new_short * long / short
+            )
+
+            if max_size is not None:
+                if max_size <= requested_new_short:
+                    raise ValueError(
+                        f"max_size = {max_size} must be strictly greater than the requested "
+                        f"size for the smaller edge size = {size}"
+                    )
+                if new_long > max_size:
+                    new_short, new_long = int(max_size * new_short / new_long), max_size
+
+            new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
+        else:  # specified both h and w
+            new_w, new_h = size[1], size[0]
+        return [new_h, new_w]
+
+    def resize(self, img, size):
+        _, image_height, image_width = self.get_dimensions(img)
+        if isinstance(size, int):
+            size = [size]
+        max_size = None
+        output_size = self._compute_resized_output_size(
+            (image_height, image_width), size, max_size
+        )
+        img = img.resize(tuple(output_size[::-1]), resample=2)
+        return img
+
+    def __call__(self, data):
+        filename = data["filename"]
+        img = Image.open(filename)
+        try:
+            img = self.crop_margin(img.convert("RGB"))
+        except OSError:
+            return
+        if img.height == 0 or img.width == 0:
+            return
+        img = self.resize(img, min(self.input_size))
+        img.thumbnail((self.input_size[1], self.input_size[0]))
+        delta_width = self.input_size[1] - img.width
+        delta_height = self.input_size[0] - img.height
+        if self.random_padding:
+            pad_width = np.random.randint(low=0, high=delta_width + 1)
+            pad_height = np.random.randint(low=0, high=delta_height + 1)
+        else:
+            pad_width = delta_width // 2
+            pad_height = delta_height // 2
+        padding = (
+            pad_width,
+            pad_height,
+            delta_width - pad_width,
+            delta_height - pad_height,
+        )
+
+        data["image"] = np.array(ImageOps.expand(img, padding))
+        return data
+
+
+class UniMERNetImgDecode:
+    def __init__(
+        self,
+        input_size,
+        random_padding=False,
+        random_resize=False,
+        random_crop=False,
+        **kwargs,
+    ):
+        self.input_size = input_size
+        self.is_random_padding = random_padding
+        self.is_random_resize = random_resize
+        self.is_random_crop = random_crop
+
+    def crop_margin(self, img):
+        data = np.array(img.convert("L"))
+        data = data.astype(np.uint8)
+        max_val = data.max()
+        min_val = data.min()
+        if max_val == min_val:
+            return img
+        data = (data - min_val) / (max_val - min_val) * 255
+        gray = 255 * (data < 200).astype(np.uint8)
+        coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+        a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+        return img.crop((a, b, w + a, h + b))
+
+    def get_dimensions(self, img):
+        if hasattr(img, "getbands"):
+            channels = len(img.getbands())
+        else:
+            channels = img.channels
+        width, height = img.size
+        return [channels, height, width]
+
+    def _compute_resized_output_size(self, image_size, size, max_size=None):
+        if len(size) == 1:  # specified size only for the smallest edge
+            h, w = image_size
+            short, long = (w, h) if w <= h else (h, w)
+            requested_new_short = size if isinstance(size, int) else size[0]
+
+            new_short, new_long = requested_new_short, int(
+                requested_new_short * long / short
+            )
+
+            if max_size is not None:
+                if max_size <= requested_new_short:
+                    raise ValueError(
+                        f"max_size = {max_size} must be strictly greater than the requested "
+                        f"size for the smaller edge size = {size}"
+                    )
+                if new_long > max_size:
+                    new_short, new_long = int(max_size * new_short / new_long), max_size
+
+            new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
+        else:  # specified both h and w
+            new_w, new_h = size[1], size[0]
+        return [new_h, new_w]
+
+    def resize(self, img, size):
+        _, image_height, image_width = self.get_dimensions(img)
+        if isinstance(size, int):
+            size = [size]
+        max_size = None
+        output_size = self._compute_resized_output_size(
+            (image_height, image_width), size, max_size
+        )
+        img = img.resize(tuple(output_size[::-1]), resample=2)
+        return img
+
+    def random_resize(self, img):
+        scale = np.random.uniform(0.5, 1)
+        img = img.resize([int(scale * s) for s in img.size])
+        return img
+
+    def random_crop(self, img, crop_ratio):
+        width, height = img.width, img.height
+        max_crop_pixel = min(width, height) * crop_ratio
+        crop_left = np.random.uniform(0, max_crop_pixel)
+        crop_right = np.random.uniform(0, max_crop_pixel)
+        crop_top = np.random.uniform(0, max_crop_pixel)
+        crop_bottom = np.random.uniform(0, max_crop_pixel)
+        # 计算裁剪后的边界
+        left = crop_left
+        top = crop_top
+        right = width - crop_right
+        bottom = height - crop_bottom
+        # 裁剪图像
+        img = img.crop((left, top, right, bottom))
+
+        return img
+
+    def __call__(self, data):
+        filename = data["filename"]
+        img = Image.open(filename)
+        try:
+            if self.is_random_resize:
+                img = self.random_resize(img)
+            img = self.crop_margin(img.convert("RGB"))
+            if "label" in data and self.is_random_crop:
+                label = data["label"]
+                equation_length = len(label)
+                if equation_length < 256:
+                    img = self.random_crop(img, crop_ratio=0.1)
+                elif 256 < equation_length <= 512:
+                    img = self.random_crop(img, crop_ratio=0.05)
+                else:
+                    img = self.random_crop(img, crop_ratio=0.03)
+        except OSError:
+            return
+        if img.height == 0 or img.width == 0:
+            return
+        img = self.resize(img, min(self.input_size))
+        img.thumbnail((self.input_size[1], self.input_size[0]))
+        delta_width = self.input_size[1] - img.width
+        delta_height = self.input_size[0] - img.height
+        if self.is_random_padding:
+            pad_width = np.random.randint(low=0, high=delta_width + 1)
+            pad_height = np.random.randint(low=0, high=delta_height + 1)
+        else:
+            pad_width = delta_width // 2
+            pad_height = delta_height // 2
+        padding = (
+            pad_width,
+            pad_height,
+            delta_width - pad_width,
+            delta_height - pad_height,
+        )
+
+        data["image"] = np.array(ImageOps.expand(img, padding))
+        return data
+
+
+class UniMERNetResize:
+    def __init__(self, input_size, random_padding=False, **kwargs):
+        self.input_size = input_size
+        self.random_padding = random_padding
+
+    def crop_margin(self, img):
+        data = np.array(img.convert("L"))
+        data = data.astype(np.uint8)
+        max_val = data.max()
+        min_val = data.min()
+        if max_val == min_val:
+            return img
+        data = (data - min_val) / (max_val - min_val) * 255
+        gray = 255 * (data < 200).astype(np.uint8)
+
+        coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+        a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+        return img.crop((a, b, w + a, h + b))
+
+    def get_dimensions(self, img):
+        if hasattr(img, "getbands"):
+            channels = len(img.getbands())
+        else:
+            channels = img.channels
+        width, height = img.size
+        return [channels, height, width]
+
+    def _compute_resized_output_size(self, image_size, size, max_size=None):
+        if len(size) == 1:  # specified size only for the smallest edge
+            h, w = image_size
+            short, long = (w, h) if w <= h else (h, w)
+            requested_new_short = size if isinstance(size, int) else size[0]
+
+            new_short, new_long = requested_new_short, int(
+                requested_new_short * long / short
+            )
+
+            if max_size is not None:
+                if max_size <= requested_new_short:
+                    raise ValueError(
+                        f"max_size = {max_size} must be strictly greater than the requested "
+                        f"size for the smaller edge size = {size}"
+                    )
+                if new_long > max_size:
+                    new_short, new_long = int(max_size * new_short / new_long), max_size
+
+            new_w, new_h = (new_short, new_long) if w <= h else (new_long, new_short)
+        else:  # specified both h and w
+            new_w, new_h = size[1], size[0]
+        return [new_h, new_w]
+
+    def resize(self, img, size):
+        _, image_height, image_width = self.get_dimensions(img)
+        if isinstance(size, int):
+            size = [size]
+        max_size = None
+        output_size = self._compute_resized_output_size(
+            (image_height, image_width), size, max_size
+        )
+        img.resize(tuple(output_size[::-1]), resample=2)
+        return img
+
+    def __call__(self, data):
+        img = data["image"]
+        img = Image.fromarray(img)
+        try:
+            img = self.crop_margin(img)
+        except OSError:
+            return
+        if img.height == 0 or img.width == 0:
+            return
+        img = self.resize(img, min(self.input_size))
+        img.thumbnail((self.input_size[1], self.input_size[0]))
+        delta_width = self.input_size[1] - img.width
+        delta_height = self.input_size[0] - img.height
+        if self.random_padding:
+            pad_width = np.random.randint(low=0, high=delta_width + 1)
+            pad_height = np.random.randint(low=0, high=delta_height + 1)
+        else:
+            pad_width = delta_width // 2
+            pad_height = delta_height // 2
+        padding = (
+            pad_width,
+            pad_height,
+            delta_width - pad_width,
+            delta_height - pad_height,
+        )
+        data["image"] = np.array(ImageOps.expand(img, padding))
+        return data
+
+
+class UniMERNetImageFormat:
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        img = data["image"]
+        im_h, im_w = img.shape[:2]
+        divide_h = math.ceil(im_h / 32) * 32
+        divide_w = math.ceil(im_w / 32) * 32
+        img = img[:, :, 0]
+        img = np.pad(
+            img, ((0, divide_h - im_h), (0, divide_w - im_w)), constant_values=(1, 1)
+        )
+        img_expanded = img[:, :, np.newaxis].transpose(2, 0, 1)
+        data["image"] = img_expanded
+        return data

ppocr/data/imaug/vqa/__init__.py

@@ -0,0 +1,29 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 .token import (
+    VQATokenPad,
+    VQASerTokenChunk,
+    VQAReTokenChunk,
+    VQAReTokenRelation,
+    TensorizeEntitiesRelations,
+)
+
+__all__ = [
+    "VQATokenPad",
+    "VQASerTokenChunk",
+    "VQAReTokenChunk",
+    "VQAReTokenRelation",
+    "TensorizeEntitiesRelations",
+]

ppocr/data/imaug/vqa/augment.py

@@ -0,0 +1,34 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import os
+import sys
+import numpy as np
+import random
+from copy import deepcopy
+
+
+def order_by_tbyx(ocr_info):
+    res = sorted(ocr_info, key=lambda r: (r["bbox"][1], r["bbox"][0]))
+    for i in range(len(res) - 1):
+        for j in range(i, 0, -1):
+            if abs(res[j + 1]["bbox"][1] - res[j]["bbox"][1]) < 20 and (
+                res[j + 1]["bbox"][0] < res[j]["bbox"][0]
+            ):
+                tmp = deepcopy(res[j])
+                res[j] = deepcopy(res[j + 1])
+                res[j + 1] = deepcopy(tmp)
+            else:
+                break
+    return res

ppocr/data/imaug/vqa/token/__init__.py

@@ -0,0 +1,18 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 .vqa_token_chunk import VQASerTokenChunk, VQAReTokenChunk
+from .vqa_token_pad import VQATokenPad
+from .vqa_token_relation import VQAReTokenRelation
+from .vqa_re_convert import TensorizeEntitiesRelations

ppocr/data/imaug/vqa/token/vqa_re_convert.py

@@ -0,0 +1,49 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import numpy as np
+
+
+class TensorizeEntitiesRelations(object):
+    def __init__(self, max_seq_len=512, infer_mode=False, **kwargs):
+        self.max_seq_len = max_seq_len
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        entities = data["entities"]
+        relations = data["relations"]
+
+        entities_new = np.full(
+            shape=[self.max_seq_len + 1, 3], fill_value=-1, dtype="int64"
+        )
+        entities_new[0, 0] = len(entities["start"])
+        entities_new[0, 1] = len(entities["end"])
+        entities_new[0, 2] = len(entities["label"])
+        entities_new[1 : len(entities["start"]) + 1, 0] = np.array(entities["start"])
+        entities_new[1 : len(entities["end"]) + 1, 1] = np.array(entities["end"])
+        entities_new[1 : len(entities["label"]) + 1, 2] = np.array(entities["label"])
+
+        relations_new = np.full(
+            shape=[self.max_seq_len * self.max_seq_len + 1, 2],
+            fill_value=-1,
+            dtype="int64",
+        )
+        relations_new[0, 0] = len(relations["head"])
+        relations_new[0, 1] = len(relations["tail"])
+        relations_new[1 : len(relations["head"]) + 1, 0] = np.array(relations["head"])
+        relations_new[1 : len(relations["tail"]) + 1, 1] = np.array(relations["tail"])
+
+        data["entities"] = entities_new
+        data["relations"] = relations_new
+        return data

ppocr/data/imaug/vqa/token/vqa_token_chunk.py

@@ -0,0 +1,134 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 collections import defaultdict
+
+
+class VQASerTokenChunk(object):
+    def __init__(self, max_seq_len=512, infer_mode=False, **kwargs):
+        self.max_seq_len = max_seq_len
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        encoded_inputs_all = []
+        seq_len = len(data["input_ids"])
+        for index in range(0, seq_len, self.max_seq_len):
+            chunk_beg = index
+            chunk_end = min(index + self.max_seq_len, seq_len)
+            encoded_inputs_example = {}
+            for key in data:
+                if key in [
+                    "label",
+                    "input_ids",
+                    "labels",
+                    "token_type_ids",
+                    "bbox",
+                    "attention_mask",
+                ]:
+                    if self.infer_mode and key == "labels":
+                        encoded_inputs_example[key] = data[key]
+                    else:
+                        encoded_inputs_example[key] = data[key][chunk_beg:chunk_end]
+                else:
+                    encoded_inputs_example[key] = data[key]
+
+            encoded_inputs_all.append(encoded_inputs_example)
+        if len(encoded_inputs_all) == 0:
+            return None
+        return encoded_inputs_all[0]
+
+
+class VQAReTokenChunk(object):
+    def __init__(
+        self, max_seq_len=512, entities_labels=None, infer_mode=False, **kwargs
+    ):
+        self.max_seq_len = max_seq_len
+        self.entities_labels = (
+            {"HEADER": 0, "QUESTION": 1, "ANSWER": 2}
+            if entities_labels is None
+            else entities_labels
+        )
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        # prepare data
+        entities = data.pop("entities")
+        relations = data.pop("relations")
+        encoded_inputs_all = []
+        for index in range(0, len(data["input_ids"]), self.max_seq_len):
+            item = {}
+            for key in data:
+                if key in [
+                    "label",
+                    "input_ids",
+                    "labels",
+                    "token_type_ids",
+                    "bbox",
+                    "attention_mask",
+                ]:
+                    if self.infer_mode and key == "labels":
+                        item[key] = data[key]
+                    else:
+                        item[key] = data[key][index : index + self.max_seq_len]
+                else:
+                    item[key] = data[key]
+            # select entity in current chunk
+            entities_in_this_span = []
+            global_to_local_map = {}  #
+            for entity_id, entity in enumerate(entities):
+                if (
+                    index <= entity["start"] < index + self.max_seq_len
+                    and index <= entity["end"] < index + self.max_seq_len
+                ):
+                    entity["start"] = entity["start"] - index
+                    entity["end"] = entity["end"] - index
+                    global_to_local_map[entity_id] = len(entities_in_this_span)
+                    entities_in_this_span.append(entity)
+
+            # select relations in current chunk
+            relations_in_this_span = []
+            for relation in relations:
+                if (
+                    index <= relation["start_index"] < index + self.max_seq_len
+                    and index <= relation["end_index"] < index + self.max_seq_len
+                ):
+                    relations_in_this_span.append(
+                        {
+                            "head": global_to_local_map[relation["head"]],
+                            "tail": global_to_local_map[relation["tail"]],
+                            "start_index": relation["start_index"] - index,
+                            "end_index": relation["end_index"] - index,
+                        }
+                    )
+            item.update(
+                {
+                    "entities": self.reformat(entities_in_this_span),
+                    "relations": self.reformat(relations_in_this_span),
+                }
+            )
+            if len(item["entities"]) > 0:
+                item["entities"]["label"] = [
+                    self.entities_labels[x] for x in item["entities"]["label"]
+                ]
+                encoded_inputs_all.append(item)
+        if len(encoded_inputs_all) == 0:
+            return None
+        return encoded_inputs_all[0]
+
+    def reformat(self, data):
+        new_data = defaultdict(list)
+        for item in data:
+            for k, v in item.items():
+                new_data[k].append(v)
+        return new_data

ppocr/data/imaug/vqa/token/vqa_token_pad.py

@@ -0,0 +1,117 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import paddle
+import numpy as np
+
+
+class VQATokenPad(object):
+    def __init__(
+        self,
+        max_seq_len=512,
+        pad_to_max_seq_len=True,
+        return_attention_mask=True,
+        return_token_type_ids=True,
+        truncation_strategy="longest_first",
+        return_overflowing_tokens=False,
+        return_special_tokens_mask=False,
+        infer_mode=False,
+        **kwargs,
+    ):
+        self.max_seq_len = max_seq_len
+        self.pad_to_max_seq_len = max_seq_len
+        self.return_attention_mask = return_attention_mask
+        self.return_token_type_ids = return_token_type_ids
+        self.truncation_strategy = truncation_strategy
+        self.return_overflowing_tokens = return_overflowing_tokens
+        self.return_special_tokens_mask = return_special_tokens_mask
+        self.pad_token_label_id = paddle.nn.CrossEntropyLoss().ignore_index
+        self.infer_mode = infer_mode
+
+    def __call__(self, data):
+        needs_to_be_padded = (
+            self.pad_to_max_seq_len and len(data["input_ids"]) < self.max_seq_len
+        )
+
+        if needs_to_be_padded:
+            if "tokenizer_params" in data:
+                tokenizer_params = data.pop("tokenizer_params")
+            else:
+                tokenizer_params = dict(
+                    padding_side="right", pad_token_type_id=0, pad_token_id=1
+                )
+
+            difference = self.max_seq_len - len(data["input_ids"])
+            if tokenizer_params["padding_side"] == "right":
+                if self.return_attention_mask:
+                    data["attention_mask"] = [1] * len(data["input_ids"]) + [
+                        0
+                    ] * difference
+                if self.return_token_type_ids:
+                    data["token_type_ids"] = (
+                        data["token_type_ids"]
+                        + [tokenizer_params["pad_token_type_id"]] * difference
+                    )
+                if self.return_special_tokens_mask:
+                    data["special_tokens_mask"] = (
+                        data["special_tokens_mask"] + [1] * difference
+                    )
+                data["input_ids"] = (
+                    data["input_ids"] + [tokenizer_params["pad_token_id"]] * difference
+                )
+                if not self.infer_mode:
+                    data["labels"] = (
+                        data["labels"] + [self.pad_token_label_id] * difference
+                    )
+                data["bbox"] = data["bbox"] + [[0, 0, 0, 0]] * difference
+            elif tokenizer_params["padding_side"] == "left":
+                if self.return_attention_mask:
+                    data["attention_mask"] = [0] * difference + [1] * len(
+                        data["input_ids"]
+                    )
+                if self.return_token_type_ids:
+                    data["token_type_ids"] = [
+                        tokenizer_params["pad_token_type_id"]
+                    ] * difference + data["token_type_ids"]
+                if self.return_special_tokens_mask:
+                    data["special_tokens_mask"] = [1] * difference + data[
+                        "special_tokens_mask"
+                    ]
+                data["input_ids"] = [
+                    tokenizer_params["pad_token_id"]
+                ] * difference + data["input_ids"]
+                if not self.infer_mode:
+                    data["labels"] = [self.pad_token_label_id] * difference + data[
+                        "labels"
+                    ]
+                data["bbox"] = [[0, 0, 0, 0]] * difference + data["bbox"]
+        else:
+            if self.return_attention_mask:
+                data["attention_mask"] = [1] * len(data["input_ids"])
+
+        for key in data:
+            if key in [
+                "input_ids",
+                "labels",
+                "token_type_ids",
+                "bbox",
+                "attention_mask",
+            ]:
+                if self.infer_mode:
+                    if key != "labels":
+                        length = min(len(data[key]), self.max_seq_len)
+                        data[key] = data[key][:length]
+                    else:
+                        continue
+                data[key] = np.array(data[key], dtype="int64")
+        return data

ppocr/data/imaug/vqa/token/vqa_token_relation.py

@@ -0,0 +1,76 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+
+class VQAReTokenRelation(object):
+    def __init__(self, **kwargs):
+        pass
+
+    def __call__(self, data):
+        """
+        build relations
+        """
+        entities = data["entities"]
+        relations = data["relations"]
+        id2label = data.pop("id2label")
+        empty_entity = data.pop("empty_entity")
+        entity_id_to_index_map = data.pop("entity_id_to_index_map")
+
+        relations = list(set(relations))
+        relations = [
+            rel
+            for rel in relations
+            if rel[0] not in empty_entity and rel[1] not in empty_entity
+        ]
+        kv_relations = []
+        for rel in relations:
+            pair = [id2label[rel[0]], id2label[rel[1]]]
+            if pair == ["question", "answer"]:
+                kv_relations.append(
+                    {
+                        "head": entity_id_to_index_map[rel[0]],
+                        "tail": entity_id_to_index_map[rel[1]],
+                    }
+                )
+            elif pair == ["answer", "question"]:
+                kv_relations.append(
+                    {
+                        "head": entity_id_to_index_map[rel[1]],
+                        "tail": entity_id_to_index_map[rel[0]],
+                    }
+                )
+            else:
+                continue
+        relations = sorted(
+            [
+                {
+                    "head": rel["head"],
+                    "tail": rel["tail"],
+                    "start_index": self.get_relation_span(rel, entities)[0],
+                    "end_index": self.get_relation_span(rel, entities)[1],
+                }
+                for rel in kv_relations
+            ],
+            key=lambda x: x["head"],
+        )
+
+        data["relations"] = relations
+        return data
+
+    def get_relation_span(self, rel, entities):
+        bound = []
+        for entity_index in [rel["head"], rel["tail"]]:
+            bound.append(entities[entity_index]["start"])
+            bound.append(entities[entity_index]["end"])
+        return min(bound), max(bound)

ppocr/data/latexocr_dataset.py

@@ -0,0 +1,174 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+"""
+This code is refer from:
+https://github.com/lukas-blecher/LaTeX-OCR/blob/main/pix2tex/dataset/dataset.py
+"""
+
+import numpy as np
+import cv2
+import math
+import os
+import json
+import pickle
+import random
+import traceback
+import paddle
+from paddle.io import Dataset
+from .imaug.label_ops import LatexOCRLabelEncode
+from .imaug import transform, create_operators
+
+
+class LaTeXOCRDataSet(Dataset):
+    def __init__(self, config, mode, logger, seed=None):
+        super(LaTeXOCRDataSet, self).__init__()
+        self.logger = logger
+        self.mode = mode.lower()
+
+        global_config = config["Global"]
+        dataset_config = config[mode]["dataset"]
+        loader_config = config[mode]["loader"]
+
+        pkl_path = dataset_config.pop("data")
+        self.data_dir = dataset_config["data_dir"]
+        self.min_dimensions = dataset_config.pop("min_dimensions")
+        self.max_dimensions = dataset_config.pop("max_dimensions")
+        self.batchsize = dataset_config.pop("batch_size_per_pair")
+        self.keep_smaller_batches = dataset_config.pop("keep_smaller_batches")
+        self.max_seq_len = global_config.pop("max_seq_len")
+        self.rec_char_dict_path = global_config.pop("rec_char_dict_path")
+        self.tokenizer = LatexOCRLabelEncode(self.rec_char_dict_path)
+
+        file = open(pkl_path, "rb")
+        data = pickle.load(file)
+        temp = {}
+        for k in data:
+            if (
+                self.min_dimensions[0] <= k[0] <= self.max_dimensions[0]
+                and self.min_dimensions[1] <= k[1] <= self.max_dimensions[1]
+            ):
+                temp[k] = data[k]
+        self.data = temp
+        self.do_shuffle = loader_config["shuffle"]
+        self.seed = seed
+
+        if self.mode == "train" and self.do_shuffle:
+            random.seed(self.seed)
+        self.pairs = []
+        for k in self.data:
+            info = np.array(self.data[k], dtype=object)
+            p = (
+                paddle.randperm(len(info))
+                if self.mode == "train" and self.do_shuffle
+                else paddle.arange(len(info))
+            )
+            for i in range(0, len(info), self.batchsize):
+                batch = info[p[i : i + self.batchsize]]
+                if len(batch.shape) == 1:
+                    batch = batch[None, :]
+                if len(batch) < self.batchsize and not self.keep_smaller_batches:
+                    continue
+                self.pairs.append(batch)
+        if self.do_shuffle:
+            self.pairs = np.random.permutation(np.array(self.pairs, dtype=object))
+        else:
+            self.pairs = np.array(self.pairs, dtype=object)
+
+        self.size = len(self.pairs)
+        self.set_epoch_as_seed(self.seed, dataset_config)
+
+        self.ops = create_operators(dataset_config["transforms"], global_config)
+        self.ext_op_transform_idx = dataset_config.get("ext_op_transform_idx", 2)
+        self.need_reset = True
+
+    def set_epoch_as_seed(self, seed, dataset_config):
+        if self.mode == "train":
+            try:
+                border_map_id = [
+                    index
+                    for index, dictionary in enumerate(dataset_config["transforms"])
+                    if "MakeBorderMap" in dictionary
+                ][0]
+                shrink_map_id = [
+                    index
+                    for index, dictionary in enumerate(dataset_config["transforms"])
+                    if "MakeShrinkMap" in dictionary
+                ][0]
+                dataset_config["transforms"][border_map_id]["MakeBorderMap"][
+                    "epoch"
+                ] = (seed if seed is not None else 0)
+                dataset_config["transforms"][shrink_map_id]["MakeShrinkMap"][
+                    "epoch"
+                ] = (seed if seed is not None else 0)
+            except Exception as E:
+                print(E)
+                return
+
+    def shuffle_data_random(self):
+        random.seed(self.seed)
+        random.shuffle(self.data_lines)
+        return
+
+    def __getitem__(self, idx):
+        batch = self.pairs[idx]
+        eqs, ims = batch.T
+        try:
+            max_width, max_height, max_length = 0, 0, 0
+
+            images_transform = []
+
+            for file_name in ims:
+                img_path = os.path.join(self.data_dir, file_name)
+                data = {
+                    "img_path": img_path,
+                }
+                with open(data["img_path"], "rb") as f:
+                    img = f.read()
+                    data["image"] = img
+                    item = transform(data, self.ops)
+                    images_transform.append(np.array(item[0]))
+            image_concat = np.concatenate(images_transform, axis=0)[:, np.newaxis, :, :]
+            images_transform = image_concat.astype(np.float32)
+            labels, attention_mask, max_length = self.tokenizer(list(eqs))
+            if self.max_seq_len < max_length:
+                rnd_idx = (
+                    np.random.randint(self.__len__())
+                    if self.mode == "train"
+                    else (idx + 1) % self.__len__()
+                )
+                return self.__getitem__(rnd_idx)
+            return (images_transform, labels, attention_mask)
+
+        except:
+
+            self.logger.error(
+                "When parsing line {}, error happened with msg: {}".format(
+                    data["img_path"], traceback.format_exc()
+                )
+            )
+            outs = None
+
+        if outs is None:
+            # during evaluation, we should fix the idx to get same results for many times of evaluation.
+            rnd_idx = (
+                np.random.randint(self.__len__())
+                if self.mode == "train"
+                else (idx + 1) % self.__len__()
+            )
+            return self.__getitem__(rnd_idx)
+        return outs
+
+    def __len__(self):
+        return self.size

ppocr/data/lmdb_dataset.py

@@ -0,0 +1,301 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import numpy as np
+import io
+import os
+from paddle.io import Dataset
+import lmdb
+import cv2
+import string
+import pickle
+from PIL import Image
+
+from .imaug import transform, create_operators
+
+
+class LMDBDataSet(Dataset):
+    def __init__(self, config, mode, logger, seed=None):
+        super(LMDBDataSet, self).__init__()
+
+        global_config = config["Global"]
+        dataset_config = config[mode]["dataset"]
+        loader_config = config[mode]["loader"]
+        batch_size = loader_config["batch_size_per_card"]
+        data_dir = dataset_config["data_dir"]
+        self.do_shuffle = loader_config["shuffle"]
+
+        self.lmdb_sets = self.load_hierarchical_lmdb_dataset(data_dir)
+        logger.info("Initialize indexes of datasets:%s" % data_dir)
+        self.data_idx_order_list = self.dataset_traversal()
+        if self.do_shuffle:
+            np.random.shuffle(self.data_idx_order_list)
+        self.ops = create_operators(dataset_config["transforms"], global_config)
+        self.ext_op_transform_idx = dataset_config.get("ext_op_transform_idx", 1)
+
+        ratio_list = dataset_config.get("ratio_list", [1.0])
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def load_hierarchical_lmdb_dataset(self, data_dir):
+        lmdb_sets = {}
+        dataset_idx = 0
+        for dirpath, dirnames, filenames in os.walk(data_dir + "/"):
+            if not dirnames:
+                env = lmdb.open(
+                    dirpath,
+                    max_readers=32,
+                    readonly=True,
+                    lock=False,
+                    readahead=False,
+                    meminit=False,
+                )
+                txn = env.begin(write=False)
+                num_samples = int(txn.get("num-samples".encode()))
+                lmdb_sets[dataset_idx] = {
+                    "dirpath": dirpath,
+                    "env": env,
+                    "txn": txn,
+                    "num_samples": num_samples,
+                }
+                dataset_idx += 1
+        return lmdb_sets
+
+    def dataset_traversal(self):
+        lmdb_num = len(self.lmdb_sets)
+        total_sample_num = 0
+        for lno in range(lmdb_num):
+            total_sample_num += self.lmdb_sets[lno]["num_samples"]
+        data_idx_order_list = np.zeros((total_sample_num, 2))
+        beg_idx = 0
+        for lno in range(lmdb_num):
+            tmp_sample_num = self.lmdb_sets[lno]["num_samples"]
+            end_idx = beg_idx + tmp_sample_num
+            data_idx_order_list[beg_idx:end_idx, 0] = lno
+            data_idx_order_list[beg_idx:end_idx, 1] = list(range(tmp_sample_num))
+            data_idx_order_list[beg_idx:end_idx, 1] += 1
+            beg_idx = beg_idx + tmp_sample_num
+        return data_idx_order_list
+
+    def get_img_data(self, value):
+        """get_img_data"""
+        if not value:
+            return None
+        imgdata = np.frombuffer(value, dtype="uint8")
+        if imgdata is None:
+            return None
+        imgori = cv2.imdecode(imgdata, 1)
+        if imgori is None:
+            return None
+        return imgori
+
+    def get_ext_data(self):
+        ext_data_num = 0
+        for op in self.ops:
+            if hasattr(op, "ext_data_num"):
+                ext_data_num = getattr(op, "ext_data_num")
+                break
+        load_data_ops = self.ops[: self.ext_op_transform_idx]
+        ext_data = []
+
+        while len(ext_data) < ext_data_num:
+            lmdb_idx, file_idx = self.data_idx_order_list[np.random.randint(len(self))]
+            lmdb_idx = int(lmdb_idx)
+            file_idx = int(file_idx)
+            sample_info = self.get_lmdb_sample_info(
+                self.lmdb_sets[lmdb_idx]["txn"], file_idx
+            )
+            if sample_info is None:
+                continue
+            img, label = sample_info
+            data = {"image": img, "label": label}
+            data = transform(data, load_data_ops)
+            if data is None:
+                continue
+            ext_data.append(data)
+        return ext_data
+
+    def get_lmdb_sample_info(self, txn, index):
+        label_key = "label-%09d".encode() % index
+        label = txn.get(label_key)
+        if label is None:
+            return None
+        label = label.decode("utf-8")
+        img_key = "image-%09d".encode() % index
+        imgbuf = txn.get(img_key)
+        return imgbuf, label
+
+    def __getitem__(self, idx):
+        lmdb_idx, file_idx = self.data_idx_order_list[idx]
+        lmdb_idx = int(lmdb_idx)
+        file_idx = int(file_idx)
+        sample_info = self.get_lmdb_sample_info(
+            self.lmdb_sets[lmdb_idx]["txn"], file_idx
+        )
+        if sample_info is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        img, label = sample_info
+        data = {"image": img, "label": label}
+        data["ext_data"] = self.get_ext_data()
+        outs = transform(data, self.ops)
+        if outs is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        return outs
+
+    def __len__(self):
+        return self.data_idx_order_list.shape[0]
+
+
+class LMDBDataSetSR(LMDBDataSet):
+    def buf2PIL(self, txn, key, type="RGB"):
+        imgbuf = txn.get(key)
+        buf = io.BytesIO()
+        buf.write(imgbuf)
+        buf.seek(0)
+        im = Image.open(buf).convert(type)
+        return im
+
+    def str_filt(self, str_, voc_type):
+        alpha_dict = {
+            "digit": string.digits,
+            "lower": string.digits + string.ascii_lowercase,
+            "upper": string.digits + string.ascii_letters,
+            "all": string.digits + string.ascii_letters + string.punctuation,
+        }
+        if voc_type == "lower":
+            str_ = str_.lower()
+        for char in str_:
+            if char not in alpha_dict[voc_type]:
+                str_ = str_.replace(char, "")
+        return str_
+
+    def get_lmdb_sample_info(self, txn, index):
+        self.voc_type = "upper"
+        self.max_len = 100
+        self.test = False
+        label_key = b"label-%09d" % index
+        word = str(txn.get(label_key).decode())
+        img_HR_key = b"image_hr-%09d" % index  # 128*32
+        img_lr_key = b"image_lr-%09d" % index  # 64*16
+        try:
+            img_HR = self.buf2PIL(txn, img_HR_key, "RGB")
+            img_lr = self.buf2PIL(txn, img_lr_key, "RGB")
+        except IOError or len(word) > self.max_len:
+            return self[index + 1]
+        label_str = self.str_filt(word, self.voc_type)
+        return img_HR, img_lr, label_str
+
+    def __getitem__(self, idx):
+        lmdb_idx, file_idx = self.data_idx_order_list[idx]
+        lmdb_idx = int(lmdb_idx)
+        file_idx = int(file_idx)
+        sample_info = self.get_lmdb_sample_info(
+            self.lmdb_sets[lmdb_idx]["txn"], file_idx
+        )
+        if sample_info is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        img_HR, img_lr, label_str = sample_info
+        data = {"image_hr": img_HR, "image_lr": img_lr, "label": label_str}
+        outs = transform(data, self.ops)
+        if outs is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        return outs
+
+
+class LMDBDataSetTableMaster(LMDBDataSet):
+    def load_hierarchical_lmdb_dataset(self, data_dir):
+        lmdb_sets = {}
+        dataset_idx = 0
+        env = lmdb.open(
+            data_dir,
+            max_readers=32,
+            readonly=True,
+            lock=False,
+            readahead=False,
+            meminit=False,
+        )
+        txn = env.begin(write=False)
+        num_samples = int(pickle.loads(txn.get(b"__len__")))
+        lmdb_sets[dataset_idx] = {
+            "dirpath": data_dir,
+            "env": env,
+            "txn": txn,
+            "num_samples": num_samples,
+        }
+        return lmdb_sets
+
+    def get_img_data(self, value):
+        """get_img_data"""
+        if not value:
+            return None
+        imgdata = np.frombuffer(value, dtype="uint8")
+        if imgdata is None:
+            return None
+        imgori = cv2.imdecode(imgdata, 1)
+        if imgori is None:
+            return None
+        return imgori
+
+    def get_lmdb_sample_info(self, txn, index):
+        def convert_bbox(bbox_str_list):
+            bbox_list = []
+            for bbox_str in bbox_str_list:
+                bbox_list.append(int(bbox_str))
+            return bbox_list
+
+        try:
+            data = pickle.loads(txn.get(str(index).encode("utf8")))
+        except:
+            return None
+
+        # img_name, img, info_lines
+        file_name = data[0]
+        bytes = data[1]
+        info_lines = data[2]  # raw data from TableMASTER annotation file.
+        # parse info_lines
+        raw_data = info_lines.strip().split("\n")
+        raw_name, text = (
+            raw_data[0],
+            raw_data[1],
+        )  # don't filter the samples's length over max_seq_len.
+        text = text.split(",")
+        bbox_str_list = raw_data[2:]
+        bbox_split = ","
+        bboxes = [
+            {"bbox": convert_bbox(bsl.strip().split(bbox_split)), "tokens": ["1", "2"]}
+            for bsl in bbox_str_list
+        ]
+
+        # advance parse bbox
+        # import pdb;pdb.set_trace()
+
+        line_info = {}
+        line_info["file_name"] = file_name
+        line_info["structure"] = text
+        line_info["cells"] = bboxes
+        line_info["image"] = bytes
+        return line_info
+
+    def __getitem__(self, idx):
+        lmdb_idx, file_idx = self.data_idx_order_list[idx]
+        lmdb_idx = int(lmdb_idx)
+        file_idx = int(file_idx)
+        data = self.get_lmdb_sample_info(self.lmdb_sets[lmdb_idx]["txn"], file_idx)
+        if data is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        outs = transform(data, self.ops)
+        if outs is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        return outs
+
+    def __len__(self):
+        return self.data_idx_order_list.shape[0]

ppocr/data/multi_scale_sampler.py

@@ -0,0 +1,171 @@
+from paddle.io import Sampler
+import paddle.distributed as dist
+
+import numpy as np
+import random
+import math
+
+
+class MultiScaleSampler(Sampler):
+    def __init__(
+        self,
+        data_source,
+        scales,
+        first_bs=128,
+        fix_bs=True,
+        divided_factor=[8, 16],
+        is_training=True,
+        ratio_wh=0.8,
+        max_w=480.0,
+        seed=None,
+    ):
+        """
+        multi scale samper
+        Args:
+            data_source(dataset)
+            scales(list): several scales for image resolution
+            first_bs(int): batch size for the first scale in scales
+            divided_factor(list[w, h]): ImageNet models down-sample images by a factor, ensure that width and height dimensions are multiples are multiple of devided_factor.
+            is_training(boolean): mode
+        """
+        # min. and max. spatial dimensions
+        self.data_source = data_source
+        self.data_idx_order_list = np.array(data_source.data_idx_order_list)
+        self.ds_width = data_source.ds_width
+        self.seed = data_source.seed
+        if self.ds_width:
+            self.wh_ratio = data_source.wh_ratio
+            self.wh_ratio_sort = data_source.wh_ratio_sort
+        self.n_data_samples = len(self.data_source)
+        self.ratio_wh = ratio_wh
+        self.max_w = max_w
+
+        if isinstance(scales[0], list):
+            width_dims = [i[0] for i in scales]
+            height_dims = [i[1] for i in scales]
+        elif isinstance(scales[0], int):
+            width_dims = scales
+            height_dims = scales
+        base_im_w = width_dims[0]
+        base_im_h = height_dims[0]
+        base_batch_size = first_bs
+
+        # Get the GPU and node related information
+        num_replicas = dist.get_world_size()
+        rank = dist.get_rank()
+        # adjust the total samples to avoid batch dropping
+        num_samples_per_replica = int(self.n_data_samples * 1.0 / num_replicas)
+
+        img_indices = [idx for idx in range(self.n_data_samples)]
+
+        self.shuffle = False
+        if is_training:
+            # compute the spatial dimensions and corresponding batch size
+            # ImageNet models down-sample images by a factor of 32.
+            # Ensure that width and height dimensions are multiples are multiple of 32.
+            width_dims = [
+                int((w // divided_factor[0]) * divided_factor[0]) for w in width_dims
+            ]
+            height_dims = [
+                int((h // divided_factor[1]) * divided_factor[1]) for h in height_dims
+            ]
+
+            img_batch_pairs = list()
+            base_elements = base_im_w * base_im_h * base_batch_size
+            for h, w in zip(height_dims, width_dims):
+                if fix_bs:
+                    batch_size = base_batch_size
+                else:
+                    batch_size = int(max(1, (base_elements / (h * w))))
+                img_batch_pairs.append((w, h, batch_size))
+            self.img_batch_pairs = img_batch_pairs
+            self.shuffle = True
+        else:
+            self.img_batch_pairs = [(base_im_w, base_im_h, base_batch_size)]
+
+        self.img_indices = img_indices
+        self.n_samples_per_replica = num_samples_per_replica
+        self.epoch = 0
+        self.rank = rank
+        self.num_replicas = num_replicas
+
+        self.batch_list = []
+        self.current = 0
+        last_index = num_samples_per_replica * num_replicas
+        indices_rank_i = self.img_indices[self.rank : last_index : self.num_replicas]
+        while self.current < self.n_samples_per_replica:
+            for curr_w, curr_h, curr_bsz in self.img_batch_pairs:
+                end_index = min(self.current + curr_bsz, self.n_samples_per_replica)
+                batch_ids = indices_rank_i[self.current : end_index]
+                n_batch_samples = len(batch_ids)
+                if n_batch_samples != curr_bsz:
+                    batch_ids += indices_rank_i[: (curr_bsz - n_batch_samples)]
+                self.current += curr_bsz
+
+                if len(batch_ids) > 0:
+                    batch = [curr_w, curr_h, len(batch_ids)]
+                    self.batch_list.append(batch)
+        random.shuffle(self.batch_list)
+        self.length = len(self.batch_list)
+        self.batchs_in_one_epoch = self.iter()
+        self.batchs_in_one_epoch_id = [i for i in range(len(self.batchs_in_one_epoch))]
+
+    def __iter__(self):
+        if self.seed is None:
+            random.seed(self.epoch)
+            self.epoch += 1
+        else:
+            random.seed(self.seed)
+        random.shuffle(self.batchs_in_one_epoch_id)
+        for batch_tuple_id in self.batchs_in_one_epoch_id:
+            yield self.batchs_in_one_epoch[batch_tuple_id]
+
+    def iter(self):
+        if self.shuffle:
+            if self.seed is not None:
+                random.seed(self.seed)
+            else:
+                random.seed(self.epoch)
+            if not self.ds_width:
+                random.shuffle(self.img_indices)
+            random.shuffle(self.img_batch_pairs)
+            indices_rank_i = self.img_indices[
+                self.rank : len(self.img_indices) : self.num_replicas
+            ]
+        else:
+            indices_rank_i = self.img_indices[
+                self.rank : len(self.img_indices) : self.num_replicas
+            ]
+
+        start_index = 0
+        batchs_in_one_epoch = []
+        for batch_tuple in self.batch_list:
+            curr_w, curr_h, curr_bsz = batch_tuple
+            end_index = min(start_index + curr_bsz, self.n_samples_per_replica)
+            batch_ids = indices_rank_i[start_index:end_index]
+            n_batch_samples = len(batch_ids)
+            if n_batch_samples != curr_bsz:
+                batch_ids += indices_rank_i[: (curr_bsz - n_batch_samples)]
+            start_index += curr_bsz
+
+            if len(batch_ids) > 0:
+                if self.ds_width:
+                    wh_ratio_current = self.wh_ratio[self.wh_ratio_sort[batch_ids]]
+                    ratio_current = wh_ratio_current.mean()
+                    ratio_current = (
+                        ratio_current
+                        if ratio_current * curr_h < self.max_w
+                        else self.max_w / curr_h
+                    )
+                else:
+                    ratio_current = None
+                batch = [(curr_w, curr_h, b_id, ratio_current) for b_id in batch_ids]
+                # yield batch
+                batchs_in_one_epoch.append(batch)
+        return batchs_in_one_epoch
+
+    def set_epoch(self, epoch: int):
+        self.epoch = epoch
+
+    def __len__(self):
+        return self.length

ppocr/data/pgnet_dataset.py

@@ -0,0 +1,107 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import numpy as np
+import os
+from paddle.io import Dataset
+from .imaug import transform, create_operators
+import random
+
+
+class PGDataSet(Dataset):
+    def __init__(self, config, mode, logger, seed=None):
+        super(PGDataSet, self).__init__()
+
+        self.logger = logger
+        self.seed = seed
+        self.mode = mode
+        global_config = config["Global"]
+        dataset_config = config[mode]["dataset"]
+        loader_config = config[mode]["loader"]
+
+        self.delimiter = dataset_config.get("delimiter", "\t")
+        label_file_list = dataset_config.pop("label_file_list")
+        data_source_num = len(label_file_list)
+        ratio_list = dataset_config.get("ratio_list", [1.0])
+        if isinstance(ratio_list, (float, int)):
+            ratio_list = [float(ratio_list)] * int(data_source_num)
+        assert (
+            len(ratio_list) == data_source_num
+        ), "The length of ratio_list should be the same as the file_list."
+        self.data_dir = dataset_config["data_dir"]
+        self.do_shuffle = loader_config["shuffle"]
+
+        logger.info("Initialize indexes of datasets:%s" % label_file_list)
+        self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
+        self.data_idx_order_list = list(range(len(self.data_lines)))
+        if mode.lower() == "train":
+            self.shuffle_data_random()
+
+        self.ops = create_operators(dataset_config["transforms"], global_config)
+
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def shuffle_data_random(self):
+        if self.do_shuffle:
+            random.seed(self.seed)
+            random.shuffle(self.data_lines)
+        return
+
+    def get_image_info_list(self, file_list, ratio_list):
+        if isinstance(file_list, str):
+            file_list = [file_list]
+        data_lines = []
+        for idx, file in enumerate(file_list):
+            with open(file, "rb") as f:
+                lines = f.readlines()
+                if self.mode == "train" or ratio_list[idx] < 1.0:
+                    random.seed(self.seed)
+                    lines = random.sample(lines, round(len(lines) * ratio_list[idx]))
+                data_lines.extend(lines)
+        return data_lines
+
+    def __getitem__(self, idx):
+        file_idx = self.data_idx_order_list[idx]
+        data_line = self.data_lines[file_idx]
+        img_id = 0
+        try:
+            data_line = data_line.decode("utf-8")
+            substr = data_line.strip("\n").split(self.delimiter)
+            file_name = substr[0]
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            if self.mode.lower() == "eval":
+                try:
+                    img_id = int(data_line.split(".")[0][7:])
+                except:
+                    img_id = 0
+            data = {"img_path": img_path, "label": label, "img_id": img_id}
+            if not os.path.exists(img_path):
+                raise Exception("{} does not exist!".format(img_path))
+            with open(data["img_path"], "rb") as f:
+                img = f.read()
+                data["image"] = img
+            outs = transform(data, self.ops)
+        except Exception as e:
+            self.logger.error(
+                "When parsing line {}, error happened with msg: {}".format(
+                    self.data_idx_order_list[idx], e
+                )
+            )
+            outs = None
+        if outs is None:
+            return self.__getitem__(np.random.randint(self.__len__()))
+        return outs
+
+    def __len__(self):
+        return len(self.data_idx_order_list)

ppocr/data/pubtab_dataset.py

@@ -0,0 +1,138 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import numpy as np
+import os
+import random
+from paddle.io import Dataset
+import json
+from copy import deepcopy
+
+from .imaug import transform, create_operators
+
+
+class PubTabDataSet(Dataset):
+    def __init__(self, config, mode, logger, seed=None):
+        super(PubTabDataSet, self).__init__()
+        self.logger = logger
+
+        global_config = config["Global"]
+        dataset_config = config[mode]["dataset"]
+        loader_config = config[mode]["loader"]
+
+        label_file_list = dataset_config.pop("label_file_list")
+        data_source_num = len(label_file_list)
+        ratio_list = dataset_config.get("ratio_list", [1.0])
+        if isinstance(ratio_list, (float, int)):
+            ratio_list = [float(ratio_list)] * int(data_source_num)
+
+        assert (
+            len(ratio_list) == data_source_num
+        ), "The length of ratio_list should be the same as the file_list."
+
+        self.data_dir = dataset_config["data_dir"]
+        self.do_shuffle = loader_config["shuffle"]
+
+        self.seed = seed
+        self.mode = mode.lower()
+        logger.info("Initialize indexes of datasets:%s" % label_file_list)
+        self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
+        # self.check(config['Global']['max_text_length'])
+
+        if mode.lower() == "train" and self.do_shuffle:
+            self.shuffle_data_random()
+        self.ops = create_operators(dataset_config["transforms"], global_config)
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def get_image_info_list(self, file_list, ratio_list):
+        if isinstance(file_list, str):
+            file_list = [file_list]
+        data_lines = []
+        for idx, file in enumerate(file_list):
+            with open(file, "rb") as f:
+                lines = f.readlines()
+                if self.mode == "train" or ratio_list[idx] < 1.0:
+                    random.seed(self.seed)
+                    lines = random.sample(lines, round(len(lines) * ratio_list[idx]))
+                data_lines.extend(lines)
+        return data_lines
+
+    def check(self, max_text_length):
+        data_lines = []
+        for line in self.data_lines:
+            data_line = line.decode("utf-8").strip("\n")
+            info = json.loads(data_line)
+            file_name = info["filename"]
+            cells = info["html"]["cells"].copy()
+            structure = info["html"]["structure"]["tokens"].copy()
+
+            img_path = os.path.join(self.data_dir, file_name)
+            if not os.path.exists(img_path):
+                self.logger.warning("{} does not exist!".format(img_path))
+                continue
+            if len(structure) == 0 or len(structure) > max_text_length:
+                continue
+            # data = {'img_path': img_path, 'cells': cells, 'structure':structure,'file_name':file_name}
+            data_lines.append(line)
+        self.data_lines = data_lines
+
+    def shuffle_data_random(self):
+        if self.do_shuffle:
+            random.seed(self.seed)
+            random.shuffle(self.data_lines)
+        return
+
+    def __getitem__(self, idx):
+        try:
+            data_line = self.data_lines[idx]
+            data_line = data_line.decode("utf-8").strip("\n")
+            info = json.loads(data_line)
+            file_name = info["filename"]
+            cells = info["html"]["cells"].copy()
+            structure = info["html"]["structure"]["tokens"].copy()
+
+            img_path = os.path.join(self.data_dir, file_name)
+            if not os.path.exists(img_path):
+                raise Exception("{} does not exist!".format(img_path))
+            data = {
+                "img_path": img_path,
+                "cells": cells,
+                "structure": structure,
+                "file_name": file_name,
+            }
+
+            with open(data["img_path"], "rb") as f:
+                img = f.read()
+                data["image"] = img
+            outs = transform(data, self.ops)
+        except:
+            import traceback
+
+            err = traceback.format_exc()
+            self.logger.error(
+                "When parsing line {}, error happened with msg: {}".format(
+                    data_line, err
+                )
+            )
+            outs = None
+        if outs is None:
+            rnd_idx = (
+                np.random.randint(self.__len__())
+                if self.mode == "train"
+                else (idx + 1) % self.__len__()
+            )
+            return self.__getitem__(rnd_idx)
+        return outs
+
+    def __len__(self):
+        return len(self.data_lines)

ppocr/data/simple_dataset.py

@@ -0,0 +1,253 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import numpy as np
+import cv2
+import math
+import os
+import json
+import random
+import traceback
+from paddle.io import Dataset
+from .imaug import transform, create_operators
+
+
+class SimpleDataSet(Dataset):
+    def __init__(self, config, mode, logger, seed=None):
+        super(SimpleDataSet, self).__init__()
+        self.logger = logger
+        self.mode = mode.lower()
+
+        global_config = config["Global"]
+        dataset_config = config[mode]["dataset"]
+        loader_config = config[mode]["loader"]
+
+        self.delimiter = dataset_config.get("delimiter", "\t")
+        label_file_list = dataset_config.pop("label_file_list")
+        data_source_num = len(label_file_list)
+        ratio_list = dataset_config.get("ratio_list", 1.0)
+        if isinstance(ratio_list, (float, int)):
+            ratio_list = [float(ratio_list)] * int(data_source_num)
+
+        assert (
+            len(ratio_list) == data_source_num
+        ), "The length of ratio_list should be the same as the file_list."
+        self.data_dir = dataset_config["data_dir"]
+        self.do_shuffle = loader_config["shuffle"]
+        self.seed = seed
+        logger.info("Initialize indexes of datasets:%s" % label_file_list)
+        self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
+        self.data_idx_order_list = list(range(len(self.data_lines)))
+        if self.mode == "train" and self.do_shuffle:
+            self.shuffle_data_random()
+        self.ops = create_operators(dataset_config["transforms"], global_config)
+        self.ext_op_transform_idx = dataset_config.get("ext_op_transform_idx", 2)
+        self.need_reset = True in [x < 1 for x in ratio_list]
+
+    def get_image_info_list(self, file_list, ratio_list):
+        if isinstance(file_list, str):
+            file_list = [file_list]
+        data_lines = []
+        for idx, file in enumerate(file_list):
+            with open(file, "rb") as f:
+                lines = f.readlines()
+                if self.mode == "train" or ratio_list[idx] < 1.0:
+                    random.seed(self.seed)
+                    lines = random.sample(lines, round(len(lines) * ratio_list[idx]))
+                data_lines.extend(lines)
+        return data_lines
+
+    def shuffle_data_random(self):
+        random.seed(self.seed)
+        random.shuffle(self.data_lines)
+        return
+
+    def _try_parse_filename_list(self, file_name):
+        # multiple images -> one gt label
+        if len(file_name) > 0 and file_name[0] == "[":
+            try:
+                info = json.loads(file_name)
+                file_name = random.choice(info)
+            except:
+                pass
+        return file_name
+
+    def get_ext_data(self):
+        ext_data_num = 0
+        for op in self.ops:
+            if hasattr(op, "ext_data_num"):
+                ext_data_num = getattr(op, "ext_data_num")
+                break
+        load_data_ops = self.ops[: self.ext_op_transform_idx]
+        ext_data = []
+
+        while len(ext_data) < ext_data_num:
+            file_idx = self.data_idx_order_list[np.random.randint(self.__len__())]
+            data_line = self.data_lines[file_idx]
+            data_line = data_line.decode("utf-8")
+            substr = data_line.strip("\n").split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {"img_path": img_path, "label": label}
+            if not os.path.exists(img_path):
+                continue
+            with open(data["img_path"], "rb") as f:
+                img = f.read()
+                data["image"] = img
+            data = transform(data, load_data_ops)
+
+            if data is None:
+                continue
+            if "polys" in data.keys():
+                if data["polys"].shape[1] != 4:
+                    continue
+            ext_data.append(data)
+        return ext_data
+
+    def __getitem__(self, idx):
+        file_idx = self.data_idx_order_list[idx]
+        data_line = self.data_lines[file_idx]
+        try:
+            data_line = data_line.decode("utf-8")
+            substr = data_line.strip("\n").split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {"img_path": img_path, "label": label}
+            if not os.path.exists(img_path):
+                raise Exception("{} does not exist!".format(img_path))
+            with open(data["img_path"], "rb") as f:
+                img = f.read()
+                data["image"] = img
+            data["ext_data"] = self.get_ext_data()
+            data["filename"] = data["img_path"]
+            outs = transform(data, self.ops)
+        except:
+            self.logger.error(
+                "When parsing line {}, error happened with msg: {}".format(
+                    data_line, traceback.format_exc()
+                )
+            )
+            outs = None
+        if outs is None:
+            # during evaluation, we should fix the idx to get same results for many times of evaluation.
+            rnd_idx = (
+                np.random.randint(self.__len__())
+                if self.mode == "train"
+                else (idx + 1) % self.__len__()
+            )
+            return self.__getitem__(rnd_idx)
+        return outs
+
+    def __len__(self):
+        return len(self.data_idx_order_list)
+
+
+class MultiScaleDataSet(SimpleDataSet):
+    def __init__(self, config, mode, logger, seed=None):
+        super(MultiScaleDataSet, self).__init__(config, mode, logger, seed)
+        self.ds_width = config[mode]["dataset"].get("ds_width", False)
+        if self.ds_width:
+            self.wh_aware()
+
+    def wh_aware(self):
+        data_line_new = []
+        wh_ratio = []
+        for lins in self.data_lines:
+            data_line_new.append(lins)
+            lins = lins.decode("utf-8")
+            name, label, w, h = lins.strip("\n").split(self.delimiter)
+            wh_ratio.append(float(w) / float(h))
+
+        self.data_lines = data_line_new
+        self.wh_ratio = np.array(wh_ratio)
+        self.wh_ratio_sort = np.argsort(self.wh_ratio)
+        self.data_idx_order_list = list(range(len(self.data_lines)))
+
+    def resize_norm_img(self, data, imgW, imgH, padding=True):
+        img = data["image"]
+        h = img.shape[0]
+        w = img.shape[1]
+        if not padding:
+            resized_image = cv2.resize(
+                img, (imgW, imgH), interpolation=cv2.INTER_LINEAR
+            )
+            resized_w = imgW
+        else:
+            ratio = w / float(h)
+            if math.ceil(imgH * ratio) > imgW:
+                resized_w = imgW
+            else:
+                resized_w = int(math.ceil(imgH * ratio))
+            resized_image = cv2.resize(img, (resized_w, imgH))
+        resized_image = resized_image.astype("float32")
+
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+        resized_image -= 0.5
+        resized_image /= 0.5
+        padding_im = np.zeros((3, imgH, imgW), dtype=np.float32)
+        padding_im[:, :, :resized_w] = resized_image
+        valid_ratio = min(1.0, float(resized_w / imgW))
+        data["image"] = padding_im
+        data["valid_ratio"] = valid_ratio
+        return data
+
+    def __getitem__(self, properties):
+        # properties is a tuple, contains (width, height, index)
+        img_height = properties[1]
+        idx = properties[2]
+        if self.ds_width and properties[3] is not None:
+            wh_ratio = properties[3]
+            img_width = img_height * (
+                1 if int(round(wh_ratio)) == 0 else int(round(wh_ratio))
+            )
+            file_idx = self.wh_ratio_sort[idx]
+        else:
+            file_idx = self.data_idx_order_list[idx]
+            img_width = properties[0]
+            wh_ratio = None
+
+        data_line = self.data_lines[file_idx]
+        try:
+            data_line = data_line.decode("utf-8")
+            substr = data_line.strip("\n").split(self.delimiter)
+            file_name = substr[0]
+            file_name = self._try_parse_filename_list(file_name)
+            label = substr[1]
+            img_path = os.path.join(self.data_dir, file_name)
+            data = {"img_path": img_path, "label": label}
+            if not os.path.exists(img_path):
+                raise Exception("{} does not exist!".format(img_path))
+            with open(data["img_path"], "rb") as f:
+                img = f.read()
+                data["image"] = img
+            data["ext_data"] = self.get_ext_data()
+            outs = transform(data, self.ops[:-1])
+            if outs is not None:
+                outs = self.resize_norm_img(outs, img_width, img_height)
+                outs = transform(outs, self.ops[-1:])
+        except:
+            self.logger.error(
+                "When parsing line {}, error happened with msg: {}".format(
+                    data_line, traceback.format_exc()
+                )
+            )
+            outs = None
+        if outs is None:
+            # during evaluation, we should fix the idx to get same results for many times of evaluation.
+            rnd_idx = (idx + 1) % self.__len__()
+            return self.__getitem__([img_width, img_height, rnd_idx, wh_ratio])
+        return outs

ppocr/ext_op/__init__.py

@@ -0,0 +1 @@
+from .roi_align_rotated.roi_align_rotated import RoIAlignRotated

ppocr/ext_op/roi_align_rotated/roi_align_rotated.cc

@@ -0,0 +1,528 @@
+
+// This code is refer from:
+// https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/pytorch/cpu/roi_align_rotated.cpp
+
+#include <cassert>
+#include <cmath>
+#include <vector>
+
+#include "paddle/extension.h"
+
+#define PADDLE_WITH_CUDA
+#define CHECK_INPUT_SAME(x1, x2)                                               \
+  PD_CHECK(x1.place() == x2.place(), "input must be same place.")
+#define CHECK_INPUT_CPU(x) PD_CHECK(x.is_cpu(), #x " must be a CPU Tensor.")
+
+template <typename T> struct PreCalc {
+  int pos1;
+  int pos2;
+  int pos3;
+  int pos4;
+  T w1;
+  T w2;
+  T w3;
+  T w4;
+};
+
+template <typename T>
+void pre_calc_for_bilinear_interpolate(
+    const int height, const int width, const int pooled_height,
+    const int pooled_width, const int iy_upper, const int ix_upper,
+    T roi_start_h, T roi_start_w, T bin_size_h, T bin_size_w,
+    int roi_bin_grid_h, int roi_bin_grid_w, T roi_center_h, T roi_center_w,
+    T cos_theta, T sin_theta, std::vector<PreCalc<T>> &pre_calc) {
+  int pre_calc_index = 0;
+  for (int ph = 0; ph < pooled_height; ph++) {
+    for (int pw = 0; pw < pooled_width; pw++) {
+      for (int iy = 0; iy < iy_upper; iy++) {
+        const T yy = roi_start_h + ph * bin_size_h +
+                     static_cast<T>(iy + .5f) * bin_size_h /
+                         static_cast<T>(roi_bin_grid_h); // e.g., 0.5, 1.5
+        for (int ix = 0; ix < ix_upper; ix++) {
+          const T xx = roi_start_w + pw * bin_size_w +
+                       static_cast<T>(ix + .5f) * bin_size_w /
+                           static_cast<T>(roi_bin_grid_w);
+
+          // Rotate by theta around the center and translate
+          // In image space, (y, x) is the order for Right Handed System,
+          // and this is essentially multiplying the point by a rotation matrix
+          // to rotate it counterclockwise through angle theta.
+          T y = yy * cos_theta - xx * sin_theta + roi_center_h;
+          T x = yy * sin_theta + xx * cos_theta + roi_center_w;
+          // deal with: inverse elements are out of feature map boundary
+          if (y < -1.0 || y > height || x < -1.0 || x > width) {
+            // empty
+            PreCalc<T> pc;
+            pc.pos1 = 0;
+            pc.pos2 = 0;
+            pc.pos3 = 0;
+            pc.pos4 = 0;
+            pc.w1 = 0;
+            pc.w2 = 0;
+            pc.w3 = 0;
+            pc.w4 = 0;
+            pre_calc[pre_calc_index] = pc;
+            pre_calc_index += 1;
+            continue;
+          }
+
+          if (y < 0) {
+            y = 0;
+          }
+          if (x < 0) {
+            x = 0;
+          }
+
+          int y_low = (int)y;
+          int x_low = (int)x;
+          int y_high;
+          int x_high;
+
+          if (y_low >= height - 1) {
+            y_high = y_low = height - 1;
+            y = (T)y_low;
+          } else {
+            y_high = y_low + 1;
+          }
+
+          if (x_low >= width - 1) {
+            x_high = x_low = width - 1;
+            x = (T)x_low;
+          } else {
+            x_high = x_low + 1;
+          }
+
+          T ly = y - y_low;
+          T lx = x - x_low;
+          T hy = 1. - ly, hx = 1. - lx;
+          T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+
+          // save weights and indices
+          PreCalc<T> pc;
+          pc.pos1 = y_low * width + x_low;
+          pc.pos2 = y_low * width + x_high;
+          pc.pos3 = y_high * width + x_low;
+          pc.pos4 = y_high * width + x_high;
+          pc.w1 = w1;
+          pc.w2 = w2;
+          pc.w3 = w3;
+          pc.w4 = w4;
+          pre_calc[pre_calc_index] = pc;
+
+          pre_calc_index += 1;
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+void roi_align_rotated_cpu_forward(const int nthreads, const T *input,
+                                   const T &spatial_scale, const bool aligned,
+                                   const bool clockwise, const int channels,
+                                   const int height, const int width,
+                                   const int pooled_height,
+                                   const int pooled_width,
+                                   const int sampling_ratio, const T *rois,
+                                   T *output) {
+  int n_rois = nthreads / channels / pooled_width / pooled_height;
+  // (n, c, ph, pw) is an element in the pooled output
+  // can be parallelized using omp
+  // #pragma omp parallel for num_threads(32)
+  for (int n = 0; n < n_rois; n++) {
+    int index_n = n * channels * pooled_width * pooled_height;
+
+    const T *current_roi = rois + n * 6;
+    int roi_batch_ind = current_roi[0];
+
+    // Do not use rounding; this implementation detail is critical
+    T offset = aligned ? (T)0.5 : (T)0.0;
+    T roi_center_w = current_roi[1] * spatial_scale - offset;
+    T roi_center_h = current_roi[2] * spatial_scale - offset;
+    T roi_width = current_roi[3] * spatial_scale;
+    T roi_height = current_roi[4] * spatial_scale;
+    T theta = current_roi[5];
+    if (clockwise) {
+      theta = -theta; // If clockwise, the angle needs to be reversed.
+    }
+    T cos_theta = cos(theta);
+    T sin_theta = sin(theta);
+
+    if (aligned) {
+      assert(roi_width >= 0 && roi_height >= 0);
+    } else { // for backward-compatibility only
+      roi_width = std::max(roi_width, (T)1.);
+      roi_height = std::max(roi_height, (T)1.);
+    }
+
+    T bin_size_h = static_cast<T>(roi_height) / static_cast<T>(pooled_height);
+    T bin_size_w = static_cast<T>(roi_width) / static_cast<T>(pooled_width);
+
+    // We use roi_bin_grid to sample the grid and mimic integral
+    int roi_bin_grid_h = (sampling_ratio > 0)
+                             ? sampling_ratio
+                             : ceilf(roi_height / pooled_height); // e.g., = 2
+    int roi_bin_grid_w =
+        (sampling_ratio > 0) ? sampling_ratio : ceilf(roi_width / pooled_width);
+
+    // We do average (integral) pooling inside a bin
+    const T count = std::max(roi_bin_grid_h * roi_bin_grid_w, 1); // e.g. = 4
+
+    // we want to precalculate indices and weights shared by all channels,
+    // this is the key point of optimization
+    std::vector<PreCalc<T>> pre_calc(roi_bin_grid_h * roi_bin_grid_w *
+                                     pooled_width * pooled_height);
+
+    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
+    // Appropriate translation needs to be applied after.
+    T roi_start_h = -roi_height / 2.0;
+    T roi_start_w = -roi_width / 2.0;
+
+    pre_calc_for_bilinear_interpolate(
+        height, width, pooled_height, pooled_width, roi_bin_grid_h,
+        roi_bin_grid_w, roi_start_h, roi_start_w, bin_size_h, bin_size_w,
+        roi_bin_grid_h, roi_bin_grid_w, roi_center_h, roi_center_w, cos_theta,
+        sin_theta, pre_calc);
+
+    for (int c = 0; c < channels; c++) {
+      int index_n_c = index_n + c * pooled_width * pooled_height;
+      const T *offset_input =
+          input + (roi_batch_ind * channels + c) * height * width;
+      int pre_calc_index = 0;
+
+      for (int ph = 0; ph < pooled_height; ph++) {
+        for (int pw = 0; pw < pooled_width; pw++) {
+          int index = index_n_c + ph * pooled_width + pw;
+
+          T output_val = 0.;
+          for (int iy = 0; iy < roi_bin_grid_h; iy++) {
+            for (int ix = 0; ix < roi_bin_grid_w; ix++) {
+              PreCalc<T> pc = pre_calc[pre_calc_index];
+              output_val += pc.w1 * offset_input[pc.pos1] +
+                            pc.w2 * offset_input[pc.pos2] +
+                            pc.w3 * offset_input[pc.pos3] +
+                            pc.w4 * offset_input[pc.pos4];
+
+              pre_calc_index += 1;
+            }
+          }
+          output_val /= count;
+
+          output[index] = output_val;
+        } // for pw
+      }   // for ph
+    }     // for c
+  }       // for n
+}
+
+template <typename T>
+void bilinear_interpolate_gradient(const int height, const int width, T y, T x,
+                                   T &w1, T &w2, T &w3, T &w4, int &x_low,
+                                   int &x_high, int &y_low, int &y_high) {
+  // deal with cases that inverse elements are out of feature map boundary
+  if (y < -1.0 || y > height || x < -1.0 || x > width) {
+    // empty
+    w1 = w2 = w3 = w4 = 0.;
+    x_low = x_high = y_low = y_high = -1;
+    return;
+  }
+
+  if (y < 0) {
+    y = 0;
+  }
+
+  if (x < 0) {
+    x = 0;
+  }
+
+  y_low = (int)y;
+  x_low = (int)x;
+
+  if (y_low >= height - 1) {
+    y_high = y_low = height - 1;
+    y = (T)y_low;
+  } else {
+    y_high = y_low + 1;
+  }
+
+  if (x_low >= width - 1) {
+    x_high = x_low = width - 1;
+    x = (T)x_low;
+  } else {
+    x_high = x_low + 1;
+  }
+
+  T ly = y - y_low;
+  T lx = x - x_low;
+  T hy = 1. - ly, hx = 1. - lx;
+
+  // reference in forward
+  // T v1 = input[y_low * width + x_low];
+  // T v2 = input[y_low * width + x_high];
+  // T v3 = input[y_high * width + x_low];
+  // T v4 = input[y_high * width + x_high];
+  // T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+
+  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+
+  return;
+}
+
+template <class T> inline void add(T *address, const T &val) {
+  *address += val;
+}
+
+template <typename T>
+void roi_align_rotated_cpu_backward(
+    const int nthreads,
+    // may not be contiguous. should index using n_stride, etc
+    const T *grad_output, const T &spatial_scale, const bool aligned,
+    const bool clockwise, const int channels, const int height, const int width,
+    const int pooled_height, const int pooled_width, const int sampling_ratio,
+    T *grad_input, const T *rois, const int n_stride, const int c_stride,
+    const int h_stride, const int w_stride) {
+  for (int index = 0; index < nthreads; index++) {
+    // (n, c, ph, pw) is an element in the pooled output
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int c = (index / pooled_width / pooled_height) % channels;
+    int n = index / pooled_width / pooled_height / channels;
+
+    const T *current_roi = rois + n * 6;
+    int roi_batch_ind = current_roi[0];
+
+    // Do not use rounding; this implementation detail is critical
+    T offset = aligned ? (T)0.5 : (T)0.0;
+    T roi_center_w = current_roi[1] * spatial_scale - offset;
+    T roi_center_h = current_roi[2] * spatial_scale - offset;
+    T roi_width = current_roi[3] * spatial_scale;
+    T roi_height = current_roi[4] * spatial_scale;
+    T theta = current_roi[5];
+    if (clockwise) {
+      theta = -theta; // If clockwise, the angle needs to be reversed.
+    }
+    T cos_theta = cos(theta);
+    T sin_theta = sin(theta);
+
+    if (aligned) {
+      assert(roi_width >= 0 && roi_height >= 0);
+    } else { // for backward-compatibility only
+      roi_width = std::max(roi_width, (T)1.);
+      roi_height = std::max(roi_height, (T)1.);
+    }
+
+    T bin_size_h = static_cast<T>(roi_height) / static_cast<T>(pooled_height);
+    T bin_size_w = static_cast<T>(roi_width) / static_cast<T>(pooled_width);
+
+    T *offset_grad_input =
+        grad_input + ((roi_batch_ind * channels + c) * height * width);
+
+    int output_offset = n * n_stride + c * c_stride;
+    const T *offset_grad_output = grad_output + output_offset;
+    const T grad_output_this_bin =
+        offset_grad_output[ph * h_stride + pw * w_stride];
+
+    // We use roi_bin_grid to sample the grid and mimic integral
+    int roi_bin_grid_h = (sampling_ratio > 0)
+                             ? sampling_ratio
+                             : ceilf(roi_height / pooled_height); // e.g., = 2
+    int roi_bin_grid_w =
+        (sampling_ratio > 0) ? sampling_ratio : ceilf(roi_width / pooled_width);
+
+    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
+    // Appropriate translation needs to be applied after.
+    T roi_start_h = -roi_height / 2.0;
+    T roi_start_w = -roi_width / 2.0;
+
+    // We do average (integral) pooling inside a bin
+    const T count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
+
+    for (int iy = 0; iy < roi_bin_grid_h; iy++) {
+      const T yy = roi_start_h + ph * bin_size_h +
+                   static_cast<T>(iy + .5f) * bin_size_h /
+                       static_cast<T>(roi_bin_grid_h); // e.g., 0.5, 1.5
+      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
+        const T xx = roi_start_w + pw * bin_size_w +
+                     static_cast<T>(ix + .5f) * bin_size_w /
+                         static_cast<T>(roi_bin_grid_w);
+
+        // Rotate by theta around the center and translate
+        T y = yy * cos_theta - xx * sin_theta + roi_center_h;
+        T x = yy * sin_theta + xx * cos_theta + roi_center_w;
+
+        T w1, w2, w3, w4;
+        int x_low, x_high, y_low, y_high;
+
+        bilinear_interpolate_gradient(height, width, y, x, w1, w2, w3, w4,
+                                      x_low, x_high, y_low, y_high);
+
+        T g1 = grad_output_this_bin * w1 / count;
+        T g2 = grad_output_this_bin * w2 / count;
+        T g3 = grad_output_this_bin * w3 / count;
+        T g4 = grad_output_this_bin * w4 / count;
+
+        if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
+          // atomic add is not needed for now since it is single threaded
+          add(offset_grad_input + y_low * width + x_low, static_cast<T>(g1));
+          add(offset_grad_input + y_low * width + x_high, static_cast<T>(g2));
+          add(offset_grad_input + y_high * width + x_low, static_cast<T>(g3));
+          add(offset_grad_input + y_high * width + x_high, static_cast<T>(g4));
+        } // if
+      }   // ix
+    }     // iy
+  }       // for
+} // ROIAlignRotatedBackward
+
+std::vector<paddle::Tensor>
+RoIAlignRotatedCPUForward(const paddle::Tensor &input,
+                          const paddle::Tensor &rois, int aligned_height,
+                          int aligned_width, float spatial_scale,
+                          int sampling_ratio, bool aligned, bool clockwise) {
+  CHECK_INPUT_CPU(input);
+  CHECK_INPUT_CPU(rois);
+
+  auto num_rois = rois.shape()[0];
+
+  auto channels = input.shape()[1];
+  auto height = input.shape()[2];
+  auto width = input.shape()[3];
+
+  auto output =
+      paddle::empty({num_rois, channels, aligned_height, aligned_width},
+                    input.type(), paddle::CPUPlace());
+  auto output_size = output.numel();
+
+  PD_DISPATCH_FLOATING_TYPES(
+      input.type(), "roi_align_rotated_cpu_forward", ([&] {
+        roi_align_rotated_cpu_forward<data_t>(
+            output_size, input.data<data_t>(),
+            static_cast<data_t>(spatial_scale), aligned, clockwise, channels,
+            height, width, aligned_height, aligned_width, sampling_ratio,
+            rois.data<data_t>(), output.data<data_t>());
+      }));
+
+  return {output};
+}
+
+std::vector<paddle::Tensor> RoIAlignRotatedCPUBackward(
+    const paddle::Tensor &input, const paddle::Tensor &rois,
+    const paddle::Tensor &grad_output, int aligned_height, int aligned_width,
+    float spatial_scale, int sampling_ratio, bool aligned, bool clockwise) {
+
+  auto batch_size = input.shape()[0];
+  auto channels = input.shape()[1];
+  auto height = input.shape()[2];
+  auto width = input.shape()[3];
+
+  auto grad_input = paddle::full({batch_size, channels, height, width}, 0.0,
+                                 input.type(), paddle::CPUPlace());
+
+  // get stride values to ensure indexing into gradients is correct.
+  int n_stride = grad_output.shape()[0];
+  int c_stride = grad_output.shape()[1];
+  int h_stride = grad_output.shape()[2];
+  int w_stride = grad_output.shape()[3];
+
+  PD_DISPATCH_FLOATING_TYPES(
+      grad_output.type(), "roi_align_rotated_cpu_backward", [&] {
+        roi_align_rotated_cpu_backward<data_t>(
+            grad_output.numel(), grad_output.data<data_t>(),
+            static_cast<data_t>(spatial_scale), aligned, clockwise, channels,
+            height, width, aligned_height, aligned_width, sampling_ratio,
+            grad_input.data<data_t>(), rois.data<data_t>(), n_stride, c_stride,
+            h_stride, w_stride);
+      });
+  return {grad_input};
+}
+
+#ifdef PADDLE_WITH_CUDA
+std::vector<paddle::Tensor>
+RoIAlignRotatedCUDAForward(const paddle::Tensor &input,
+                           const paddle::Tensor &rois, int aligned_height,
+                           int aligned_width, float spatial_scale,
+                           int sampling_ratio, bool aligned, bool clockwise);
+#endif
+
+#ifdef PADDLE_WITH_CUDA
+std::vector<paddle::Tensor> RoIAlignRotatedCUDABackward(
+    const paddle::Tensor &input, const paddle::Tensor &rois,
+    const paddle::Tensor &grad_output, int aligned_height, int aligned_width,
+    float spatial_scale, int sampling_ratio, bool aligned, bool clockwise);
+#endif
+
+std::vector<paddle::Tensor>
+RoIAlignRotatedForward(const paddle::Tensor &input, const paddle::Tensor &rois,
+                       int aligned_height, int aligned_width,
+                       float spatial_scale, int sampling_ratio, bool aligned,
+                       bool clockwise) {
+  CHECK_INPUT_SAME(input, rois);
+  if (input.is_cpu()) {
+    return RoIAlignRotatedCPUForward(input, rois, aligned_height, aligned_width,
+                                     spatial_scale, sampling_ratio, aligned,
+                                     clockwise);
+#ifdef PADDLE_WITH_CUDA
+  } else if (input.is_gpu()) {
+    return RoIAlignRotatedCUDAForward(input, rois, aligned_height,
+                                      aligned_width, spatial_scale,
+                                      sampling_ratio, aligned, clockwise);
+#endif
+  } else {
+    PD_THROW("Unsupported device type for forward function of roi align "
+             "rotated operator.");
+  }
+}
+
+std::vector<paddle::Tensor>
+RoIAlignRotatedBackward(const paddle::Tensor &input, const paddle::Tensor &rois,
+                        const paddle::Tensor &grad_output, int aligned_height,
+                        int aligned_width, float spatial_scale,
+                        int sampling_ratio, bool aligned, bool clockwise) {
+  CHECK_INPUT_SAME(input, rois);
+  if (input.is_cpu()) {
+    return RoIAlignRotatedCPUBackward(input, rois, grad_output, aligned_height,
+                                      aligned_width, spatial_scale,
+                                      sampling_ratio, aligned, clockwise);
+#ifdef PADDLE_WITH_CUDA
+  } else if (input.is_gpu()) {
+    return RoIAlignRotatedCUDABackward(input, rois, grad_output, aligned_height,
+                                       aligned_width, spatial_scale,
+                                       sampling_ratio, aligned, clockwise);
+#endif
+  } else {
+    PD_THROW("Unsupported device type for forward function of roi align "
+             "rotated operator.");
+  }
+}
+
+std::vector<std::vector<int64_t>> InferShape(std::vector<int64_t> input_shape,
+                                             std::vector<int64_t> rois_shape) {
+  return {{rois_shape[0], input_shape[1], input_shape[2], input_shape[3]}};
+}
+
+std::vector<std::vector<int64_t>>
+InferBackShape(std::vector<int64_t> input_shape,
+               std::vector<int64_t> rois_shape) {
+  return {input_shape};
+}
+
+std::vector<paddle::DataType> InferDtype(paddle::DataType input_dtype,
+                                         paddle::DataType rois_dtype) {
+  return {input_dtype};
+}
+
+PD_BUILD_OP(roi_align_rotated)
+    .Inputs({"Input", "Rois"})
+    .Outputs({"Output"})
+    .Attrs({"aligned_height: int", "aligned_width: int", "spatial_scale: float",
+            "sampling_ratio: int", "aligned: bool", "clockwise: bool"})
+    .SetKernelFn(PD_KERNEL(RoIAlignRotatedForward))
+    .SetInferShapeFn(PD_INFER_SHAPE(InferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(InferDtype));
+
+PD_BUILD_GRAD_OP(roi_align_rotated)
+    .Inputs({"Input", "Rois", paddle::Grad("Output")})
+    .Attrs({"aligned_height: int", "aligned_width: int", "spatial_scale: float",
+            "sampling_ratio: int", "aligned: bool", "clockwise: bool"})
+    .Outputs({paddle::Grad("Input")})
+    .SetKernelFn(PD_KERNEL(RoIAlignRotatedBackward))
+    .SetInferShapeFn(PD_INFER_SHAPE(InferBackShape));

ppocr/ext_op/roi_align_rotated/roi_align_rotated.cu

@@ -0,0 +1,381 @@
+
+// This code is refer from:
+// https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/csrc/common/cuda/roi_align_rotated_cuda_kernel.cuh
+
+#include <cassert>
+#include <cmath>
+#include <vector>
+
+#include "paddle/extension.h"
+#include <cuda.h>
+
+#define CUDA_1D_KERNEL_LOOP(i, n)                                              \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n);                 \
+       i += blockDim.x * gridDim.x)
+
+#define THREADS_PER_BLOCK 512
+
+inline int GET_BLOCKS(const int N) {
+  int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+  int max_block_num = 4096;
+  return min(optimal_block_num, max_block_num);
+}
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600
+
+static __inline__ __device__ double atomicAdd(double *address, double val) {
+  unsigned long long int *address_as_ull = (unsigned long long int *)address;
+  unsigned long long int old = *address_as_ull, assumed;
+  if (val == 0.0)
+    return __longlong_as_double(old);
+  do {
+    assumed = old;
+    old = atomicCAS(address_as_ull, assumed,
+                    __double_as_longlong(val + __longlong_as_double(assumed)));
+  } while (assumed != old);
+  return __longlong_as_double(old);
+}
+
+#endif
+
+template <typename T>
+__device__ T bilinear_interpolate(const T *input, const int height,
+                                  const int width, T y, T x,
+                                  const int index /* index for debug only*/) {
+  // deal with cases that inverse elements are out of feature map boundary
+  if (y < -1.0 || y > height || x < -1.0 || x > width)
+    return 0;
+
+  if (y <= 0)
+    y = 0;
+  if (x <= 0)
+    x = 0;
+
+  int y_low = (int)y;
+  int x_low = (int)x;
+  int y_high;
+  int x_high;
+
+  if (y_low >= height - 1) {
+    y_high = y_low = height - 1;
+    y = (T)y_low;
+  } else {
+    y_high = y_low + 1;
+  }
+
+  if (x_low >= width - 1) {
+    x_high = x_low = width - 1;
+    x = (T)x_low;
+  } else {
+    x_high = x_low + 1;
+  }
+
+  T ly = y - y_low;
+  T lx = x - x_low;
+  T hy = 1. - ly, hx = 1. - lx;
+  // do bilinear interpolation
+  T v1 = input[y_low * width + x_low];
+  T v2 = input[y_low * width + x_high];
+  T v3 = input[y_high * width + x_low];
+  T v4 = input[y_high * width + x_high];
+  T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+
+  T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+
+  return val;
+}
+
+template <typename T>
+__device__ void
+bilinear_interpolate_gradient(const int height, const int width, T y, T x,
+                              T &w1, T &w2, T &w3, T &w4, int &x_low,
+                              int &x_high, int &y_low, int &y_high,
+                              const int index /* index for debug only*/) {
+  // deal with cases that inverse elements are out of feature map boundary
+  if (y < -1.0 || y > height || x < -1.0 || x > width) {
+    // empty
+    w1 = w2 = w3 = w4 = 0.;
+    x_low = x_high = y_low = y_high = -1;
+    return;
+  }
+
+  if (y <= 0)
+    y = 0;
+  if (x <= 0)
+    x = 0;
+
+  y_low = (int)y;
+  x_low = (int)x;
+
+  if (y_low >= height - 1) {
+    y_high = y_low = height - 1;
+    y = (T)y_low;
+  } else {
+    y_high = y_low + 1;
+  }
+
+  if (x_low >= width - 1) {
+    x_high = x_low = width - 1;
+    x = (T)x_low;
+  } else {
+    x_high = x_low + 1;
+  }
+
+  T ly = y - y_low;
+  T lx = x - x_low;
+  T hy = 1. - ly, hx = 1. - lx;
+
+  // reference in forward
+  // T v1 = input[y_low * width + x_low];
+  // T v2 = input[y_low * width + x_high];
+  // T v3 = input[y_high * width + x_low];
+  // T v4 = input[y_high * width + x_high];
+  // T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+
+  w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+
+  return;
+}
+
+/*** Forward ***/
+template <typename scalar_t>
+__global__ void roi_align_rotated_cuda_forward_kernel(
+    const int nthreads, const scalar_t *bottom_data,
+    const scalar_t *bottom_rois, const scalar_t spatial_scale,
+    const int sample_num, const bool aligned, const bool clockwise,
+    const int channels, const int height, const int width,
+    const int pooled_height, const int pooled_width, scalar_t *top_data) {
+  CUDA_1D_KERNEL_LOOP(index, nthreads) {
+    // (n, c, ph, pw) is an element in the pooled output
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int c = (index / pooled_width / pooled_height) % channels;
+    int n = index / pooled_width / pooled_height / channels;
+
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 6;
+    int roi_batch_ind = offset_bottom_rois[0];
+
+    // Do not using rounding; this implementation detail is critical
+    scalar_t offset = aligned ? (scalar_t)0.5 : (scalar_t)0.0;
+    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale - offset;
+    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale - offset;
+    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
+    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
+    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
+    scalar_t theta = offset_bottom_rois[5];
+    if (clockwise) {
+      theta = -theta; // If clockwise, the angle needs to be reversed.
+    }
+    if (!aligned) { // for backward-compatibility only
+      // Force malformed ROIs to be 1x1
+      roi_width = max(roi_width, (scalar_t)1.);
+      roi_height = max(roi_height, (scalar_t)1.);
+    }
+    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) /
+                          static_cast<scalar_t>(pooled_height);
+    scalar_t bin_size_w =
+        static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);
+
+    const scalar_t *offset_bottom_data =
+        bottom_data + (roi_batch_ind * channels + c) * height * width;
+
+    // We use roi_bin_grid to sample the grid and mimic integral
+    int roi_bin_grid_h = (sample_num > 0)
+                             ? sample_num
+                             : ceilf(roi_height / pooled_height); // e.g., = 2
+    int roi_bin_grid_w =
+        (sample_num > 0) ? sample_num : ceilf(roi_width / pooled_width);
+
+    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
+    // Appropriate translation needs to be applied after.
+    scalar_t roi_start_h = -roi_height / 2.0;
+    scalar_t roi_start_w = -roi_width / 2.0;
+    scalar_t cosscalar_theta = cos(theta);
+    scalar_t sinscalar_theta = sin(theta);
+
+    // We do average (integral) pooling inside a bin
+    const scalar_t count = max(roi_bin_grid_h * roi_bin_grid_w, 1); // e.g. = 4
+
+    scalar_t output_val = 0.;
+    for (int iy = 0; iy < roi_bin_grid_h; iy++) { // e.g., iy = 0, 1
+      const scalar_t yy =
+          roi_start_h + ph * bin_size_h +
+          static_cast<scalar_t>(iy + .5f) * bin_size_h /
+              static_cast<scalar_t>(roi_bin_grid_h); // e.g., 0.5, 1.5
+      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
+        const scalar_t xx = roi_start_w + pw * bin_size_w +
+                            static_cast<scalar_t>(ix + .5f) * bin_size_w /
+                                static_cast<scalar_t>(roi_bin_grid_w);
+
+        // Rotate by theta (counterclockwise) around the center and translate
+        scalar_t y = yy * cosscalar_theta - xx * sinscalar_theta + roi_center_h;
+        scalar_t x = yy * sinscalar_theta + xx * cosscalar_theta + roi_center_w;
+
+        scalar_t val = bilinear_interpolate<scalar_t>(
+            offset_bottom_data, height, width, y, x, index);
+        output_val += val;
+      }
+    }
+    output_val /= count;
+
+    top_data[index] = output_val;
+  }
+}
+
+/*** Backward ***/
+template <typename scalar_t>
+__global__ void roi_align_rotated_backward_cuda_kernel(
+    const int nthreads, const scalar_t *top_diff, const scalar_t *bottom_rois,
+    const scalar_t spatial_scale, const int sample_num, const bool aligned,
+    const bool clockwise, const int channels, const int height, const int width,
+    const int pooled_height, const int pooled_width, scalar_t *bottom_diff) {
+  CUDA_1D_KERNEL_LOOP(index, nthreads) {
+    // (n, c, ph, pw) is an element in the pooled output
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int c = (index / pooled_width / pooled_height) % channels;
+    int n = index / pooled_width / pooled_height / channels;
+
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 6;
+    int roi_batch_ind = offset_bottom_rois[0];
+
+    // Do not round
+    scalar_t offset = aligned ? (scalar_t)0.5 : (scalar_t)0.0;
+    scalar_t roi_center_w = offset_bottom_rois[1] * spatial_scale - offset;
+    scalar_t roi_center_h = offset_bottom_rois[2] * spatial_scale - offset;
+    scalar_t roi_width = offset_bottom_rois[3] * spatial_scale;
+    scalar_t roi_height = offset_bottom_rois[4] * spatial_scale;
+    // scalar_t theta = offset_bottom_rois[5] * M_PI / 180.0;
+    scalar_t theta = offset_bottom_rois[5];
+    if (clockwise) {
+      theta = -theta; // If clockwise, the angle needs to be reversed.
+    }
+    if (!aligned) { // for backward-compatibility only
+      // Force malformed ROIs to be 1x1
+      roi_width = max(roi_width, (scalar_t)1.);
+      roi_height = max(roi_height, (scalar_t)1.);
+    }
+    scalar_t bin_size_h = static_cast<scalar_t>(roi_height) /
+                          static_cast<scalar_t>(pooled_height);
+    scalar_t bin_size_w =
+        static_cast<scalar_t>(roi_width) / static_cast<scalar_t>(pooled_width);
+
+    scalar_t *offset_bottom_diff =
+        bottom_diff + (roi_batch_ind * channels + c) * height * width;
+
+    int top_offset = (n * channels + c) * pooled_height * pooled_width;
+    const scalar_t *offset_top_diff = top_diff + top_offset;
+    const scalar_t top_diff_this_bin = offset_top_diff[ph * pooled_width + pw];
+
+    // We use roi_bin_grid to sample the grid and mimic integral
+    int roi_bin_grid_h = (sample_num > 0)
+                             ? sample_num
+                             : ceilf(roi_height / pooled_height); // e.g., = 2
+    int roi_bin_grid_w =
+        (sample_num > 0) ? sample_num : ceilf(roi_width / pooled_width);
+
+    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
+    // Appropriate translation needs to be applied after.
+    scalar_t roi_start_h = -roi_height / 2.0;
+    scalar_t roi_start_w = -roi_width / 2.0;
+    scalar_t cosTheta = cos(theta);
+    scalar_t sinTheta = sin(theta);
+
+    // We do average (integral) pooling inside a bin
+    const scalar_t count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4
+
+    for (int iy = 0; iy < roi_bin_grid_h; iy++) { // e.g., iy = 0, 1
+      const scalar_t yy =
+          roi_start_h + ph * bin_size_h +
+          static_cast<scalar_t>(iy + .5f) * bin_size_h /
+              static_cast<scalar_t>(roi_bin_grid_h); // e.g., 0.5, 1.5
+      for (int ix = 0; ix < roi_bin_grid_w; ix++) {
+        const scalar_t xx = roi_start_w + pw * bin_size_w +
+                            static_cast<scalar_t>(ix + .5f) * bin_size_w /
+                                static_cast<scalar_t>(roi_bin_grid_w);
+
+        // Rotate by theta around the center and translate
+        scalar_t y = yy * cosTheta - xx * sinTheta + roi_center_h;
+        scalar_t x = yy * sinTheta + xx * cosTheta + roi_center_w;
+
+        scalar_t w1, w2, w3, w4;
+        int x_low, x_high, y_low, y_high;
+
+        bilinear_interpolate_gradient<scalar_t>(height, width, y, x, w1, w2, w3,
+                                                w4, x_low, x_high, y_low,
+                                                y_high, index);
+
+        scalar_t g1 = top_diff_this_bin * w1 / count;
+        scalar_t g2 = top_diff_this_bin * w2 / count;
+        scalar_t g3 = top_diff_this_bin * w3 / count;
+        scalar_t g4 = top_diff_this_bin * w4 / count;
+
+        if (x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0) {
+          atomicAdd(offset_bottom_diff + y_low * width + x_low, g1);
+          atomicAdd(offset_bottom_diff + y_low * width + x_high, g2);
+          atomicAdd(offset_bottom_diff + y_high * width + x_low, g3);
+          atomicAdd(offset_bottom_diff + y_high * width + x_high, g4);
+        } // if
+      }   // ix
+    }     // iy
+  }       // CUDA_1D_KERNEL_LOOP
+} // RoIAlignBackward
+
+std::vector<paddle::Tensor>
+RoIAlignRotatedCUDAForward(const paddle::Tensor &input,
+                           const paddle::Tensor &rois, int aligned_height,
+                           int aligned_width, float spatial_scale,
+                           int sampling_ratio, bool aligned, bool clockwise) {
+
+  auto num_rois = rois.shape()[0];
+
+  auto channels = input.shape()[1];
+  auto height = input.shape()[2];
+  auto width = input.shape()[3];
+
+  auto output =
+      paddle::empty({num_rois, channels, aligned_height, aligned_width},
+                    input.type(), paddle::GPUPlace());
+  auto output_size = output.numel();
+
+  PD_DISPATCH_FLOATING_TYPES(
+      input.type(), "roi_align_rotated_cuda_forward_kernel", ([&] {
+        roi_align_rotated_cuda_forward_kernel<data_t>
+            <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>>(
+                output_size, input.data<data_t>(), rois.data<data_t>(),
+                static_cast<data_t>(spatial_scale), sampling_ratio, aligned,
+                clockwise, channels, height, width, aligned_height,
+                aligned_width, output.data<data_t>());
+      }));
+
+  return {output};
+}
+
+std::vector<paddle::Tensor> RoIAlignRotatedCUDABackward(
+    const paddle::Tensor &input, const paddle::Tensor &rois,
+    const paddle::Tensor &grad_output, int aligned_height, int aligned_width,
+    float spatial_scale, int sampling_ratio, bool aligned, bool clockwise) {
+
+  auto num_rois = rois.shape()[0];
+
+  auto batch_size = input.shape()[0];
+  auto channels = input.shape()[1];
+  auto height = input.shape()[2];
+  auto width = input.shape()[3];
+
+  auto grad_input = paddle::full({batch_size, channels, height, width}, 0.0,
+                                 input.type(), paddle::GPUPlace());
+
+  const int output_size = num_rois * aligned_height * aligned_width * channels;
+
+  PD_DISPATCH_FLOATING_TYPES(
+      grad_output.type(), "roi_align_rotated_backward_cuda_kernel", ([&] {
+        roi_align_rotated_backward_cuda_kernel<data_t>
+            <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>>(
+                output_size, grad_output.data<data_t>(), rois.data<data_t>(),
+                spatial_scale, sampling_ratio, aligned, clockwise, channels,
+                height, width, aligned_height, aligned_width,
+                grad_input.data<data_t>());
+      }));
+  return {grad_input};
+}

ppocr/ext_op/roi_align_rotated/roi_align_rotated.py

@@ -0,0 +1,69 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmcv/blob/master/mmcv/ops/roi_align_rotated.py
+"""
+
+import paddle
+import paddle.nn as nn
+from paddle.utils.cpp_extension import load
+
+custom_ops = load(
+    name="custom_jit_ops",
+    sources=[
+        "ppocr/ext_op/roi_align_rotated/roi_align_rotated.cc",
+        "ppocr/ext_op/roi_align_rotated/roi_align_rotated.cu",
+    ],
+)
+
+roi_align_rotated = custom_ops.roi_align_rotated
+
+
+class RoIAlignRotated(nn.Layer):
+    """RoI align pooling layer for rotated proposals."""
+
+    def __init__(
+        self, out_size, spatial_scale, sample_num=0, aligned=True, clockwise=False
+    ):
+        super(RoIAlignRotated, self).__init__()
+
+        if isinstance(out_size, int):
+            self.out_h = out_size
+            self.out_w = out_size
+        elif isinstance(out_size, tuple):
+            assert len(out_size) == 2
+            assert isinstance(out_size[0], int)
+            assert isinstance(out_size[1], int)
+            self.out_h, self.out_w = out_size
+        else:
+            raise TypeError('"out_size" must be an integer or tuple of integers')
+
+        self.spatial_scale = float(spatial_scale)
+        self.sample_num = int(sample_num)
+        self.aligned = aligned
+        self.clockwise = clockwise
+
+    def forward(self, feats, rois):
+        output = roi_align_rotated(
+            feats,
+            rois,
+            self.out_h,
+            self.out_w,
+            self.spatial_scale,
+            self.sample_num,
+            self.aligned,
+            self.clockwise,
+        )
+        return output

ppocr/losses/__init__.py

@@ -0,0 +1,124 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+import copy
+import paddle
+import paddle.nn as nn
+
+# basic_loss
+from .basic_loss import LossFromOutput
+
+# det loss
+from .det_db_loss import DBLoss
+from .det_east_loss import EASTLoss
+from .det_sast_loss import SASTLoss
+from .det_pse_loss import PSELoss
+from .det_fce_loss import FCELoss
+from .det_ct_loss import CTLoss
+from .det_drrg_loss import DRRGLoss
+
+# rec loss
+from .rec_ctc_loss import CTCLoss
+from .rec_att_loss import AttentionLoss
+from .rec_srn_loss import SRNLoss
+from .rec_ce_loss import CELoss
+from .rec_sar_loss import SARLoss
+from .rec_aster_loss import AsterLoss
+from .rec_pren_loss import PRENLoss
+from .rec_multi_loss import MultiLoss
+from .rec_vl_loss import VLLoss
+from .rec_spin_att_loss import SPINAttentionLoss
+from .rec_rfl_loss import RFLLoss
+from .rec_can_loss import CANLoss
+from .rec_satrn_loss import SATRNLoss
+from .rec_nrtr_loss import NRTRLoss
+from .rec_parseq_loss import ParseQLoss
+from .rec_cppd_loss import CPPDLoss
+from .rec_latexocr_loss import LaTeXOCRLoss
+from .rec_unimernet_loss import UniMERNetLoss
+from .rec_ppformulanet_loss import PPFormulaNet_S_Loss, PPFormulaNet_L_Loss
+
+# cls loss
+from .cls_loss import ClsLoss
+
+# e2e loss
+from .e2e_pg_loss import PGLoss
+from .kie_sdmgr_loss import SDMGRLoss
+
+# basic loss function
+from .basic_loss import DistanceLoss
+
+# combined loss function
+from .combined_loss import CombinedLoss
+
+# table loss
+from .table_att_loss import TableAttentionLoss, SLALoss
+from .table_master_loss import TableMasterLoss
+
+# vqa token loss
+from .vqa_token_layoutlm_loss import VQASerTokenLayoutLMLoss
+
+# sr loss
+from .stroke_focus_loss import StrokeFocusLoss
+from .text_focus_loss import TelescopeLoss
+
+
+def build_loss(config):
+    support_dict = [
+        "DBLoss",
+        "PSELoss",
+        "EASTLoss",
+        "SASTLoss",
+        "FCELoss",
+        "CTCLoss",
+        "ClsLoss",
+        "AttentionLoss",
+        "SRNLoss",
+        "PGLoss",
+        "CombinedLoss",
+        "CELoss",
+        "TableAttentionLoss",
+        "SARLoss",
+        "AsterLoss",
+        "SDMGRLoss",
+        "VQASerTokenLayoutLMLoss",
+        "LossFromOutput",
+        "PRENLoss",
+        "MultiLoss",
+        "TableMasterLoss",
+        "SPINAttentionLoss",
+        "VLLoss",
+        "StrokeFocusLoss",
+        "SLALoss",
+        "CTLoss",
+        "RFLLoss",
+        "DRRGLoss",
+        "CANLoss",
+        "TelescopeLoss",
+        "SATRNLoss",
+        "NRTRLoss",
+        "ParseQLoss",
+        "CPPDLoss",
+        "LaTeXOCRLoss",
+        "UniMERNetLoss",
+        "PPFormulaNet_S_Loss",
+        "PPFormulaNet_L_Loss",
+    ]
+    config = copy.deepcopy(config)
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "loss only support {}".format(support_dict)
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/losses/ace_loss.py

@@ -0,0 +1,49 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+# This code is refer from: https://github.com/viig99/LS-ACELoss
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+import paddle.nn as nn
+
+
+class ACELoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.loss_func = nn.CrossEntropyLoss(
+            weight=None, ignore_index=0, reduction="none", soft_label=True, axis=-1
+        )
+
+    def __call__(self, predicts, batch):
+        if isinstance(predicts, (list, tuple)):
+            predicts = predicts[-1]
+
+        B, N = predicts.shape[:2]
+        div = paddle.to_tensor([N]).astype("float32")
+
+        predicts = nn.functional.softmax(predicts, axis=-1)
+        aggregation_preds = paddle.sum(predicts, axis=1)
+        aggregation_preds = paddle.divide(aggregation_preds, div)
+
+        length = batch[2].astype("float32")
+        batch = batch[3].astype("float32")
+        batch[:, 0] = paddle.subtract(div, length)
+        batch = paddle.divide(batch, div)
+
+        loss = self.loss_func(aggregation_preds, batch)
+        return {"loss_ace": loss}

ppocr/losses/basic_loss.py

@@ -0,0 +1,247 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+from paddle.nn import L1Loss
+from paddle.nn import MSELoss as L2Loss
+from paddle.nn import SmoothL1Loss
+
+
+class CELoss(nn.Layer):
+    def __init__(self, epsilon=None):
+        super().__init__()
+        if epsilon is not None and (epsilon <= 0 or epsilon >= 1):
+            epsilon = None
+        self.epsilon = epsilon
+
+    def _labelsmoothing(self, target, class_num):
+        if target.shape[-1] != class_num:
+            one_hot_target = F.one_hot(target, class_num)
+        else:
+            one_hot_target = target
+        soft_target = F.label_smooth(one_hot_target, epsilon=self.epsilon)
+        soft_target = paddle.reshape(soft_target, shape=[-1, class_num])
+        return soft_target
+
+    def forward(self, x, label):
+        loss_dict = {}
+        if self.epsilon is not None:
+            class_num = x.shape[-1]
+            label = self._labelsmoothing(label, class_num)
+            x = -F.log_softmax(x, axis=-1)
+            loss = paddle.sum(x * label, axis=-1)
+        else:
+            if label.shape[-1] == x.shape[-1]:
+                label = F.softmax(label, axis=-1)
+                soft_label = True
+            else:
+                soft_label = False
+            loss = F.cross_entropy(x, label=label, soft_label=soft_label)
+        return loss
+
+
+class KLJSLoss(object):
+    def __init__(self, mode="kl"):
+        assert mode in [
+            "kl",
+            "js",
+            "KL",
+            "JS",
+        ], "mode can only be one of ['kl', 'KL', 'js', 'JS']"
+        self.mode = mode
+
+    def __call__(self, p1, p2, reduction="mean", eps=1e-5):
+        if self.mode.lower() == "kl":
+            loss = paddle.multiply(p2, paddle.log((p2 + eps) / (p1 + eps) + eps))
+            loss += paddle.multiply(p1, paddle.log((p1 + eps) / (p2 + eps) + eps))
+            loss *= 0.5
+        elif self.mode.lower() == "js":
+            loss = paddle.multiply(
+                p2, paddle.log((2 * p2 + eps) / (p1 + p2 + eps) + eps)
+            )
+            loss += paddle.multiply(
+                p1, paddle.log((2 * p1 + eps) / (p1 + p2 + eps) + eps)
+            )
+            loss *= 0.5
+        else:
+            raise ValueError(
+                "The mode.lower() if KLJSLoss should be one of ['kl', 'js']"
+            )
+
+        if reduction == "mean":
+            loss = paddle.mean(loss, axis=[1, 2])
+        elif reduction == "none" or reduction is None:
+            return loss
+        else:
+            loss = paddle.sum(loss, axis=[1, 2])
+
+        return loss
+
+
+class DMLLoss(nn.Layer):
+    """
+    DMLLoss
+    """
+
+    def __init__(self, act=None, use_log=False):
+        super().__init__()
+        if act is not None:
+            assert act in ["softmax", "sigmoid"]
+        if act == "softmax":
+            self.act = nn.Softmax(axis=-1)
+        elif act == "sigmoid":
+            self.act = nn.Sigmoid()
+        else:
+            self.act = None
+
+        self.use_log = use_log
+        self.jskl_loss = KLJSLoss(mode="kl")
+
+    def _kldiv(self, x, target):
+        eps = 1.0e-10
+        loss = target * (paddle.log(target + eps) - x)
+        # batch mean loss
+        loss = paddle.sum(loss) / loss.shape[0]
+        return loss
+
+    def forward(self, out1, out2):
+        if self.act is not None:
+            out1 = self.act(out1) + 1e-10
+            out2 = self.act(out2) + 1e-10
+        if self.use_log:
+            # for recognition distillation, log is needed for feature map
+            log_out1 = paddle.log(out1)
+            log_out2 = paddle.log(out2)
+            loss = (self._kldiv(log_out1, out2) + self._kldiv(log_out2, out1)) / 2.0
+        else:
+            # for detection distillation log is not needed
+            loss = self.jskl_loss(out1, out2)
+        return loss
+
+
+class DistanceLoss(nn.Layer):
+    """
+    DistanceLoss:
+        mode: loss mode
+    """
+
+    def __init__(self, mode="l2", **kargs):
+        super().__init__()
+        assert mode in ["l1", "l2", "smooth_l1"]
+        if mode == "l1":
+            self.loss_func = nn.L1Loss(**kargs)
+        elif mode == "l2":
+            self.loss_func = nn.MSELoss(**kargs)
+        elif mode == "smooth_l1":
+            self.loss_func = nn.SmoothL1Loss(**kargs)
+
+    def forward(self, x, y):
+        return self.loss_func(x, y)
+
+
+class LossFromOutput(nn.Layer):
+    def __init__(self, key="loss", reduction="none"):
+        super().__init__()
+        self.key = key
+        self.reduction = reduction
+
+    def forward(self, predicts, batch):
+        loss = predicts
+        if self.key is not None and isinstance(predicts, dict):
+            loss = loss[self.key]
+        if self.reduction == "mean":
+            loss = paddle.mean(loss)
+        elif self.reduction == "sum":
+            loss = paddle.sum(loss)
+        return {"loss": loss}
+
+
+class KLDivLoss(nn.Layer):
+    """
+    KLDivLoss
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def _kldiv(self, x, target, mask=None):
+        eps = 1.0e-10
+        loss = target * (paddle.log(target + eps) - x)
+        if mask is not None:
+            loss = loss.flatten(0, 1).sum(axis=1)
+            loss = loss.masked_select(mask).mean()
+        else:
+            # batch mean loss
+            loss = paddle.sum(loss) / loss.shape[0]
+        return loss
+
+    def forward(self, logits_s, logits_t, mask=None):
+        log_out_s = F.log_softmax(logits_s, axis=-1)
+        out_t = F.softmax(logits_t, axis=-1)
+        loss = self._kldiv(log_out_s, out_t, mask)
+        return loss
+
+
+class DKDLoss(nn.Layer):
+    """
+    KLDivLoss
+    """
+
+    def __init__(self, temperature=1.0, alpha=1.0, beta=1.0):
+        super().__init__()
+        self.temperature = temperature
+        self.alpha = alpha
+        self.beta = beta
+
+    def _cat_mask(self, t, mask1, mask2):
+        t1 = (t * mask1).sum(axis=1, keepdim=True)
+        t2 = (t * mask2).sum(axis=1, keepdim=True)
+        rt = paddle.concat([t1, t2], axis=1)
+        return rt
+
+    def _kl_div(self, x, label, mask=None):
+        y = (label * (paddle.log(label + 1e-10) - x)).sum(axis=1)
+        if mask is not None:
+            y = y.masked_select(mask).mean()
+        else:
+            y = y.mean()
+        return y
+
+    def forward(self, logits_student, logits_teacher, target, mask=None):
+        gt_mask = F.one_hot(target.reshape([-1]), num_classes=logits_student.shape[-1])
+        other_mask = 1 - gt_mask
+        logits_student = logits_student.flatten(0, 1)
+        logits_teacher = logits_teacher.flatten(0, 1)
+        pred_student = F.softmax(logits_student / self.temperature, axis=1)
+        pred_teacher = F.softmax(logits_teacher / self.temperature, axis=1)
+        pred_student = self._cat_mask(pred_student, gt_mask, other_mask)
+        pred_teacher = self._cat_mask(pred_teacher, gt_mask, other_mask)
+        log_pred_student = paddle.log(pred_student)
+        tckd_loss = self._kl_div(log_pred_student, pred_teacher) * (self.temperature**2)
+        pred_teacher_part2 = F.softmax(
+            logits_teacher / self.temperature - 1000.0 * gt_mask, axis=1
+        )
+        log_pred_student_part2 = F.log_softmax(
+            logits_student / self.temperature - 1000.0 * gt_mask, axis=1
+        )
+        nckd_loss = self._kl_div(log_pred_student_part2, pred_teacher_part2) * (
+            self.temperature**2
+        )
+
+        loss = self.alpha * tckd_loss + self.beta * nckd_loss
+
+        return loss

ppocr/losses/center_loss.py

@@ -0,0 +1,89 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+# This code is refer from: https://github.com/KaiyangZhou/pytorch-center-loss
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+import os
+import pickle
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+
+class CenterLoss(nn.Layer):
+    """
+    Reference: Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
+    """
+
+    def __init__(self, num_classes=6625, feat_dim=96, center_file_path=None):
+        super().__init__()
+        self.num_classes = num_classes
+        self.feat_dim = feat_dim
+        self.centers = paddle.randn(shape=[self.num_classes, self.feat_dim]).astype(
+            "float64"
+        )
+
+        if center_file_path is not None:
+            assert os.path.exists(
+                center_file_path
+            ), f"center path({center_file_path}) must exist when it is not None."
+            with open(center_file_path, "rb") as f:
+                char_dict = pickle.load(f)
+                for key in char_dict.keys():
+                    self.centers[key] = paddle.to_tensor(char_dict[key])
+
+    def __call__(self, predicts, batch):
+        assert isinstance(predicts, (list, tuple))
+        features, predicts = predicts
+
+        feats_reshape = paddle.reshape(features, [-1, features.shape[-1]]).astype(
+            "float64"
+        )
+        label = paddle.argmax(predicts, axis=2)
+        label = paddle.reshape(label, [label.shape[0] * label.shape[1]])
+
+        batch_size = feats_reshape.shape[0]
+
+        # calc l2 distance between feats and centers
+        square_feat = paddle.sum(paddle.square(feats_reshape), axis=1, keepdim=True)
+        square_feat = paddle.expand(square_feat, [batch_size, self.num_classes])
+
+        square_center = paddle.sum(paddle.square(self.centers), axis=1, keepdim=True)
+        square_center = paddle.expand(
+            square_center, [self.num_classes, batch_size]
+        ).astype("float64")
+        square_center = paddle.transpose(square_center, [1, 0])
+
+        distmat = paddle.add(square_feat, square_center)
+        feat_dot_center = paddle.matmul(
+            feats_reshape, paddle.transpose(self.centers, [1, 0])
+        )
+        distmat = distmat - 2.0 * feat_dot_center
+
+        # generate the mask
+        classes = paddle.arange(self.num_classes).astype("int64")
+        label = paddle.expand(
+            paddle.unsqueeze(label, 1), (batch_size, self.num_classes)
+        )
+        mask = paddle.equal(
+            paddle.expand(classes, [batch_size, self.num_classes]), label
+        ).astype("float64")
+        dist = paddle.multiply(distmat, mask)
+
+        loss = paddle.sum(paddle.clip(dist, min=1e-12, max=1e12)) / batch_size
+        return {"loss_center": loss}

ppocr/losses/cls_loss.py

@@ -0,0 +1,30 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+
+
+class ClsLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(ClsLoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(reduction="mean")
+
+    def forward(self, predicts, batch):
+        label = batch[1].astype("int64")
+        loss = self.loss_func(input=predicts, label=label)
+        return {"loss": loss}

ppocr/losses/combined_loss.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+
+import paddle
+import paddle.nn as nn
+
+from .rec_ctc_loss import CTCLoss
+from .center_loss import CenterLoss
+from .ace_loss import ACELoss
+from .rec_sar_loss import SARLoss
+
+from .distillation_loss import DistillationCTCLoss, DistillCTCLogits
+from .distillation_loss import DistillationSARLoss, DistillationNRTRLoss
+from .distillation_loss import (
+    DistillationDMLLoss,
+    DistillationKLDivLoss,
+    DistillationDKDLoss,
+)
+from .distillation_loss import (
+    DistillationDistanceLoss,
+    DistillationDBLoss,
+    DistillationDilaDBLoss,
+)
+from .distillation_loss import (
+    DistillationVQASerTokenLayoutLMLoss,
+    DistillationSERDMLLoss,
+)
+from .distillation_loss import DistillationLossFromOutput
+from .distillation_loss import DistillationVQADistanceLoss
+
+
+class CombinedLoss(nn.Layer):
+    """
+    CombinedLoss:
+        a combionation of loss function
+    """
+
+    def __init__(self, loss_config_list=None):
+        super().__init__()
+        self.loss_func = []
+        self.loss_weight = []
+        assert isinstance(loss_config_list, list), "operator config should be a list"
+        for config in loss_config_list:
+            assert isinstance(config, dict) and len(config) == 1, "yaml format error"
+            name = list(config)[0]
+            param = config[name]
+            assert (
+                "weight" in param
+            ), "weight must be in param, but param just contains {}".format(
+                param.keys()
+            )
+            self.loss_weight.append(param.pop("weight"))
+            self.loss_func.append(eval(name)(**param))
+
+    def forward(self, input, batch, **kargs):
+        loss_dict = {}
+        loss_all = 0.0
+        for idx, loss_func in enumerate(self.loss_func):
+            loss = loss_func(input, batch, **kargs)
+            if isinstance(loss, paddle.Tensor):
+                loss = {"loss_{}_{}".format(str(loss), idx): loss}
+
+            weight = self.loss_weight[idx]
+
+            loss = {key: loss[key] * weight for key in loss}
+
+            if "loss" in loss:
+                loss_all += loss["loss"]
+            else:
+                loss_all += paddle.add_n(list(loss.values()))
+            loss_dict.update(loss)
+        loss_dict["loss"] = loss_all
+        return loss_dict

ppocr/losses/det_basic_loss.py

@@ -0,0 +1,161 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/models/losses/basic_loss.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+
+
+class BalanceLoss(nn.Layer):
+    def __init__(
+        self,
+        balance_loss=True,
+        main_loss_type="DiceLoss",
+        negative_ratio=3,
+        return_origin=False,
+        eps=1e-6,
+        **kwargs,
+    ):
+        """
+        The BalanceLoss for Differentiable Binarization text detection
+        args:
+            balance_loss (bool): whether balance loss or not, default is True
+            main_loss_type (str): can only be one of ['CrossEntropy','DiceLoss',
+                'Euclidean','BCELoss', 'MaskL1Loss'], default is  'DiceLoss'.
+            negative_ratio (int|float): float, default is 3.
+            return_origin (bool): whether return unbalanced loss or not, default is False.
+            eps (float): default is 1e-6.
+        """
+        super(BalanceLoss, self).__init__()
+        self.balance_loss = balance_loss
+        self.main_loss_type = main_loss_type
+        self.negative_ratio = negative_ratio
+        self.return_origin = return_origin
+        self.eps = eps
+
+        if self.main_loss_type == "CrossEntropy":
+            self.loss = nn.CrossEntropyLoss()
+        elif self.main_loss_type == "Euclidean":
+            self.loss = nn.MSELoss()
+        elif self.main_loss_type == "DiceLoss":
+            self.loss = DiceLoss(self.eps)
+        elif self.main_loss_type == "BCELoss":
+            self.loss = BCELoss(reduction="none")
+        elif self.main_loss_type == "MaskL1Loss":
+            self.loss = MaskL1Loss(self.eps)
+        else:
+            loss_type = [
+                "CrossEntropy",
+                "DiceLoss",
+                "Euclidean",
+                "BCELoss",
+                "MaskL1Loss",
+            ]
+            raise Exception(
+                "main_loss_type in BalanceLoss() can only be one of {}".format(
+                    loss_type
+                )
+            )
+
+    def forward(self, pred, gt, mask=None):
+        """
+        The BalanceLoss for Differentiable Binarization text detection
+        args:
+            pred (variable): predicted feature maps.
+            gt (variable): ground truth feature maps.
+            mask (variable): masked maps.
+        return: (variable) balanced loss
+        """
+        positive = gt * mask
+        negative = (1 - gt) * mask
+
+        positive_count = int(positive.sum())
+        negative_count = int(min(negative.sum(), positive_count * self.negative_ratio))
+        loss = self.loss(pred, gt, mask=mask)
+
+        if not self.balance_loss:
+            return loss
+
+        positive_loss = positive * loss
+        negative_loss = negative * loss
+        negative_loss = paddle.reshape(negative_loss, shape=[-1])
+        if negative_count > 0:
+            sort_loss = negative_loss.sort(descending=True)
+            negative_loss = sort_loss[:negative_count]
+            # negative_loss, _ = paddle.topk(negative_loss, k=negative_count_int)
+            balance_loss = (positive_loss.sum() + negative_loss.sum()) / (
+                positive_count + negative_count + self.eps
+            )
+        else:
+            balance_loss = positive_loss.sum() / (positive_count + self.eps)
+        if self.return_origin:
+            return balance_loss, loss
+
+        return balance_loss
+
+
+class DiceLoss(nn.Layer):
+    def __init__(self, eps=1e-6):
+        super(DiceLoss, self).__init__()
+        self.eps = eps
+
+    def forward(self, pred, gt, mask, weights=None):
+        """
+        DiceLoss function.
+        """
+
+        assert pred.shape == gt.shape
+        assert pred.shape == mask.shape
+        if weights is not None:
+            assert weights.shape == mask.shape
+            mask = weights * mask
+        intersection = paddle.sum(pred * gt * mask)
+
+        union = paddle.sum(pred * mask) + paddle.sum(gt * mask) + self.eps
+        loss = 1 - 2.0 * intersection / union
+        assert loss <= 1
+        return loss
+
+
+class MaskL1Loss(nn.Layer):
+    def __init__(self, eps=1e-6):
+        super(MaskL1Loss, self).__init__()
+        self.eps = eps
+
+    def forward(self, pred, gt, mask):
+        """
+        Mask L1 Loss
+        """
+        loss = (paddle.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps)
+        loss = paddle.mean(loss)
+        return loss
+
+
+class BCELoss(nn.Layer):
+    def __init__(self, reduction="mean"):
+        super(BCELoss, self).__init__()
+        self.reduction = reduction
+
+    def forward(self, input, label, mask=None, weight=None, name=None):
+        loss = F.binary_cross_entropy(input, label, reduction=self.reduction)
+        return loss

ppocr/losses/det_ct_loss.py

@@ -0,0 +1,302 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/shengtao96/CentripetalText/tree/main/models/loss
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+import numpy as np
+
+
+def ohem_single(score, gt_text, training_mask):
+    # online hard example mining
+
+    pos_num = int(paddle.sum(gt_text > 0.5)) - int(
+        paddle.sum((gt_text > 0.5) & (training_mask <= 0.5))
+    )
+
+    if pos_num == 0:
+        # selected_mask = gt_text.copy() * 0 # may be not good
+        selected_mask = training_mask
+        selected_mask = paddle.cast(
+            selected_mask.reshape((1, selected_mask.shape[0], selected_mask.shape[1])),
+            "float32",
+        )
+        return selected_mask
+
+    neg_num = int(paddle.sum((gt_text <= 0.5) & (training_mask > 0.5)))
+    neg_num = int(min(pos_num * 3, neg_num))
+
+    if neg_num == 0:
+        selected_mask = training_mask
+        selected_mask = paddle.cast(
+            selected_mask.reshape((1, selected_mask.shape[0], selected_mask.shape[1])),
+            "float32",
+        )
+        return selected_mask
+
+    # hard example
+    neg_score = score[(gt_text <= 0.5) & (training_mask > 0.5)]
+    neg_score_sorted = paddle.sort(-neg_score)
+    threshold = -neg_score_sorted[neg_num - 1]
+
+    selected_mask = ((score >= threshold) | (gt_text > 0.5)) & (training_mask > 0.5)
+    selected_mask = paddle.cast(
+        selected_mask.reshape((1, selected_mask.shape[0], selected_mask.shape[1])),
+        "float32",
+    )
+    return selected_mask
+
+
+def ohem_batch(scores, gt_texts, training_masks):
+    selected_masks = []
+    for i in range(scores.shape[0]):
+        selected_masks.append(
+            ohem_single(scores[i, :, :], gt_texts[i, :, :], training_masks[i, :, :])
+        )
+
+    selected_masks = paddle.cast(paddle.concat(selected_masks, 0), "float32")
+    return selected_masks
+
+
+def iou_single(a, b, mask, n_class):
+    EPS = 1e-6
+    valid = mask == 1
+    a = a[valid]
+    b = b[valid]
+    miou = []
+
+    # iou of each class
+    for i in range(n_class):
+        inter = paddle.cast(((a == i) & (b == i)), "float32")
+        union = paddle.cast(((a == i) | (b == i)), "float32")
+
+        miou.append(paddle.sum(inter) / (paddle.sum(union) + EPS))
+    miou = sum(miou) / len(miou)
+    return miou
+
+
+def iou(a, b, mask, n_class=2, reduce=True):
+    batch_size = a.shape[0]
+
+    a = a.reshape((batch_size, -1))
+    b = b.reshape((batch_size, -1))
+    mask = mask.reshape((batch_size, -1))
+
+    iou = paddle.zeros((batch_size,), dtype="float32")
+    for i in range(batch_size):
+        iou[i] = iou_single(a[i], b[i], mask[i], n_class)
+
+    if reduce:
+        iou = paddle.mean(iou)
+    return iou
+
+
+class DiceLoss(nn.Layer):
+    def __init__(self, loss_weight=1.0):
+        super(DiceLoss, self).__init__()
+        self.loss_weight = loss_weight
+
+    def forward(self, input, target, mask, reduce=True):
+        batch_size = input.shape[0]
+        input = F.sigmoid(input)  # scale to 0-1
+
+        input = input.reshape((batch_size, -1))
+        target = paddle.cast(target.reshape((batch_size, -1)), "float32")
+        mask = paddle.cast(mask.reshape((batch_size, -1)), "float32")
+
+        input = input * mask
+        target = target * mask
+
+        a = paddle.sum(input * target, axis=1)
+        b = paddle.sum(input * input, axis=1) + 0.001
+        c = paddle.sum(target * target, axis=1) + 0.001
+        d = (2 * a) / (b + c)
+        loss = 1 - d
+
+        loss = self.loss_weight * loss
+
+        if reduce:
+            loss = paddle.mean(loss)
+
+        return loss
+
+
+class SmoothL1Loss(nn.Layer):
+    def __init__(self, beta=1.0, loss_weight=1.0):
+        super(SmoothL1Loss, self).__init__()
+        self.beta = beta
+        self.loss_weight = loss_weight
+
+        np_coord = np.zeros(shape=[640, 640, 2], dtype=np.int64)
+        for i in range(640):
+            for j in range(640):
+                np_coord[i, j, 0] = j
+                np_coord[i, j, 1] = i
+        np_coord = np_coord.reshape((-1, 2))
+
+        self.coord = self.create_parameter(
+            shape=[640 * 640, 2],
+            dtype="int32",  # NOTE: not support "int64" before paddle 2.3.1
+            default_initializer=nn.initializer.Assign(value=np_coord),
+        )
+        self.coord.stop_gradient = True
+
+    def forward_single(self, input, target, mask, beta=1.0, eps=1e-6):
+        batch_size = input.shape[0]
+
+        diff = paddle.abs(input - target) * mask.unsqueeze(1)
+        loss = paddle.where(diff < beta, 0.5 * diff * diff / beta, diff - 0.5 * beta)
+        loss = paddle.cast(loss.reshape((batch_size, -1)), "float32")
+        mask = paddle.cast(mask.reshape((batch_size, -1)), "float32")
+        loss = paddle.sum(loss, axis=-1)
+        loss = loss / (mask.sum(axis=-1) + eps)
+
+        return loss
+
+    def select_single(self, distance, gt_instance, gt_kernel_instance, training_mask):
+        with paddle.no_grad():
+            # paddle 2.3.1, paddle.slice not support:
+            # distance[:, self.coord[:, 1], self.coord[:, 0]]
+            select_distance_list = []
+            for i in range(2):
+                tmp1 = distance[i, :]
+                tmp2 = tmp1[self.coord[:, 1], self.coord[:, 0]]
+                select_distance_list.append(tmp2.unsqueeze(0))
+            select_distance = paddle.concat(select_distance_list, axis=0)
+
+            off_points = paddle.cast(
+                self.coord, "float32"
+            ) + 10 * select_distance.transpose((1, 0))
+
+            off_points = paddle.cast(off_points, "int64")
+            off_points = paddle.clip(off_points, 0, distance.shape[-1] - 1)
+
+            selected_mask = (
+                gt_instance[self.coord[:, 1], self.coord[:, 0]]
+                != gt_kernel_instance[off_points[:, 1], off_points[:, 0]]
+            )
+            selected_mask = paddle.cast(
+                selected_mask.reshape((1, -1, distance.shape[-1])), "int64"
+            )
+            selected_training_mask = selected_mask * training_mask
+
+            return selected_training_mask
+
+    def forward(
+        self,
+        distances,
+        gt_instances,
+        gt_kernel_instances,
+        training_masks,
+        gt_distances,
+        reduce=True,
+    ):
+        selected_training_masks = []
+        for i in range(distances.shape[0]):
+            selected_training_masks.append(
+                self.select_single(
+                    distances[i, :, :, :],
+                    gt_instances[i, :, :],
+                    gt_kernel_instances[i, :, :],
+                    training_masks[i, :, :],
+                )
+            )
+        selected_training_masks = paddle.cast(
+            paddle.concat(selected_training_masks, 0), "float32"
+        )
+
+        loss = self.forward_single(
+            distances, gt_distances, selected_training_masks, self.beta
+        )
+        loss = self.loss_weight * loss
+
+        with paddle.no_grad():
+            batch_size = distances.shape[0]
+            false_num = selected_training_masks.reshape((batch_size, -1))
+            false_num = false_num.sum(axis=-1)
+            total_num = paddle.cast(training_masks.reshape((batch_size, -1)), "float32")
+            total_num = total_num.sum(axis=-1)
+            iou_text = (total_num - false_num) / (total_num + 1e-6)
+
+        if reduce:
+            loss = paddle.mean(loss)
+
+        return loss, iou_text
+
+
+class CTLoss(nn.Layer):
+    def __init__(self):
+        super(CTLoss, self).__init__()
+        self.kernel_loss = DiceLoss()
+        self.loc_loss = SmoothL1Loss(beta=0.1, loss_weight=0.05)
+
+    def forward(self, preds, batch):
+        imgs = batch[0]
+        out = preds["maps"]
+        (
+            gt_kernels,
+            training_masks,
+            gt_instances,
+            gt_kernel_instances,
+            training_mask_distances,
+            gt_distances,
+        ) = batch[1:]
+
+        kernels = out[:, 0, :, :]
+        distances = out[:, 1:, :, :]
+
+        # kernel loss
+        selected_masks = ohem_batch(kernels, gt_kernels, training_masks)
+
+        loss_kernel = self.kernel_loss(
+            kernels, gt_kernels, selected_masks, reduce=False
+        )
+
+        iou_kernel = iou(
+            paddle.cast((kernels > 0), "int64"),
+            gt_kernels,
+            training_masks,
+            reduce=False,
+        )
+        losses = dict(
+            loss_kernels=loss_kernel,
+        )
+
+        # loc loss
+        loss_loc, iou_text = self.loc_loss(
+            distances,
+            gt_instances,
+            gt_kernel_instances,
+            training_mask_distances,
+            gt_distances,
+            reduce=False,
+        )
+        losses.update(
+            dict(
+                loss_loc=loss_loc,
+            )
+        )
+
+        loss_all = loss_kernel + loss_loc
+        losses = {"loss": loss_all}
+
+        return losses

ppocr/losses/det_db_loss.py

@@ -0,0 +1,99 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/WenmuZhou/DBNet.pytorch/blob/master/models/losses/DB_loss.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+from .det_basic_loss import BalanceLoss, MaskL1Loss, DiceLoss
+
+
+class DBLoss(nn.Layer):
+    """
+    Differentiable Binarization (DB) Loss Function
+    args:
+        param (dict): the super parameter for DB Loss
+    """
+
+    def __init__(
+        self,
+        balance_loss=True,
+        main_loss_type="DiceLoss",
+        alpha=5,
+        beta=10,
+        ohem_ratio=3,
+        eps=1e-6,
+        **kwargs,
+    ):
+        super(DBLoss, self).__init__()
+        self.alpha = alpha
+        self.beta = beta
+        self.dice_loss = DiceLoss(eps=eps)
+        self.l1_loss = MaskL1Loss(eps=eps)
+        self.bce_loss = BalanceLoss(
+            balance_loss=balance_loss,
+            main_loss_type=main_loss_type,
+            negative_ratio=ohem_ratio,
+        )
+
+    def forward(self, predicts, labels):
+        predict_maps = predicts["maps"]
+        (
+            label_threshold_map,
+            label_threshold_mask,
+            label_shrink_map,
+            label_shrink_mask,
+        ) = labels[1:]
+        shrink_maps = predict_maps[:, 0, :, :]
+        threshold_maps = predict_maps[:, 1, :, :]
+        binary_maps = predict_maps[:, 2, :, :]
+
+        loss_shrink_maps = self.bce_loss(
+            shrink_maps, label_shrink_map, label_shrink_mask
+        )
+        loss_threshold_maps = self.l1_loss(
+            threshold_maps, label_threshold_map, label_threshold_mask
+        )
+        loss_binary_maps = self.dice_loss(
+            binary_maps, label_shrink_map, label_shrink_mask
+        )
+        loss_shrink_maps = self.alpha * loss_shrink_maps
+        loss_threshold_maps = self.beta * loss_threshold_maps
+        # CBN loss
+        if "distance_maps" in predicts.keys():
+            distance_maps = predicts["distance_maps"]
+            cbn_maps = predicts["cbn_maps"]
+            cbn_loss = self.bce_loss(
+                cbn_maps[:, 0, :, :], label_shrink_map, label_shrink_mask
+            )
+        else:
+            dis_loss = paddle.to_tensor([0.0])
+            cbn_loss = paddle.to_tensor([0.0])
+
+        loss_all = loss_shrink_maps + loss_threshold_maps + loss_binary_maps
+        losses = {
+            "loss": loss_all + cbn_loss,
+            "loss_shrink_maps": loss_shrink_maps,
+            "loss_threshold_maps": loss_threshold_maps,
+            "loss_binary_maps": loss_binary_maps,
+            "loss_cbn": cbn_loss,
+        }
+        return losses

ppocr/losses/det_drrg_loss.py

@@ -0,0 +1,234 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textdet/losses/drrg_loss.py
+"""
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+
+
+class DRRGLoss(nn.Layer):
+    def __init__(self, ohem_ratio=3.0):
+        super().__init__()
+        self.ohem_ratio = ohem_ratio
+        self.downsample_ratio = 1.0
+
+    def balance_bce_loss(self, pred, gt, mask):
+        """Balanced Binary-CrossEntropy Loss.
+
+        Args:
+            pred (Tensor): Shape of :math:`(1, H, W)`.
+            gt (Tensor): Shape of :math:`(1, H, W)`.
+            mask (Tensor): Shape of :math:`(1, H, W)`.
+
+        Returns:
+            Tensor: Balanced bce loss.
+        """
+        assert pred.shape == gt.shape == mask.shape
+        assert paddle.all(pred >= 0) and paddle.all(pred <= 1)
+        assert paddle.all(gt >= 0) and paddle.all(gt <= 1)
+        positive = gt * mask
+        negative = (1 - gt) * mask
+        positive_count = int(positive.sum())
+
+        if positive_count > 0:
+            loss = F.binary_cross_entropy(pred, gt, reduction="none")
+            positive_loss = paddle.sum(loss * positive)
+            negative_loss = loss * negative
+            negative_count = min(
+                int(negative.sum()), int(positive_count * self.ohem_ratio)
+            )
+        else:
+            positive_loss = paddle.to_tensor(0.0)
+            loss = F.binary_cross_entropy(pred, gt, reduction="none")
+            negative_loss = loss * negative
+            negative_count = 100
+        negative_loss, _ = paddle.topk(negative_loss.reshape([-1]), negative_count)
+
+        balance_loss = (positive_loss + paddle.sum(negative_loss)) / (
+            float(positive_count + negative_count) + 1e-5
+        )
+
+        return balance_loss
+
+    def gcn_loss(self, gcn_data):
+        """CrossEntropy Loss from gcn module.
+
+        Args:
+            gcn_data (tuple(Tensor, Tensor)): The first is the
+                prediction with shape :math:`(N, 2)` and the
+                second is the gt label with shape :math:`(m, n)`
+                where :math:`m * n = N`.
+
+        Returns:
+            Tensor: CrossEntropy loss.
+        """
+        gcn_pred, gt_labels = gcn_data
+        gt_labels = gt_labels.reshape([-1])
+        loss = F.cross_entropy(gcn_pred, gt_labels)
+
+        return loss
+
+    def bitmasks2tensor(self, bitmasks, target_sz):
+        """Convert Bitmasks to tensor.
+
+        Args:
+            bitmasks (list[BitmapMasks]): The BitmapMasks list. Each item is
+                for one img.
+            target_sz (tuple(int, int)): The target tensor of size
+                :math:`(H, W)`.
+
+        Returns:
+            list[Tensor]: The list of kernel tensors. Each element stands for
+            one kernel level.
+        """
+        batch_size = len(bitmasks)
+        results = []
+
+        kernel = []
+        for batch_inx in range(batch_size):
+            mask = bitmasks[batch_inx]
+            # hxw
+            mask_sz = mask.shape
+            # left, right, top, bottom
+            pad = [0, target_sz[1] - mask_sz[1], 0, target_sz[0] - mask_sz[0]]
+            mask = F.pad(mask, pad, mode="constant", value=0)
+            kernel.append(mask)
+        kernel = paddle.stack(kernel)
+        results.append(kernel)
+
+        return results
+
+    def forward(self, preds, labels):
+        """Compute Drrg loss."""
+
+        assert isinstance(preds, tuple)
+        (
+            gt_text_mask,
+            gt_center_region_mask,
+            gt_mask,
+            gt_top_height_map,
+            gt_bot_height_map,
+            gt_sin_map,
+            gt_cos_map,
+        ) = labels[1:8]
+
+        downsample_ratio = self.downsample_ratio
+
+        pred_maps, gcn_data = preds
+        pred_text_region = pred_maps[:, 0, :, :]
+        pred_center_region = pred_maps[:, 1, :, :]
+        pred_sin_map = pred_maps[:, 2, :, :]
+        pred_cos_map = pred_maps[:, 3, :, :]
+        pred_top_height_map = pred_maps[:, 4, :, :]
+        pred_bot_height_map = pred_maps[:, 5, :, :]
+        feature_sz = pred_maps.shape
+
+        # bitmask 2 tensor
+        mapping = {
+            "gt_text_mask": paddle.cast(gt_text_mask, "float32"),
+            "gt_center_region_mask": paddle.cast(gt_center_region_mask, "float32"),
+            "gt_mask": paddle.cast(gt_mask, "float32"),
+            "gt_top_height_map": paddle.cast(gt_top_height_map, "float32"),
+            "gt_bot_height_map": paddle.cast(gt_bot_height_map, "float32"),
+            "gt_sin_map": paddle.cast(gt_sin_map, "float32"),
+            "gt_cos_map": paddle.cast(gt_cos_map, "float32"),
+        }
+        gt = {}
+        for key, value in mapping.items():
+            gt[key] = value
+            if abs(downsample_ratio - 1.0) < 1e-2:
+                gt[key] = self.bitmasks2tensor(gt[key], feature_sz[2:])
+            else:
+                gt[key] = [item.rescale(downsample_ratio) for item in gt[key]]
+                gt[key] = self.bitmasks2tensor(gt[key], feature_sz[2:])
+                if key in ["gt_top_height_map", "gt_bot_height_map"]:
+                    gt[key] = [item * downsample_ratio for item in gt[key]]
+            gt[key] = [item for item in gt[key]]
+
+        scale = paddle.sqrt(1.0 / (pred_sin_map**2 + pred_cos_map**2 + 1e-8))
+        pred_sin_map = pred_sin_map * scale
+        pred_cos_map = pred_cos_map * scale
+
+        loss_text = self.balance_bce_loss(
+            F.sigmoid(pred_text_region), gt["gt_text_mask"][0], gt["gt_mask"][0]
+        )
+
+        text_mask = gt["gt_text_mask"][0] * gt["gt_mask"][0]
+        negative_text_mask = (1 - gt["gt_text_mask"][0]) * gt["gt_mask"][0]
+        loss_center_map = F.binary_cross_entropy(
+            F.sigmoid(pred_center_region),
+            gt["gt_center_region_mask"][0],
+            reduction="none",
+        )
+        if int(text_mask.sum()) > 0:
+            loss_center_positive = paddle.sum(loss_center_map * text_mask) / paddle.sum(
+                text_mask
+            )
+        else:
+            loss_center_positive = paddle.to_tensor(0.0)
+        loss_center_negative = paddle.sum(
+            loss_center_map * negative_text_mask
+        ) / paddle.sum(negative_text_mask)
+        loss_center = loss_center_positive + 0.5 * loss_center_negative
+
+        center_mask = gt["gt_center_region_mask"][0] * gt["gt_mask"][0]
+        if int(center_mask.sum()) > 0:
+            map_sz = pred_top_height_map.shape
+            ones = paddle.ones(map_sz, dtype="float32")
+            loss_top = F.smooth_l1_loss(
+                pred_top_height_map / (gt["gt_top_height_map"][0] + 1e-2),
+                ones,
+                reduction="none",
+            )
+            loss_bot = F.smooth_l1_loss(
+                pred_bot_height_map / (gt["gt_bot_height_map"][0] + 1e-2),
+                ones,
+                reduction="none",
+            )
+            gt_height = gt["gt_top_height_map"][0] + gt["gt_bot_height_map"][0]
+            loss_height = paddle.sum(
+                (paddle.log(gt_height + 1) * (loss_top + loss_bot)) * center_mask
+            ) / paddle.sum(center_mask)
+
+            loss_sin = paddle.sum(
+                F.smooth_l1_loss(pred_sin_map, gt["gt_sin_map"][0], reduction="none")
+                * center_mask
+            ) / paddle.sum(center_mask)
+            loss_cos = paddle.sum(
+                F.smooth_l1_loss(pred_cos_map, gt["gt_cos_map"][0], reduction="none")
+                * center_mask
+            ) / paddle.sum(center_mask)
+        else:
+            loss_height = paddle.to_tensor(0.0)
+            loss_sin = paddle.to_tensor(0.0)
+            loss_cos = paddle.to_tensor(0.0)
+
+        loss_gcn = self.gcn_loss(gcn_data)
+
+        loss = loss_text + loss_center + loss_height + loss_sin + loss_cos + loss_gcn
+        results = dict(
+            loss=loss,
+            loss_text=loss_text,
+            loss_center=loss_center,
+            loss_height=loss_height,
+            loss_sin=loss_sin,
+            loss_cos=loss_cos,
+            loss_gcn=loss_gcn,
+        )
+
+        return results

ppocr/losses/det_east_loss.py

@@ -0,0 +1,62 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+from .det_basic_loss import DiceLoss
+
+
+class EASTLoss(nn.Layer):
+    """ """
+
+    def __init__(self, eps=1e-6, **kwargs):
+        super(EASTLoss, self).__init__()
+        self.dice_loss = DiceLoss(eps=eps)
+
+    def forward(self, predicts, labels):
+        l_score, l_geo, l_mask = labels[1:]
+        f_score = predicts["f_score"]
+        f_geo = predicts["f_geo"]
+
+        dice_loss = self.dice_loss(f_score, l_score, l_mask)
+
+        # smoooth_l1_loss
+        channels = 8
+        l_geo_split = paddle.split(l_geo, num_or_sections=channels + 1, axis=1)
+        f_geo_split = paddle.split(f_geo, num_or_sections=channels, axis=1)
+        smooth_l1 = 0
+        for i in range(0, channels):
+            geo_diff = l_geo_split[i] - f_geo_split[i]
+            abs_geo_diff = paddle.abs(geo_diff)
+            smooth_l1_sign = paddle.less_than(abs_geo_diff, l_score)
+            smooth_l1_sign = paddle.cast(smooth_l1_sign, dtype="float32")
+            in_loss = abs_geo_diff * abs_geo_diff * smooth_l1_sign + (
+                abs_geo_diff - 0.5
+            ) * (1.0 - smooth_l1_sign)
+            out_loss = l_geo_split[-1] / channels * in_loss * l_score
+            smooth_l1 += out_loss
+        smooth_l1_loss = paddle.mean(smooth_l1 * l_score)
+
+        dice_loss = dice_loss * 0.01
+        total_loss = dice_loss + smooth_l1_loss
+        losses = {
+            "loss": total_loss,
+            "dice_loss": dice_loss,
+            "smooth_l1_loss": smooth_l1_loss,
+        }
+        return losses

ppocr/losses/det_fce_loss.py

@@ -0,0 +1,240 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textdet/losses/fce_loss.py
+"""
+
+import numpy as np
+from paddle import nn
+import paddle
+import paddle.nn.functional as F
+from functools import partial
+
+
+def multi_apply(func, *args, **kwargs):
+    pfunc = partial(func, **kwargs) if kwargs else func
+    map_results = map(pfunc, *args)
+    return tuple(map(list, zip(*map_results)))
+
+
+class FCELoss(nn.Layer):
+    """The class for implementing FCENet loss
+    FCENet(CVPR2021): Fourier Contour Embedding for Arbitrary-shaped
+        Text Detection
+
+    [https://arxiv.org/abs/2104.10442]
+
+    Args:
+        fourier_degree (int) : The maximum Fourier transform degree k.
+        num_sample (int) : The sampling points number of regression
+            loss. If it is too small, fcenet tends to be overfitting.
+        ohem_ratio (float): the negative/positive ratio in OHEM.
+    """
+
+    def __init__(self, fourier_degree, num_sample, ohem_ratio=3.0):
+        super().__init__()
+        self.fourier_degree = fourier_degree
+        self.num_sample = num_sample
+        self.ohem_ratio = ohem_ratio
+
+    def forward(self, preds, labels):
+        assert isinstance(preds, dict)
+        preds = preds["levels"]
+
+        p3_maps, p4_maps, p5_maps = labels[1:]
+        assert (
+            p3_maps[0].shape[0] == 4 * self.fourier_degree + 5
+        ), "fourier degree not equal in FCEhead and FCEtarget"
+
+        # to tensor
+        gts = [p3_maps, p4_maps, p5_maps]
+        for idx, maps in enumerate(gts):
+            gts[idx] = paddle.to_tensor(np.stack(maps))
+
+        losses = multi_apply(self.forward_single, preds, gts)
+
+        loss_tr = paddle.to_tensor(0.0).astype("float32")
+        loss_tcl = paddle.to_tensor(0.0).astype("float32")
+        loss_reg_x = paddle.to_tensor(0.0).astype("float32")
+        loss_reg_y = paddle.to_tensor(0.0).astype("float32")
+        loss_all = paddle.to_tensor(0.0).astype("float32")
+
+        for idx, loss in enumerate(losses):
+            loss_all += sum(loss)
+            if idx == 0:
+                loss_tr += sum(loss)
+            elif idx == 1:
+                loss_tcl += sum(loss)
+            elif idx == 2:
+                loss_reg_x += sum(loss)
+            else:
+                loss_reg_y += sum(loss)
+
+        results = dict(
+            loss=loss_all,
+            loss_text=loss_tr,
+            loss_center=loss_tcl,
+            loss_reg_x=loss_reg_x,
+            loss_reg_y=loss_reg_y,
+        )
+        return results
+
+    def forward_single(self, pred, gt):
+        cls_pred = paddle.transpose(pred[0], (0, 2, 3, 1))
+        reg_pred = paddle.transpose(pred[1], (0, 2, 3, 1))
+        gt = paddle.transpose(gt, (0, 2, 3, 1))
+
+        k = 2 * self.fourier_degree + 1
+        tr_pred = paddle.reshape(cls_pred[:, :, :, :2], (-1, 2))
+        tcl_pred = paddle.reshape(cls_pred[:, :, :, 2:], (-1, 2))
+        x_pred = paddle.reshape(reg_pred[:, :, :, 0:k], (-1, k))
+        y_pred = paddle.reshape(reg_pred[:, :, :, k : 2 * k], (-1, k))
+
+        tr_mask = gt[:, :, :, :1].reshape([-1])
+        tcl_mask = gt[:, :, :, 1:2].reshape([-1])
+        train_mask = gt[:, :, :, 2:3].reshape([-1])
+        x_map = paddle.reshape(gt[:, :, :, 3 : 3 + k], (-1, k))
+        y_map = paddle.reshape(gt[:, :, :, 3 + k :], (-1, k))
+
+        tr_train_mask = (train_mask * tr_mask).astype("bool")
+        tr_train_mask2 = paddle.concat(
+            [tr_train_mask.unsqueeze(1), tr_train_mask.unsqueeze(1)], axis=1
+        )
+        # tr loss
+        loss_tr = self.ohem(tr_pred, tr_mask, train_mask)
+        # tcl loss
+        loss_tcl = paddle.to_tensor(0.0).astype("float32")
+        tr_neg_mask = tr_train_mask.logical_not()
+        tr_neg_mask2 = paddle.concat(
+            [tr_neg_mask.unsqueeze(1), tr_neg_mask.unsqueeze(1)], axis=1
+        )
+        if tr_train_mask.sum().item() > 0:
+            loss_tcl_pos = F.cross_entropy(
+                tcl_pred.masked_select(tr_train_mask2).reshape([-1, 2]),
+                tcl_mask.masked_select(tr_train_mask).astype("int64"),
+            )
+            loss_tcl_neg = F.cross_entropy(
+                tcl_pred.masked_select(tr_neg_mask2).reshape([-1, 2]),
+                tcl_mask.masked_select(tr_neg_mask).astype("int64"),
+            )
+            loss_tcl = loss_tcl_pos + 0.5 * loss_tcl_neg
+
+        # regression loss
+        loss_reg_x = paddle.to_tensor(0.0).astype("float32")
+        loss_reg_y = paddle.to_tensor(0.0).astype("float32")
+        if tr_train_mask.sum().item() > 0:
+            weight = (
+                tr_mask.masked_select(tr_train_mask.astype("bool")).astype("float32")
+                + tcl_mask.masked_select(tr_train_mask.astype("bool")).astype("float32")
+            ) / 2
+            weight = weight.reshape([-1, 1])
+
+            ft_x, ft_y = self.fourier2poly(x_map, y_map)
+            ft_x_pre, ft_y_pre = self.fourier2poly(x_pred, y_pred)
+
+            dim = ft_x.shape[1]
+
+            tr_train_mask3 = paddle.concat(
+                [tr_train_mask.unsqueeze(1) for i in range(dim)], axis=1
+            )
+
+            loss_reg_x = paddle.mean(
+                weight
+                * F.smooth_l1_loss(
+                    ft_x_pre.masked_select(tr_train_mask3).reshape([-1, dim]),
+                    ft_x.masked_select(tr_train_mask3).reshape([-1, dim]),
+                    reduction="none",
+                )
+            )
+            loss_reg_y = paddle.mean(
+                weight
+                * F.smooth_l1_loss(
+                    ft_y_pre.masked_select(tr_train_mask3).reshape([-1, dim]),
+                    ft_y.masked_select(tr_train_mask3).reshape([-1, dim]),
+                    reduction="none",
+                )
+            )
+
+        return loss_tr, loss_tcl, loss_reg_x, loss_reg_y
+
+    def ohem(self, predict, target, train_mask):
+        pos = (target * train_mask).astype("bool")
+        neg = ((1 - target) * train_mask).astype("bool")
+
+        pos2 = paddle.concat([pos.unsqueeze(1), pos.unsqueeze(1)], axis=1)
+        neg2 = paddle.concat([neg.unsqueeze(1), neg.unsqueeze(1)], axis=1)
+
+        n_pos = pos.astype("float32").sum()
+
+        if n_pos.item() > 0:
+            loss_pos = F.cross_entropy(
+                predict.masked_select(pos2).reshape([-1, 2]),
+                target.masked_select(pos).astype("int64"),
+                reduction="sum",
+            )
+            loss_neg = F.cross_entropy(
+                predict.masked_select(neg2).reshape([-1, 2]),
+                target.masked_select(neg).astype("int64"),
+                reduction="none",
+            )
+            n_neg = min(
+                int(neg.astype("float32").sum().item()),
+                int(self.ohem_ratio * n_pos.astype("float32")),
+            )
+        else:
+            loss_pos = paddle.to_tensor(0.0)
+            loss_neg = F.cross_entropy(
+                predict.masked_select(neg2).reshape([-1, 2]),
+                target.masked_select(neg).astype("int64"),
+                reduction="none",
+            )
+            n_neg = 100
+        if len(loss_neg) > n_neg:
+            loss_neg, _ = paddle.topk(loss_neg, n_neg)
+
+        return (loss_pos + loss_neg.sum()) / (n_pos + n_neg).astype("float32")
+
+    def fourier2poly(self, real_maps, imag_maps):
+        """Transform Fourier coefficient maps to polygon maps.
+
+        Args:
+            real_maps (tensor): A map composed of the real parts of the
+                Fourier coefficients, whose shape is (-1, 2k+1)
+            imag_maps (tensor):A map composed of the imag parts of the
+                Fourier coefficients, whose shape is (-1, 2k+1)
+
+        Returns
+            x_maps (tensor): A map composed of the x value of the polygon
+                represented by n sample points (xn, yn), whose shape is (-1, n)
+            y_maps (tensor): A map composed of the y value of the polygon
+                represented by n sample points (xn, yn), whose shape is (-1, n)
+        """
+
+        k_vect = paddle.arange(
+            -self.fourier_degree, self.fourier_degree + 1, dtype="float32"
+        ).reshape([-1, 1])
+        i_vect = paddle.arange(0, self.num_sample, dtype="float32").reshape([1, -1])
+
+        transform_matrix = 2 * np.pi / self.num_sample * paddle.matmul(k_vect, i_vect)
+
+        x1 = paddle.einsum("ak, kn-> an", real_maps, paddle.cos(transform_matrix))
+        x2 = paddle.einsum("ak, kn-> an", imag_maps, paddle.sin(transform_matrix))
+        y1 = paddle.einsum("ak, kn-> an", real_maps, paddle.sin(transform_matrix))
+        y2 = paddle.einsum("ak, kn-> an", imag_maps, paddle.cos(transform_matrix))
+
+        x_maps = x1 - x2
+        y_maps = y1 + y2
+
+        return x_maps, y_maps

ppocr/losses/det_pse_loss.py

@@ -0,0 +1,158 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/whai362/PSENet/blob/python3/models/head/psenet_head.py
+"""
+
+import paddle
+from paddle import nn
+from paddle.nn import functional as F
+import numpy as np
+from ppocr.utils.iou import iou
+
+
+class PSELoss(nn.Layer):
+    def __init__(
+        self,
+        alpha,
+        ohem_ratio=3,
+        kernel_sample_mask="pred",
+        reduction="sum",
+        eps=1e-6,
+        **kwargs,
+    ):
+        """Implement PSE Loss."""
+        super(PSELoss, self).__init__()
+        assert reduction in ["sum", "mean", "none"]
+        self.alpha = alpha
+        self.ohem_ratio = ohem_ratio
+        self.kernel_sample_mask = kernel_sample_mask
+        self.reduction = reduction
+        self.eps = eps
+
+    def forward(self, outputs, labels):
+        predicts = outputs["maps"]
+        predicts = F.interpolate(predicts, scale_factor=4)
+
+        texts = predicts[:, 0, :, :]
+        kernels = predicts[:, 1:, :, :]
+        gt_texts, gt_kernels, training_masks = labels[1:]
+
+        # text loss
+        selected_masks = self.ohem_batch(texts, gt_texts, training_masks)
+
+        loss_text = self.dice_loss(texts, gt_texts, selected_masks)
+        iou_text = iou(
+            (texts > 0).astype("int64"), gt_texts, training_masks, reduce=False
+        )
+        losses = dict(loss_text=loss_text, iou_text=iou_text)
+
+        # kernel loss
+        loss_kernels = []
+        if self.kernel_sample_mask == "gt":
+            selected_masks = gt_texts * training_masks
+        elif self.kernel_sample_mask == "pred":
+            selected_masks = (F.sigmoid(texts) > 0.5).astype("float32") * training_masks
+
+        for i in range(kernels.shape[1]):
+            kernel_i = kernels[:, i, :, :]
+            gt_kernel_i = gt_kernels[:, i, :, :]
+            loss_kernel_i = self.dice_loss(kernel_i, gt_kernel_i, selected_masks)
+            loss_kernels.append(loss_kernel_i)
+        loss_kernels = paddle.mean(paddle.stack(loss_kernels, axis=1), axis=1)
+        iou_kernel = iou(
+            (kernels[:, -1, :, :] > 0).astype("int64"),
+            gt_kernels[:, -1, :, :],
+            training_masks * gt_texts,
+            reduce=False,
+        )
+        losses.update(dict(loss_kernels=loss_kernels, iou_kernel=iou_kernel))
+        loss = self.alpha * loss_text + (1 - self.alpha) * loss_kernels
+        losses["loss"] = loss
+        if self.reduction == "sum":
+            losses = {x: paddle.sum(v) for x, v in losses.items()}
+        elif self.reduction == "mean":
+            losses = {x: paddle.mean(v) for x, v in losses.items()}
+        return losses
+
+    def dice_loss(self, input, target, mask):
+        input = F.sigmoid(input)
+
+        input = input.reshape([input.shape[0], -1])
+        target = target.reshape([target.shape[0], -1])
+        mask = mask.reshape([mask.shape[0], -1])
+
+        input = input * mask
+        target = target * mask
+
+        a = paddle.sum(input * target, 1)
+        b = paddle.sum(input * input, 1) + self.eps
+        c = paddle.sum(target * target, 1) + self.eps
+        d = (2 * a) / (b + c)
+        return 1 - d
+
+    def ohem_single(self, score, gt_text, training_mask, ohem_ratio=3):
+        pos_num = int(paddle.sum((gt_text > 0.5).astype("float32"))) - int(
+            paddle.sum(
+                paddle.logical_and((gt_text > 0.5), (training_mask <= 0.5)).astype(
+                    "float32"
+                )
+            )
+        )
+
+        if pos_num == 0:
+            selected_mask = training_mask
+            selected_mask = selected_mask.reshape(
+                [1, selected_mask.shape[0], selected_mask.shape[1]]
+            ).astype("float32")
+            return selected_mask
+
+        neg_num = int(paddle.sum((gt_text <= 0.5).astype("float32")))
+        neg_num = int(min(pos_num * ohem_ratio, neg_num))
+
+        if neg_num == 0:
+            selected_mask = training_mask
+            selected_mask = selected_mask.reshape(
+                [1, selected_mask.shape[0], selected_mask.shape[1]]
+            ).astype("float32")
+            return selected_mask
+
+        neg_score = paddle.masked_select(score, gt_text <= 0.5)
+        neg_score_sorted = paddle.sort(-neg_score)
+        threshold = -neg_score_sorted[neg_num - 1]
+
+        selected_mask = paddle.logical_and(
+            paddle.logical_or((score >= threshold), (gt_text > 0.5)),
+            (training_mask > 0.5),
+        )
+        selected_mask = selected_mask.reshape(
+            [1, selected_mask.shape[0], selected_mask.shape[1]]
+        ).astype("float32")
+        return selected_mask
+
+    def ohem_batch(self, scores, gt_texts, training_masks, ohem_ratio=3):
+        selected_masks = []
+        for i in range(scores.shape[0]):
+            selected_masks.append(
+                self.ohem_single(
+                    scores[i, :, :],
+                    gt_texts[i, :, :],
+                    training_masks[i, :, :],
+                    ohem_ratio,
+                )
+            )
+
+        selected_masks = paddle.concat(selected_masks, 0).astype("float32")
+        return selected_masks

ppocr/losses/det_sast_loss.py

@@ -0,0 +1,133 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+from .det_basic_loss import DiceLoss
+import numpy as np
+
+
+class SASTLoss(nn.Layer):
+    """ """
+
+    def __init__(self, eps=1e-6, **kwargs):
+        super(SASTLoss, self).__init__()
+        self.dice_loss = DiceLoss(eps=eps)
+
+    def forward(self, predicts, labels):
+        """
+        tcl_pos: N x 128 x 3
+        tcl_mask: N x 128 x 1
+        tcl_label: N x X list or LoDTensor
+        """
+
+        f_score = predicts["f_score"]
+        f_border = predicts["f_border"]
+        f_tvo = predicts["f_tvo"]
+        f_tco = predicts["f_tco"]
+
+        l_score, l_border, l_mask, l_tvo, l_tco = labels[1:]
+
+        # score_loss
+        intersection = paddle.sum(f_score * l_score * l_mask)
+        union = paddle.sum(f_score * l_mask) + paddle.sum(l_score * l_mask)
+        score_loss = 1.0 - 2 * intersection / (union + 1e-5)
+
+        # border loss
+        l_border_split, l_border_norm = paddle.split(
+            l_border, num_or_sections=[4, 1], axis=1
+        )
+        f_border_split = f_border
+        border_ex_shape = l_border_norm.shape * np.array([1, 4, 1, 1])
+        l_border_norm_split = paddle.expand(x=l_border_norm, shape=border_ex_shape)
+        l_border_score = paddle.expand(x=l_score, shape=border_ex_shape)
+        l_border_mask = paddle.expand(x=l_mask, shape=border_ex_shape)
+
+        border_diff = l_border_split - f_border_split
+        abs_border_diff = paddle.abs(border_diff)
+        border_sign = abs_border_diff < 1.0
+        border_sign = paddle.cast(border_sign, dtype="float32")
+        border_sign.stop_gradient = True
+        border_in_loss = 0.5 * abs_border_diff * abs_border_diff * border_sign + (
+            abs_border_diff - 0.5
+        ) * (1.0 - border_sign)
+        border_out_loss = l_border_norm_split * border_in_loss
+        border_loss = paddle.sum(border_out_loss * l_border_score * l_border_mask) / (
+            paddle.sum(l_border_score * l_border_mask) + 1e-5
+        )
+
+        # tvo_loss
+        l_tvo_split, l_tvo_norm = paddle.split(l_tvo, num_or_sections=[8, 1], axis=1)
+        f_tvo_split = f_tvo
+        tvo_ex_shape = l_tvo_norm.shape * np.array([1, 8, 1, 1])
+        l_tvo_norm_split = paddle.expand(x=l_tvo_norm, shape=tvo_ex_shape)
+        l_tvo_score = paddle.expand(x=l_score, shape=tvo_ex_shape)
+        l_tvo_mask = paddle.expand(x=l_mask, shape=tvo_ex_shape)
+        #
+        tvo_geo_diff = l_tvo_split - f_tvo_split
+        abs_tvo_geo_diff = paddle.abs(tvo_geo_diff)
+        tvo_sign = abs_tvo_geo_diff < 1.0
+        tvo_sign = paddle.cast(tvo_sign, dtype="float32")
+        tvo_sign.stop_gradient = True
+        tvo_in_loss = 0.5 * abs_tvo_geo_diff * abs_tvo_geo_diff * tvo_sign + (
+            abs_tvo_geo_diff - 0.5
+        ) * (1.0 - tvo_sign)
+        tvo_out_loss = l_tvo_norm_split * tvo_in_loss
+        tvo_loss = paddle.sum(tvo_out_loss * l_tvo_score * l_tvo_mask) / (
+            paddle.sum(l_tvo_score * l_tvo_mask) + 1e-5
+        )
+
+        # tco_loss
+        l_tco_split, l_tco_norm = paddle.split(l_tco, num_or_sections=[2, 1], axis=1)
+        f_tco_split = f_tco
+        tco_ex_shape = l_tco_norm.shape * np.array([1, 2, 1, 1])
+        l_tco_norm_split = paddle.expand(x=l_tco_norm, shape=tco_ex_shape)
+        l_tco_score = paddle.expand(x=l_score, shape=tco_ex_shape)
+        l_tco_mask = paddle.expand(x=l_mask, shape=tco_ex_shape)
+
+        tco_geo_diff = l_tco_split - f_tco_split
+        abs_tco_geo_diff = paddle.abs(tco_geo_diff)
+        tco_sign = abs_tco_geo_diff < 1.0
+        tco_sign = paddle.cast(tco_sign, dtype="float32")
+        tco_sign.stop_gradient = True
+        tco_in_loss = 0.5 * abs_tco_geo_diff * abs_tco_geo_diff * tco_sign + (
+            abs_tco_geo_diff - 0.5
+        ) * (1.0 - tco_sign)
+        tco_out_loss = l_tco_norm_split * tco_in_loss
+        tco_loss = paddle.sum(tco_out_loss * l_tco_score * l_tco_mask) / (
+            paddle.sum(l_tco_score * l_tco_mask) + 1e-5
+        )
+
+        # total loss
+        tvo_lw, tco_lw = 1.5, 1.5
+        score_lw, border_lw = 1.0, 1.0
+        total_loss = (
+            score_loss * score_lw
+            + border_loss * border_lw
+            + tvo_loss * tvo_lw
+            + tco_loss * tco_lw
+        )
+
+        losses = {
+            "loss": total_loss,
+            "score_loss": score_loss,
+            "border_loss": border_loss,
+            "tvo_loss": tvo_loss,
+            "tco_loss": tco_loss,
+        }
+        return losses

ppocr/losses/e2e_pg_loss.py

@@ -0,0 +1,165 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+import paddle
+
+from .det_basic_loss import DiceLoss
+from ppocr.utils.e2e_utils.extract_batchsize import pre_process
+
+
+class PGLoss(nn.Layer):
+    def __init__(
+        self, tcl_bs, max_text_length, max_text_nums, pad_num, eps=1e-6, **kwargs
+    ):
+        super(PGLoss, self).__init__()
+        self.tcl_bs = tcl_bs
+        self.max_text_nums = max_text_nums
+        self.max_text_length = max_text_length
+        self.pad_num = pad_num
+        self.dice_loss = DiceLoss(eps=eps)
+
+    def border_loss(self, f_border, l_border, l_score, l_mask):
+        l_border_split, l_border_norm = paddle.tensor.split(
+            l_border, num_or_sections=[4, 1], axis=1
+        )
+        f_border_split = f_border
+        b, c, h, w = l_border_norm.shape
+        l_border_norm_split = paddle.expand(x=l_border_norm, shape=[b, 4 * c, h, w])
+        b, c, h, w = l_score.shape
+        l_border_score = paddle.expand(x=l_score, shape=[b, 4 * c, h, w])
+        b, c, h, w = l_mask.shape
+        l_border_mask = paddle.expand(x=l_mask, shape=[b, 4 * c, h, w])
+        border_diff = l_border_split - f_border_split
+        abs_border_diff = paddle.abs(border_diff)
+        border_sign = abs_border_diff < 1.0
+        border_sign = paddle.cast(border_sign, dtype="float32")
+        border_sign.stop_gradient = True
+        border_in_loss = 0.5 * abs_border_diff * abs_border_diff * border_sign + (
+            abs_border_diff - 0.5
+        ) * (1.0 - border_sign)
+        border_out_loss = l_border_norm_split * border_in_loss
+        border_loss = paddle.sum(border_out_loss * l_border_score * l_border_mask) / (
+            paddle.sum(l_border_score * l_border_mask) + 1e-5
+        )
+        return border_loss
+
+    def direction_loss(self, f_direction, l_direction, l_score, l_mask):
+        l_direction_split, l_direction_norm = paddle.tensor.split(
+            l_direction, num_or_sections=[2, 1], axis=1
+        )
+        f_direction_split = f_direction
+        b, c, h, w = l_direction_norm.shape
+        l_direction_norm_split = paddle.expand(
+            x=l_direction_norm, shape=[b, 2 * c, h, w]
+        )
+        b, c, h, w = l_score.shape
+        l_direction_score = paddle.expand(x=l_score, shape=[b, 2 * c, h, w])
+        b, c, h, w = l_mask.shape
+        l_direction_mask = paddle.expand(x=l_mask, shape=[b, 2 * c, h, w])
+        direction_diff = l_direction_split - f_direction_split
+        abs_direction_diff = paddle.abs(direction_diff)
+        direction_sign = abs_direction_diff < 1.0
+        direction_sign = paddle.cast(direction_sign, dtype="float32")
+        direction_sign.stop_gradient = True
+        direction_in_loss = (
+            0.5 * abs_direction_diff * abs_direction_diff * direction_sign
+            + (abs_direction_diff - 0.5) * (1.0 - direction_sign)
+        )
+        direction_out_loss = l_direction_norm_split * direction_in_loss
+        direction_loss = paddle.sum(
+            direction_out_loss * l_direction_score * l_direction_mask
+        ) / (paddle.sum(l_direction_score * l_direction_mask) + 1e-5)
+        return direction_loss
+
+    def ctcloss(self, f_char, tcl_pos, tcl_mask, tcl_label, label_t):
+        f_char = paddle.transpose(f_char, [0, 2, 3, 1])
+        tcl_pos = paddle.reshape(tcl_pos, [-1, 3])
+        tcl_pos = paddle.cast(tcl_pos, dtype=int)
+        f_tcl_char = paddle.gather_nd(f_char, tcl_pos)
+        f_tcl_char = paddle.reshape(
+            f_tcl_char, [-1, 64, self.pad_num + 1]
+        )  # len(Lexicon_Table)+1
+        f_tcl_char_fg, f_tcl_char_bg = paddle.split(
+            f_tcl_char, [self.pad_num, 1], axis=2
+        )
+        f_tcl_char_bg = f_tcl_char_bg * tcl_mask + (1.0 - tcl_mask) * 20.0
+        b, c, l = tcl_mask.shape
+        tcl_mask_fg = paddle.expand(x=tcl_mask, shape=[b, c, self.pad_num * l])
+        tcl_mask_fg.stop_gradient = True
+        f_tcl_char_fg = f_tcl_char_fg * tcl_mask_fg + (1.0 - tcl_mask_fg) * (-20.0)
+        f_tcl_char_mask = paddle.concat([f_tcl_char_fg, f_tcl_char_bg], axis=2)
+        f_tcl_char_ld = paddle.transpose(f_tcl_char_mask, (1, 0, 2))
+        N, B, _ = f_tcl_char_ld.shape
+        input_lengths = paddle.to_tensor([N] * B, dtype="int64")
+        cost = paddle.nn.functional.ctc_loss(
+            log_probs=f_tcl_char_ld,
+            labels=tcl_label,
+            input_lengths=input_lengths,
+            label_lengths=label_t,
+            blank=self.pad_num,
+            reduction="none",
+        )
+        cost = cost.mean()
+        return cost
+
+    def forward(self, predicts, labels):
+        (
+            images,
+            tcl_maps,
+            tcl_label_maps,
+            border_maps,
+            direction_maps,
+            training_masks,
+            label_list,
+            pos_list,
+            pos_mask,
+        ) = labels
+        # for all the batch_size
+        pos_list, pos_mask, label_list, label_t = pre_process(
+            label_list,
+            pos_list,
+            pos_mask,
+            self.max_text_length,
+            self.max_text_nums,
+            self.pad_num,
+            self.tcl_bs,
+        )
+
+        f_score, f_border, f_direction, f_char = (
+            predicts["f_score"],
+            predicts["f_border"],
+            predicts["f_direction"],
+            predicts["f_char"],
+        )
+        score_loss = self.dice_loss(f_score, tcl_maps, training_masks)
+        border_loss = self.border_loss(f_border, border_maps, tcl_maps, training_masks)
+        direction_loss = self.direction_loss(
+            f_direction, direction_maps, tcl_maps, training_masks
+        )
+        ctc_loss = self.ctcloss(f_char, pos_list, pos_mask, label_list, label_t)
+        loss_all = score_loss + border_loss + direction_loss + 5 * ctc_loss
+
+        losses = {
+            "loss": loss_all,
+            "score_loss": score_loss,
+            "border_loss": border_loss,
+            "direction_loss": direction_loss,
+            "ctc_loss": ctc_loss,
+        }
+        return losses

ppocr/losses/kie_sdmgr_loss.py

@@ -0,0 +1,116 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+# reference from : https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/kie/losses/sdmgr_loss.py
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+import paddle
+
+
+class SDMGRLoss(nn.Layer):
+    def __init__(self, node_weight=1.0, edge_weight=1.0, ignore=0):
+        super().__init__()
+        self.loss_node = nn.CrossEntropyLoss(ignore_index=ignore)
+        self.loss_edge = nn.CrossEntropyLoss(ignore_index=-1)
+        self.node_weight = node_weight
+        self.edge_weight = edge_weight
+        self.ignore = ignore
+
+    def pre_process(self, gts, tag):
+        gts, tag = gts.numpy(), tag.numpy().tolist()
+        temp_gts = []
+        batch = len(tag)
+        for i in range(batch):
+            num, recoder_len = tag[i][0], tag[i][1]
+            temp_gts.append(paddle.to_tensor(gts[i, :num, : num + 1], dtype="int64"))
+        return temp_gts
+
+    def accuracy(self, pred, target, topk=1, thresh=None):
+        """Calculate accuracy according to the prediction and target.
+
+        Args:
+            pred (torch.Tensor): The model prediction, shape (N, num_class)
+            target (torch.Tensor): The target of each prediction, shape (N, )
+            topk (int | tuple[int], optional): If the predictions in ``topk``
+                matches the target, the predictions will be regarded as
+                correct ones. Defaults to 1.
+            thresh (float, optional): If not None, predictions with scores under
+                this threshold are considered incorrect. Default to None.
+
+        Returns:
+            float | tuple[float]: If the input ``topk`` is a single integer,
+                the function will return a single float as accuracy. If
+                ``topk`` is a tuple containing multiple integers, the
+                function will return a tuple containing accuracies of
+                each ``topk`` number.
+        """
+        assert isinstance(topk, (int, tuple))
+        if isinstance(topk, int):
+            topk = (topk,)
+            return_single = True
+        else:
+            return_single = False
+
+        maxk = max(topk)
+        if pred.shape[0] == 0:
+            accu = [pred.new_tensor(0.0) for i in range(len(topk))]
+            return accu[0] if return_single else accu
+        pred_value, pred_label = paddle.topk(pred, maxk, axis=1)
+        pred_label = pred_label.transpose([1, 0])  # transpose to shape (maxk, N)
+        correct = paddle.equal(
+            pred_label, (target.reshape([1, -1]).expand_as(pred_label))
+        )
+        res = []
+        for k in topk:
+            correct_k = paddle.sum(
+                correct[:k].reshape([-1]).astype("float32"), axis=0, keepdim=True
+            )
+            res.append(
+                paddle.multiply(correct_k, paddle.to_tensor(100.0 / pred.shape[0]))
+            )
+        return res[0] if return_single else res
+
+    def forward(self, pred, batch):
+        node_preds, edge_preds = pred
+        gts, tag = batch[4], batch[5]
+        gts = self.pre_process(gts, tag)
+        node_gts, edge_gts = [], []
+        for gt in gts:
+            node_gts.append(gt[:, 0])
+            edge_gts.append(gt[:, 1:].reshape([-1]))
+        node_gts = paddle.concat(node_gts)
+        edge_gts = paddle.concat(edge_gts)
+
+        node_valids = paddle.nonzero(node_gts != self.ignore).reshape([-1])
+        edge_valids = paddle.nonzero(edge_gts != -1).reshape([-1])
+        loss_node = self.loss_node(node_preds, node_gts)
+        loss_edge = self.loss_edge(edge_preds, edge_gts)
+        loss = self.node_weight * loss_node + self.edge_weight * loss_edge
+        return dict(
+            loss=loss,
+            loss_node=loss_node,
+            loss_edge=loss_edge,
+            acc_node=self.accuracy(
+                paddle.gather(node_preds, node_valids),
+                paddle.gather(node_gts, node_valids),
+            ),
+            acc_edge=self.accuracy(
+                paddle.gather(edge_preds, edge_valids),
+                paddle.gather(edge_gts, edge_valids),
+            ),
+        )

ppocr/losses/rec_aster_loss.py

@@ -0,0 +1,103 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class CosineEmbeddingLoss(nn.Layer):
+    def __init__(self, margin=0.0):
+        super(CosineEmbeddingLoss, self).__init__()
+        self.margin = margin
+        self.epsilon = 1e-12
+
+    def forward(self, x1, x2, target):
+        similarity = paddle.sum(x1 * x2, axis=-1) / (
+            paddle.norm(x1, axis=-1) * paddle.norm(x2, axis=-1) + self.epsilon
+        )
+        one_list = paddle.full_like(target, fill_value=1)
+        out = paddle.mean(
+            paddle.where(
+                paddle.equal(target, one_list),
+                1.0 - similarity,
+                paddle.maximum(paddle.zeros_like(similarity), similarity - self.margin),
+            )
+        )
+
+        return out
+
+
+class AsterLoss(nn.Layer):
+    def __init__(
+        self,
+        weight=None,
+        size_average=True,
+        ignore_index=-100,
+        sequence_normalize=False,
+        sample_normalize=True,
+        **kwargs,
+    ):
+        super(AsterLoss, self).__init__()
+        self.weight = weight
+        self.size_average = size_average
+        self.ignore_index = ignore_index
+        self.sequence_normalize = sequence_normalize
+        self.sample_normalize = sample_normalize
+        self.loss_sem = CosineEmbeddingLoss()
+        self.is_cosin_loss = True
+        self.loss_func_rec = nn.CrossEntropyLoss(weight=None, reduction="none")
+
+    def forward(self, predicts, batch):
+        targets = batch[1].astype("int64")
+        label_lengths = batch[2].astype("int64")
+        sem_target = batch[3].astype("float32")
+        embedding_vectors = predicts["embedding_vectors"]
+        rec_pred = predicts["rec_pred"]
+
+        if not self.is_cosin_loss:
+            sem_loss = paddle.sum(self.loss_sem(embedding_vectors, sem_target))
+        else:
+            label_target = paddle.ones([embedding_vectors.shape[0]])
+            sem_loss = paddle.sum(
+                self.loss_sem(embedding_vectors, sem_target, label_target)
+            )
+
+        # rec loss
+        batch_size, def_max_length = targets.shape[0], targets.shape[1]
+
+        mask = paddle.zeros([batch_size, def_max_length])
+        for i in range(batch_size):
+            mask[i, : label_lengths[i]] = 1
+        mask = paddle.cast(mask, "float32")
+        max_length = max(label_lengths)
+        assert max_length == rec_pred.shape[1]
+        targets = targets[:, :max_length]
+        mask = mask[:, :max_length]
+        rec_pred = paddle.reshape(rec_pred, [-1, rec_pred.shape[2]])
+        input = nn.functional.log_softmax(rec_pred, axis=1)
+        targets = paddle.reshape(targets, [-1, 1])
+        mask = paddle.reshape(mask, [-1, 1])
+        output = -paddle.index_sample(input, index=targets) * mask
+        output = paddle.sum(output)
+        if self.sequence_normalize:
+            output = output / paddle.sum(mask)
+        if self.sample_normalize:
+            output = output / batch_size
+
+        loss = output + sem_loss * 0.1
+        return {"loss": loss}

ppocr/losses/rec_att_loss.py

@@ -0,0 +1,43 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class AttentionLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(AttentionLoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(weight=None, reduction="none")
+
+    def forward(self, predicts, batch):
+        targets = batch[1].astype("int64")
+        label_lengths = batch[2].astype("int64")
+        batch_size, num_steps, num_classes = (
+            predicts.shape[0],
+            predicts.shape[1],
+            predicts.shape[2],
+        )
+        assert (
+            len(targets.shape) == len(list(predicts.shape)) - 1
+        ), "The target's shape and inputs's shape is [N, d] and [N, num_steps]"
+
+        inputs = paddle.reshape(predicts, [-1, predicts.shape[-1]])
+        targets = paddle.reshape(targets, [-1])
+
+        return {"loss": paddle.sum(self.loss_func(inputs, targets))}

ppocr/losses/rec_can_loss.py

@@ -0,0 +1,88 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/LBH1024/CAN/models/can.py
+"""
+
+import paddle
+import paddle.nn as nn
+import numpy as np
+
+
+class CANLoss(nn.Layer):
+    """
+    CANLoss is consist of two part:
+        word_average_loss: average accuracy of the symbol
+        counting_loss: counting loss of every symbol
+    """
+
+    def __init__(self):
+        super(CANLoss, self).__init__()
+
+        self.use_label_mask = False
+        self.out_channel = 111
+        self.cross = (
+            nn.CrossEntropyLoss(reduction="none")
+            if self.use_label_mask
+            else nn.CrossEntropyLoss()
+        )
+        self.counting_loss = nn.SmoothL1Loss(reduction="mean")
+        self.ratio = 16
+
+    def forward(self, preds, batch):
+        word_probs = preds[0]
+        counting_preds = preds[1]
+        counting_preds1 = preds[2]
+        counting_preds2 = preds[3]
+        labels = batch[2]
+        labels_mask = batch[3]
+        counting_labels = gen_counting_label(labels, self.out_channel, True)
+        counting_loss = (
+            self.counting_loss(counting_preds1, counting_labels)
+            + self.counting_loss(counting_preds2, counting_labels)
+            + self.counting_loss(counting_preds, counting_labels)
+        )
+
+        word_loss = self.cross(
+            paddle.reshape(word_probs, [-1, word_probs.shape[-1]]),
+            paddle.reshape(labels, [-1]),
+        )
+        word_average_loss = (
+            paddle.sum(paddle.reshape(word_loss * labels_mask, [-1]))
+            / (paddle.sum(labels_mask) + 1e-10)
+            if self.use_label_mask
+            else word_loss
+        )
+        loss = word_average_loss + counting_loss
+        return {"loss": loss}
+
+
+def gen_counting_label(labels, channel, tag):
+    b, t = labels.shape
+    counting_labels = np.zeros([b, channel])
+
+    if tag:
+        ignore = [0, 1, 107, 108, 109, 110]
+    else:
+        ignore = []
+    for i in range(b):
+        for j in range(t):
+            k = labels[i][j]
+            if k in ignore:
+                continue
+            else:
+                counting_labels[i][k] += 1
+    counting_labels = paddle.to_tensor(counting_labels, dtype="float32")
+    return counting_labels

ppocr/losses/rec_ce_loss.py

@@ -0,0 +1,61 @@
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+
+
+class CELoss(nn.Layer):
+    def __init__(self, smoothing=False, with_all=False, ignore_index=-1, **kwargs):
+        super(CELoss, self).__init__()
+        if ignore_index >= 0:
+            self.loss_func = nn.CrossEntropyLoss(
+                reduction="mean", ignore_index=ignore_index
+            )
+        else:
+            self.loss_func = nn.CrossEntropyLoss(reduction="mean")
+        self.smoothing = smoothing
+        self.with_all = with_all
+
+    def forward(self, pred, batch):
+        if isinstance(pred, dict):  # for ABINet
+            loss = {}
+            loss_sum = []
+            for name, logits in pred.items():
+                if isinstance(logits, list):
+                    logit_num = len(logits)
+                    all_tgt = paddle.concat([batch[1]] * logit_num, 0)
+                    all_logits = paddle.concat(logits, 0)
+                    flt_logtis = all_logits.reshape([-1, all_logits.shape[2]])
+                    flt_tgt = all_tgt.reshape([-1])
+                else:
+                    flt_logtis = logits.reshape([-1, logits.shape[2]])
+                    flt_tgt = batch[1].reshape([-1])
+                loss[name + "_loss"] = self.loss_func(flt_logtis, flt_tgt)
+                loss_sum.append(loss[name + "_loss"])
+            loss["loss"] = sum(loss_sum)
+            return loss
+        else:
+            if self.with_all:  # for ViTSTR
+                tgt = batch[1]
+                pred = pred.reshape([-1, pred.shape[2]])
+                tgt = tgt.reshape([-1])
+                loss = self.loss_func(pred, tgt)
+                return {"loss": loss}
+            else:  # for NRTR
+                max_len = batch[2].max()
+                tgt = batch[1][:, 1 : 2 + max_len]
+                pred = pred.reshape([-1, pred.shape[2]])
+                tgt = tgt.reshape([-1])
+                if self.smoothing:
+                    eps = 0.1
+                    n_class = pred.shape[1]
+                    one_hot = F.one_hot(tgt, pred.shape[1])
+                    one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+                    log_prb = F.log_softmax(pred, axis=1)
+                    non_pad_mask = paddle.not_equal(
+                        tgt, paddle.zeros(tgt.shape, dtype=tgt.dtype)
+                    )
+                    loss = -(one_hot * log_prb).sum(axis=1)
+                    loss = loss.masked_select(non_pad_mask).mean()
+                else:
+                    loss = self.loss_func(pred, tgt)
+                return {"loss": loss}

ppocr/losses/rec_cppd_loss.py

@@ -0,0 +1,76 @@
+# copyright (c) 2023 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+
+
+class CPPDLoss(nn.Layer):
+    def __init__(
+        self, smoothing=False, ignore_index=100, sideloss_weight=1.0, **kwargs
+    ):
+        super(CPPDLoss, self).__init__()
+        self.edge_ce = nn.CrossEntropyLoss(reduction="mean", ignore_index=ignore_index)
+        self.char_node_ce = nn.CrossEntropyLoss(reduction="mean")
+        self.pos_node_ce = nn.BCEWithLogitsLoss(reduction="mean")
+        self.smoothing = smoothing
+        self.ignore_index = ignore_index
+        self.sideloss_weight = sideloss_weight
+
+    def label_smoothing_ce(self, preds, targets):
+        non_pad_mask = paddle.not_equal(
+            targets,
+            paddle.zeros(targets.shape, dtype=targets.dtype) + self.ignore_index,
+        )
+        tgts = paddle.where(
+            targets
+            == (paddle.zeros(targets.shape, dtype=targets.dtype) + self.ignore_index),
+            paddle.zeros(targets.shape, dtype=targets.dtype),
+            targets,
+        )
+        eps = 0.1
+        n_class = preds.shape[1]
+        one_hot = F.one_hot(tgts, preds.shape[1])
+        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+        log_prb = F.log_softmax(preds, axis=1)
+        loss = -(one_hot * log_prb).sum(axis=1)
+        loss = loss.masked_select(non_pad_mask).mean()
+        return loss
+
+    def forward(self, pred, batch):
+        node_feats, edge_feats = pred
+        node_tgt = batch[2]
+        char_tgt = batch[1]
+
+        loss_char_node = self.char_node_ce(
+            node_feats[0].flatten(0, 1), node_tgt[:, :-26].flatten(0, 1)
+        )
+        loss_pos_node = self.pos_node_ce(
+            node_feats[1].flatten(0, 1), node_tgt[:, -26:].flatten(0, 1).cast("float32")
+        )
+        loss_node = loss_char_node + loss_pos_node
+
+        edge_feats = edge_feats.flatten(0, 1)
+        char_tgt = char_tgt.flatten(0, 1)
+        if self.smoothing:
+            loss_edge = self.label_smoothing_ce(edge_feats, char_tgt)
+        else:
+            loss_edge = self.edge_ce(edge_feats, char_tgt)
+
+        return {
+            "loss": self.sideloss_weight * loss_node + loss_edge,
+            "loss_node": self.sideloss_weight * loss_node,
+            "loss_edge": loss_edge,
+        }

ppocr/losses/rec_ctc_loss.py

@@ -0,0 +1,46 @@
+# copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class CTCLoss(nn.Layer):
+    def __init__(self, use_focal_loss=False, **kwargs):
+        super(CTCLoss, self).__init__()
+        self.loss_func = nn.CTCLoss(blank=0, reduction="none")
+        self.use_focal_loss = use_focal_loss
+
+    def forward(self, predicts, batch):
+        if isinstance(predicts, (list, tuple)):
+            predicts = predicts[-1]
+        predicts = predicts.transpose((1, 0, 2))
+        N, B, _ = predicts.shape
+        preds_lengths = paddle.to_tensor(
+            [N] * B, dtype="int64", place=paddle.CPUPlace()
+        )
+        labels = batch[1].astype("int32")
+        label_lengths = batch[2].astype("int64")
+        loss = self.loss_func(predicts, labels, preds_lengths, label_lengths)
+        if self.use_focal_loss:
+            weight = paddle.exp(-loss)
+            weight = paddle.subtract(paddle.to_tensor([1.0]), weight)
+            weight = paddle.square(weight)
+            loss = paddle.multiply(loss, weight)
+        loss = loss.mean()
+        return {"loss": loss}

ppocr/losses/rec_enhanced_ctc_loss.py

@@ -0,0 +1,76 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+from .ace_loss import ACELoss
+from .center_loss import CenterLoss
+from .rec_ctc_loss import CTCLoss
+
+
+class EnhancedCTCLoss(nn.Layer):
+    def __init__(
+        self,
+        use_focal_loss=False,
+        use_ace_loss=False,
+        ace_loss_weight=0.1,
+        use_center_loss=False,
+        center_loss_weight=0.05,
+        num_classes=6625,
+        feat_dim=96,
+        init_center=False,
+        center_file_path=None,
+        **kwargs,
+    ):
+        super(EnhancedCTCLoss, self).__init__()
+        self.ctc_loss_func = CTCLoss(use_focal_loss=use_focal_loss)
+
+        self.use_ace_loss = False
+        if use_ace_loss:
+            self.use_ace_loss = use_ace_loss
+            self.ace_loss_func = ACELoss()
+            self.ace_loss_weight = ace_loss_weight
+
+        self.use_center_loss = False
+        if use_center_loss:
+            self.use_center_loss = use_center_loss
+            self.center_loss_func = CenterLoss(
+                num_classes=num_classes,
+                feat_dim=feat_dim,
+                init_center=init_center,
+                center_file_path=center_file_path,
+            )
+            self.center_loss_weight = center_loss_weight
+
+    def __call__(self, predicts, batch):
+        loss = self.ctc_loss_func(predicts, batch)["loss"]
+
+        if self.use_center_loss:
+            center_loss = (
+                self.center_loss_func(predicts, batch)["loss_center"]
+                * self.center_loss_weight
+            )
+            loss = loss + center_loss
+
+        if self.use_ace_loss:
+            ace_loss = (
+                self.ace_loss_func(predicts, batch)["loss_ace"] * self.ace_loss_weight
+            )
+            loss = loss + ace_loss
+
+        return {"enhanced_ctc_loss": loss}

ppocr/losses/rec_latexocr_loss.py

@@ -0,0 +1,47 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+"""
+This code is refer from:
+https://github.com/lucidrains/x-transformers/blob/main/x_transformers/autoregressive_wrapper.py
+"""
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+
+
+class LaTeXOCRLoss(nn.Layer):
+    """
+    LaTeXOCR adopt CrossEntropyLoss for network training.
+    """
+
+    def __init__(self):
+        super(LaTeXOCRLoss, self).__init__()
+        self.ignore_index = -100
+        self.cross = nn.CrossEntropyLoss(
+            reduction="mean", ignore_index=self.ignore_index
+        )
+
+    def forward(self, preds, batch):
+        word_probs = preds
+        labels = batch[1][:, 1:]
+        word_loss = self.cross(
+            paddle.reshape(word_probs, [-1, word_probs.shape[-1]]),
+            paddle.reshape(labels, [-1]),
+        )
+
+        loss = word_loss
+        return {"loss": loss}

ppocr/losses/rec_multi_loss.py

@@ -0,0 +1,68 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+from .rec_ctc_loss import CTCLoss
+from .rec_sar_loss import SARLoss
+from .rec_nrtr_loss import NRTRLoss
+
+
+class MultiLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.loss_funcs = {}
+        self.loss_list = kwargs.pop("loss_config_list")
+        self.weight_1 = kwargs.get("weight_1", 1.0)
+        self.weight_2 = kwargs.get("weight_2", 1.0)
+        for loss_info in self.loss_list:
+            for name, param in loss_info.items():
+                if param is not None:
+                    kwargs.update(param)
+                loss = eval(name)(**kwargs)
+                self.loss_funcs[name] = loss
+
+    def forward(self, predicts, batch):
+        self.total_loss = {}
+        total_loss = 0.0
+        # batch [image, label_ctc, label_sar, length, valid_ratio]
+        for name, loss_func in self.loss_funcs.items():
+            if name == "CTCLoss":
+                loss = (
+                    loss_func(predicts["ctc"], batch[:2] + batch[3:])["loss"]
+                    * self.weight_1
+                )
+            elif name == "SARLoss":
+                loss = (
+                    loss_func(predicts["sar"], batch[:1] + batch[2:])["loss"]
+                    * self.weight_2
+                )
+            elif name == "NRTRLoss":
+                loss = (
+                    loss_func(predicts["gtc"], batch[:1] + batch[2:])["loss"]
+                    * self.weight_2
+                )
+            else:
+                raise NotImplementedError(
+                    "{} is not supported in MultiLoss yet".format(name)
+                )
+            self.total_loss[name] = loss
+            total_loss += loss
+        self.total_loss["loss"] = total_loss
+        return self.total_loss

ppocr/losses/rec_nrtr_loss.py

@@ -0,0 +1,33 @@
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+
+
+class NRTRLoss(nn.Layer):
+    def __init__(self, smoothing=True, ignore_index=0, **kwargs):
+        super(NRTRLoss, self).__init__()
+        if ignore_index >= 0 and not smoothing:
+            self.loss_func = nn.CrossEntropyLoss(
+                reduction="mean", ignore_index=ignore_index
+            )
+        self.smoothing = smoothing
+
+    def forward(self, pred, batch):
+        max_len = batch[2].max()
+        tgt = batch[1][:, 1 : 2 + max_len]
+        pred = pred.reshape([-1, pred.shape[2]])
+        tgt = tgt.reshape([-1])
+        if self.smoothing:
+            eps = 0.1
+            n_class = pred.shape[1]
+            one_hot = F.one_hot(tgt, pred.shape[1])
+            one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+            log_prb = F.log_softmax(pred, axis=1)
+            non_pad_mask = paddle.not_equal(
+                tgt, paddle.zeros(tgt.shape, dtype=tgt.dtype)
+            )
+            loss = -(one_hot * log_prb).sum(axis=1)
+            loss = loss.masked_select(non_pad_mask).mean()
+        else:
+            loss = self.loss_func(pred, tgt)
+        return {"loss": loss}

ppocr/losses/rec_parseq_loss.py

@@ -0,0 +1,52 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class ParseQLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(ParseQLoss, self).__init__()
+
+    def forward(self, predicts, targets):
+        label = targets[1]  # label
+        label_len = targets[2]
+        max_step = paddle.max(label_len).cpu().numpy()[0] + 2
+        tgt = label[:, :max_step]
+
+        logits_list = predicts["logits_list"]
+        pad_id = predicts["pad_id"]
+        eos_id = predicts["eos_id"]
+
+        tgt_out = tgt[:, 1:]
+        loss = 0
+        loss_numel = 0
+        n = (tgt_out != pad_id).sum().item()
+
+        for i, logits in enumerate(logits_list):
+            loss += n * paddle.nn.functional.cross_entropy(
+                input=logits, label=tgt_out.flatten(), ignore_index=pad_id
+            )
+            loss_numel += n
+            if i == 1:
+                tgt_out = paddle.where(condition=tgt_out == eos_id, x=pad_id, y=tgt_out)
+                n = (tgt_out != pad_id).sum().item()
+        loss /= loss_numel
+
+        return {"loss": loss}

ppocr/losses/rec_ppformulanet_loss.py

@@ -0,0 +1,74 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+import paddle
+import paddle.nn as nn
+
+
+class PPFormulaNet_S_Loss(nn.Layer):
+    """
+    PP=FormulaNet-S adopt CrossEntropyLoss for network training.
+    """
+
+    def __init__(self, vocab_size=50000, parallel_step=1):
+        super(PPFormulaNet_S_Loss, self).__init__()
+        self.ignore_index = -100
+        self.vocab_size = vocab_size
+        self.parallel_step = int(parallel_step)
+        self.pad_token_id = 1
+        # ignore padding characters during training
+        self.cross = nn.CrossEntropyLoss(
+            reduction="mean", ignore_index=self.ignore_index
+        )
+
+    def forward(self, preds, batch):
+        logits, masked_label = preds
+
+        word_loss = self.cross(
+            paddle.reshape(logits, [-1, logits.shape[-1]]),
+            paddle.reshape(masked_label[:, self.parallel_step :], [-1]),
+        )
+        loss = word_loss
+        return {
+            "loss": loss,
+            "word_loss": word_loss,
+        }
+
+
+class PPFormulaNet_L_Loss(nn.Layer):
+    """
+    PPFormulaNet_L adopt CrossEntropyLoss for network training.
+    """
+
+    def __init__(self, vocab_size=50000):
+        super(PPFormulaNet_L_Loss, self).__init__()
+        self.ignore_index = -100
+        self.vocab_size = vocab_size
+        self.pad_token_id = 1
+        # ignore padding characters during training
+        self.cross = nn.CrossEntropyLoss(
+            reduction="mean", ignore_index=self.ignore_index
+        )
+
+    def forward(self, preds, batch):
+        logits, masked_label = preds
+
+        word_loss = self.cross(
+            paddle.reshape(logits, [-1, logits.shape[-1]]),
+            paddle.reshape(masked_label[:, 1:], [-1]),
+        )
+        loss = word_loss
+        return {
+            "loss": loss,
+            "word_loss": word_loss,
+        }

ppocr/losses/rec_pren_loss.py

@@ -0,0 +1,30 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+
+
+class PRENLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(PRENLoss, self).__init__()
+        # note: 0 is padding idx
+        self.loss_func = nn.CrossEntropyLoss(reduction="mean", ignore_index=0)
+
+    def forward(self, predicts, batch):
+        loss = self.loss_func(predicts, batch[1].astype("int64"))
+        return {"loss": loss}

ppocr/losses/rec_rfl_loss.py

@@ -0,0 +1,70 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/hikopensource/DAVAR-Lab-OCR/blob/main/davarocr/davar_common/models/loss/cross_entropy_loss.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+from .basic_loss import CELoss, DistanceLoss
+
+
+class RFLLoss(nn.Layer):
+    def __init__(self, ignore_index=-100, **kwargs):
+        super().__init__()
+
+        self.cnt_loss = nn.MSELoss(**kwargs)
+        self.seq_loss = nn.CrossEntropyLoss(ignore_index=ignore_index)
+
+    def forward(self, predicts, batch):
+        self.total_loss = {}
+        total_loss = 0.0
+        if isinstance(predicts, tuple) or isinstance(predicts, list):
+            cnt_outputs, seq_outputs = predicts
+        else:
+            cnt_outputs, seq_outputs = predicts, None
+        # batch [image, label, length, cnt_label]
+        if cnt_outputs is not None:
+            cnt_loss = self.cnt_loss(cnt_outputs, paddle.cast(batch[3], paddle.float32))
+            self.total_loss["cnt_loss"] = cnt_loss
+            total_loss += cnt_loss
+
+        if seq_outputs is not None:
+            targets = batch[1].astype("int64")
+            label_lengths = batch[2].astype("int64")
+            batch_size, num_steps, num_classes = (
+                seq_outputs.shape[0],
+                seq_outputs.shape[1],
+                seq_outputs.shape[2],
+            )
+            assert (
+                len(targets.shape) == len(list(seq_outputs.shape)) - 1
+            ), "The target's shape and inputs's shape is [N, d] and [N, num_steps]"
+
+            inputs = seq_outputs[:, :-1, :]
+            targets = targets[:, 1:]
+
+            inputs = paddle.reshape(inputs, [-1, inputs.shape[-1]])
+            targets = paddle.reshape(targets, [-1])
+            seq_loss = self.seq_loss(inputs, targets)
+            self.total_loss["seq_loss"] = seq_loss
+            total_loss += seq_loss
+
+        self.total_loss["loss"] = total_loss
+        return self.total_loss

ppocr/losses/rec_sar_loss.py

@@ -0,0 +1,36 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class SARLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(SARLoss, self).__init__()
+        ignore_index = kwargs.get("ignore_index", 92)  # 6626
+        self.loss_func = paddle.nn.loss.CrossEntropyLoss(
+            reduction="mean", ignore_index=ignore_index
+        )
+
+    def forward(self, predicts, batch):
+        predict = predicts[
+            :, :-1, :
+        ]  # ignore last index of outputs to be in same seq_len with targets
+        label = batch[1].astype("int64")[
+            :, 1:
+        ]  # ignore first index of target in loss calculation
+        batch_size, num_steps, num_classes = (
+            predict.shape[0],
+            predict.shape[1],
+            predict.shape[2],
+        )
+        assert (
+            len(label.shape) == len(list(predict.shape)) - 1
+        ), "The target's shape and inputs's shape is [N, d] and [N, num_steps]"
+
+        inputs = paddle.reshape(predict, [-1, num_classes])
+        targets = paddle.reshape(label, [-1])
+        loss = self.loss_func(inputs, targets)
+        return {"loss": loss}

ppocr/losses/rec_satrn_loss.py

@@ -0,0 +1,53 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/open-mmlab/mmocr/blob/1.x/mmocr/models/textrecog/module_losses/ce_module_loss.py
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class SATRNLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(SATRNLoss, self).__init__()
+        ignore_index = kwargs.get("ignore_index", 92)  # 6626
+        self.loss_func = paddle.nn.loss.CrossEntropyLoss(
+            reduction="none", ignore_index=ignore_index
+        )
+
+    def forward(self, predicts, batch):
+        predict = predicts[
+            :, :-1, :
+        ]  # ignore last index of outputs to be in same seq_len with targets
+        label = batch[1].astype("int64")[
+            :, 1:
+        ]  # ignore first index of target in loss calculation
+        batch_size, num_steps, num_classes = (
+            predict.shape[0],
+            predict.shape[1],
+            predict.shape[2],
+        )
+        assert (
+            len(label.shape) == len(list(predict.shape)) - 1
+        ), "The target's shape and inputs's shape is [N, d] and [N, num_steps]"
+
+        inputs = paddle.reshape(predict, [-1, num_classes])
+        targets = paddle.reshape(label, [-1])
+        loss = self.loss_func(inputs, targets)
+        return {"loss": loss.mean()}

ppocr/losses/rec_spin_att_loss.py

@@ -0,0 +1,52 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+"""This code is refer from:
+https://github.com/hikopensource/DAVAR-Lab-OCR
+"""
+
+
+class SPINAttentionLoss(nn.Layer):
+    def __init__(self, reduction="mean", ignore_index=-100, **kwargs):
+        super(SPINAttentionLoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(
+            weight=None, reduction=reduction, ignore_index=ignore_index
+        )
+
+    def forward(self, predicts, batch):
+        targets = batch[1].astype("int64")
+        targets = targets[:, 1:]  # remove [eos] in label
+
+        label_lengths = batch[2].astype("int64")
+        batch_size, num_steps, num_classes = (
+            predicts.shape[0],
+            predicts.shape[1],
+            predicts.shape[2],
+        )
+        assert (
+            len(targets.shape) == len(list(predicts.shape)) - 1
+        ), "The target's shape and inputs's shape is [N, d] and [N, num_steps]"
+
+        inputs = paddle.reshape(predicts, [-1, predicts.shape[-1]])
+        targets = paddle.reshape(targets, [-1])
+
+        return {"loss": self.loss_func(inputs, targets)}

ppocr/losses/rec_srn_loss.py

@@ -0,0 +1,47 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class SRNLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(SRNLoss, self).__init__()
+        self.loss_func = paddle.nn.loss.CrossEntropyLoss(reduction="sum")
+
+    def forward(self, predicts, batch):
+        predict = predicts["predict"]
+        word_predict = predicts["word_out"]
+        gsrm_predict = predicts["gsrm_out"]
+        label = batch[1]
+
+        casted_label = paddle.cast(x=label, dtype="int64")
+        casted_label = paddle.reshape(x=casted_label, shape=[-1, 1])
+
+        cost_word = self.loss_func(word_predict, label=casted_label)
+        cost_gsrm = self.loss_func(gsrm_predict, label=casted_label)
+        cost_vsfd = self.loss_func(predict, label=casted_label)
+
+        cost_word = paddle.reshape(x=paddle.sum(cost_word), shape=[1])
+        cost_gsrm = paddle.reshape(x=paddle.sum(cost_gsrm), shape=[1])
+        cost_vsfd = paddle.reshape(x=paddle.sum(cost_vsfd), shape=[1])
+
+        sum_cost = cost_word * 3.0 + cost_vsfd + cost_gsrm * 0.15
+
+        return {"loss": sum_cost, "word_loss": cost_word, "img_loss": cost_vsfd}

ppocr/losses/rec_unimernet_loss.py

@@ -0,0 +1,54 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+import paddle
+import paddle.nn as nn
+import paddle.nn.functional as F
+import numpy as np
+
+
+class UniMERNetLoss(nn.Layer):
+    def __init__(self, length_aware=True, vocab_size=50000):
+        super(UniMERNetLoss, self).__init__()
+        self.ignore_index = -100
+        self.vocab_size = vocab_size
+        self.pad_token_id = 1
+        self.length_aware = length_aware
+        self.cross = nn.CrossEntropyLoss(
+            reduction="mean", ignore_index=self.ignore_index
+        )
+        self.counting_loss_fct = nn.SmoothL1Loss()
+
+    def _get_count_gt(self, labels):
+        mask = (labels != self.pad_token_id).cast("float32")
+        one_hot_labels = F.one_hot(
+            labels, num_classes=self.vocab_size
+        ) * mask.unsqueeze(-1)
+        count_gt = paddle.sum(one_hot_labels, axis=1)
+        return count_gt
+
+    def forward(self, preds, batch):
+        logits, count_pred, masked_label = preds
+        labels = batch[1][:, 1:]
+        word_loss = self.cross(
+            paddle.reshape(logits, [-1, logits.shape[-1]]),
+            paddle.reshape(masked_label[:, 1:], [-1]),
+        )
+        loss = word_loss
+        if self.length_aware:
+            count_gt = self._get_count_gt(labels)
+            count_gt = paddle.log(count_gt.cast(paddle.float32) + 1)
+            count_loss = self.counting_loss_fct(count_pred, count_gt)
+            loss += 0.5 * count_loss
+        return {"loss": loss, "word_loss": word_loss, "count_loss": count_loss}

ppocr/losses/rec_vl_loss.py

@@ -0,0 +1,70 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/wangyuxin87/VisionLAN
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+
+
+class VLLoss(nn.Layer):
+    def __init__(self, mode="LF_1", weight_res=0.5, weight_mas=0.5, **kwargs):
+        super(VLLoss, self).__init__()
+        self.loss_func = paddle.nn.loss.CrossEntropyLoss(reduction="mean")
+        assert mode in ["LF_1", "LF_2", "LA"]
+        self.mode = mode
+        self.weight_res = weight_res
+        self.weight_mas = weight_mas
+
+    def flatten_label(self, target):
+        label_flatten = []
+        label_length = []
+        for i in range(0, target.shape[0]):
+            cur_label = target[i].tolist()
+            label_flatten += cur_label[: cur_label.index(0) + 1]
+            label_length.append(cur_label.index(0) + 1)
+        label_flatten = paddle.to_tensor(label_flatten, dtype="int64")
+        label_length = paddle.to_tensor(label_length, dtype="int32")
+        return (label_flatten, label_length)
+
+    def _flatten(self, sources, lengths):
+        return paddle.concat([t[:l] for t, l in zip(sources, lengths)])
+
+    def forward(self, predicts, batch):
+        text_pre = predicts[0]
+        target = batch[1].astype("int64")
+        label_flatten, length = self.flatten_label(target)
+        text_pre = self._flatten(text_pre, length)
+        if self.mode == "LF_1":
+            loss = self.loss_func(text_pre, label_flatten)
+        else:
+            text_rem = predicts[1]
+            text_mas = predicts[2]
+            target_res = batch[2].astype("int64")
+            target_sub = batch[3].astype("int64")
+            label_flatten_res, length_res = self.flatten_label(target_res)
+            label_flatten_sub, length_sub = self.flatten_label(target_sub)
+            text_rem = self._flatten(text_rem, length_res)
+            text_mas = self._flatten(text_mas, length_sub)
+            loss_ori = self.loss_func(text_pre, label_flatten)
+            loss_res = self.loss_func(text_rem, label_flatten_res)
+            loss_mas = self.loss_func(text_mas, label_flatten_sub)
+            loss = loss_ori + loss_res * self.weight_res + loss_mas * self.weight_mas
+        return {"loss": loss}

ppocr/losses/stroke_focus_loss.py

@@ -0,0 +1,63 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/text-gestalt/loss/stroke_focus_loss.py
+"""
+import cv2
+import sys
+import time
+import string
+import random
+import numpy as np
+import paddle.nn as nn
+import paddle
+
+
+class StrokeFocusLoss(nn.Layer):
+    def __init__(self, character_dict_path=None, **kwargs):
+        super(StrokeFocusLoss, self).__init__(character_dict_path)
+        self.mse_loss = nn.MSELoss()
+        self.ce_loss = nn.CrossEntropyLoss()
+        self.l1_loss = nn.L1Loss()
+        self.english_stroke_alphabet = "0123456789"
+        self.english_stroke_dict = {}
+        for index in range(len(self.english_stroke_alphabet)):
+            self.english_stroke_dict[self.english_stroke_alphabet[index]] = index
+
+        stroke_decompose_lines = open(character_dict_path, "r").readlines()
+        self.dic = {}
+        for line in stroke_decompose_lines:
+            line = line.strip()
+            character, sequence = line.split()
+            self.dic[character] = sequence
+
+    def forward(self, pred, data):
+        sr_img = pred["sr_img"]
+        hr_img = pred["hr_img"]
+
+        mse_loss = self.mse_loss(sr_img, hr_img)
+        word_attention_map_gt = pred["word_attention_map_gt"]
+        word_attention_map_pred = pred["word_attention_map_pred"]
+
+        hr_pred = pred["hr_pred"]
+        sr_pred = pred["sr_pred"]
+
+        attention_loss = paddle.nn.functional.l1_loss(
+            word_attention_map_gt, word_attention_map_pred
+        )
+
+        loss = (mse_loss + attention_loss * 50) * 100
+
+        return {"mse_loss": mse_loss, "attention_loss": attention_loss, "loss": loss}

ppocr/losses/table_att_loss.py

@@ -0,0 +1,100 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import paddle
+from paddle import nn
+from paddle.nn import functional as F
+
+
+class TableAttentionLoss(nn.Layer):
+    def __init__(self, structure_weight=1.0, loc_weight=0.0, **kwargs):
+        super(TableAttentionLoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(weight=None, reduction="none")
+        self.structure_weight = structure_weight
+        self.loc_weight = loc_weight
+
+    def forward(self, predicts, batch):
+        structure_probs = predicts["structure_probs"]
+        structure_targets = batch[1].astype("int64")
+        structure_targets = structure_targets[:, 1:]
+        structure_probs = paddle.reshape(
+            structure_probs, [-1, structure_probs.shape[-1]]
+        )
+        structure_targets = paddle.reshape(structure_targets, [-1])
+        structure_loss = self.loss_func(structure_probs, structure_targets)
+
+        structure_loss = paddle.mean(structure_loss) * self.structure_weight
+
+        loc_preds = predicts["loc_preds"]
+        loc_targets = batch[2].astype("float32")
+        loc_targets_mask = batch[3].astype("float32")
+        loc_targets = loc_targets[:, 1:, :]
+        loc_targets_mask = loc_targets_mask[:, 1:, :]
+        loc_loss = (
+            F.mse_loss(loc_preds * loc_targets_mask, loc_targets) * self.loc_weight
+        )
+
+        total_loss = structure_loss + loc_loss
+        return {
+            "loss": total_loss,
+            "structure_loss": structure_loss,
+            "loc_loss": loc_loss,
+        }
+
+
+class SLALoss(nn.Layer):
+    def __init__(self, structure_weight=1.0, loc_weight=0.0, loc_loss="mse", **kwargs):
+        super(SLALoss, self).__init__()
+        self.loss_func = nn.CrossEntropyLoss(weight=None, reduction="mean")
+        self.structure_weight = structure_weight
+        self.loc_weight = loc_weight
+        self.loc_loss = loc_loss
+        self.eps = 1e-12
+
+    def forward(self, predicts, batch):
+        structure_probs = predicts["structure_probs"]
+        structure_targets = batch[1].astype("int64")
+        max_len = batch[-2].max().astype("int32")
+        structure_targets = structure_targets[:, 1 : max_len + 2]
+
+        structure_loss = self.loss_func(structure_probs, structure_targets)
+
+        structure_loss = paddle.mean(structure_loss) * self.structure_weight
+
+        loc_preds = predicts["loc_preds"]
+        loc_targets = batch[2].astype("float32")
+        loc_targets_mask = batch[3].astype("float32")
+        loc_targets = loc_targets[:, 1 : max_len + 2]
+        loc_targets_mask = loc_targets_mask[:, 1 : max_len + 2]
+
+        loc_loss = (
+            F.smooth_l1_loss(
+                loc_preds * loc_targets_mask,
+                loc_targets * loc_targets_mask,
+                reduction="sum",
+            )
+            * self.loc_weight
+        )
+
+        loc_loss = loc_loss / (loc_targets_mask.sum() + self.eps)
+        total_loss = structure_loss + loc_loss
+        return {
+            "loss": total_loss,
+            "structure_loss": structure_loss,
+            "loc_loss": loc_loss,
+        }

ppocr/losses/table_master_loss.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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.
+"""
+This code is refer from:
+https://github.com/JiaquanYe/TableMASTER-mmocr/tree/master/mmocr/models/textrecog/losses
+"""
+
+import paddle
+from paddle import nn
+
+
+class TableMasterLoss(nn.Layer):
+    def __init__(self, ignore_index=-1):
+        super(TableMasterLoss, self).__init__()
+        self.structure_loss = nn.CrossEntropyLoss(
+            ignore_index=ignore_index, reduction="mean"
+        )
+        self.box_loss = nn.L1Loss(reduction="sum")
+        self.eps = 1e-12
+
+    def forward(self, predicts, batch):
+        # structure_loss
+        structure_probs = predicts["structure_probs"]
+        structure_targets = batch[1]
+        structure_targets = structure_targets[:, 1:]
+        structure_probs = structure_probs.reshape([-1, structure_probs.shape[-1]])
+        structure_targets = structure_targets.reshape([-1])
+
+        structure_loss = self.structure_loss(structure_probs, structure_targets)
+        structure_loss = structure_loss.mean()
+        losses = dict(structure_loss=structure_loss)
+
+        # box loss
+        bboxes_preds = predicts["loc_preds"]
+        bboxes_targets = batch[2][:, 1:, :]
+        bbox_masks = batch[3][:, 1:]
+        # mask empty-bbox or non-bbox structure token's bbox.
+
+        masked_bboxes_preds = bboxes_preds * bbox_masks
+        masked_bboxes_targets = bboxes_targets * bbox_masks
+
+        # horizon loss (x and width)
+        horizon_sum_loss = self.box_loss(
+            masked_bboxes_preds[:, :, 0::2], masked_bboxes_targets[:, :, 0::2]
+        )
+        horizon_loss = horizon_sum_loss / (bbox_masks.sum() + self.eps)
+        # vertical loss (y and height)
+        vertical_sum_loss = self.box_loss(
+            masked_bboxes_preds[:, :, 1::2], masked_bboxes_targets[:, :, 1::2]
+        )
+        vertical_loss = vertical_sum_loss / (bbox_masks.sum() + self.eps)
+
+        horizon_loss = horizon_loss.mean()
+        vertical_loss = vertical_loss.mean()
+        all_loss = structure_loss + horizon_loss + vertical_loss
+        losses.update(
+            {
+                "loss": all_loss,
+                "horizon_bbox_loss": horizon_loss,
+                "vertical_bbox_loss": vertical_loss,
+            }
+        )
+        return losses

ppocr/losses/text_focus_loss.py

@@ -0,0 +1,91 @@
+# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+"""
+This code is refer from:
+https://github.com/FudanVI/FudanOCR/blob/main/scene-text-telescope/loss/text_focus_loss.py
+"""
+
+import paddle.nn as nn
+import paddle
+import numpy as np
+import pickle as pkl
+
+standard_alphebet = "-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+standard_dict = {}
+for index in range(len(standard_alphebet)):
+    standard_dict[standard_alphebet[index]] = index
+
+
+def load_confuse_matrix(confuse_dict_path):
+    f = open(confuse_dict_path, "rb")
+    data = pkl.load(f)
+    f.close()
+    number = data[:10]
+    upper = data[10:36]
+    lower = data[36:]
+    end = np.ones((1, 62))
+    pad = np.ones((63, 1))
+    rearrange_data = np.concatenate((end, number, lower, upper), axis=0)
+    rearrange_data = np.concatenate((pad, rearrange_data), axis=1)
+    rearrange_data = 1 / rearrange_data
+    rearrange_data[rearrange_data == np.inf] = 1
+    rearrange_data = paddle.to_tensor(rearrange_data)
+
+    lower_alpha = "abcdefghijklmnopqrstuvwxyz"
+    # upper_alpha = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    for i in range(63):
+        for j in range(63):
+            if i != j and standard_alphebet[j] in lower_alpha:
+                rearrange_data[i][j] = max(
+                    rearrange_data[i][j], rearrange_data[i][j + 26]
+                )
+    rearrange_data = rearrange_data[:37, :37]
+
+    return rearrange_data
+
+
+def weight_cross_entropy(pred, gt, weight_table):
+    batch = gt.shape[0]
+    weight = weight_table[gt]
+    pred_exp = paddle.exp(pred)
+    pred_exp_weight = weight * pred_exp
+    loss = 0
+    for i in range(len(gt)):
+        loss -= paddle.log(
+            pred_exp_weight[i][gt[i]] / paddle.sum(pred_exp_weight, 1)[i]
+        )
+    return loss / batch
+
+
+class TelescopeLoss(nn.Layer):
+    def __init__(self, confuse_dict_path):
+        super(TelescopeLoss, self).__init__()
+        self.weight_table = load_confuse_matrix(confuse_dict_path)
+        self.mse_loss = nn.MSELoss()
+        self.ce_loss = nn.CrossEntropyLoss()
+        self.l1_loss = nn.L1Loss()
+
+    def forward(self, pred, data):
+        sr_img = pred["sr_img"]
+        hr_img = pred["hr_img"]
+        sr_pred = pred["sr_pred"]
+        text_gt = pred["text_gt"]
+
+        word_attention_map_gt = pred["word_attention_map_gt"]
+        word_attention_map_pred = pred["word_attention_map_pred"]
+        mse_loss = self.mse_loss(sr_img, hr_img)
+        attention_loss = self.l1_loss(word_attention_map_gt, word_attention_map_pred)
+        recognition_loss = weight_cross_entropy(sr_pred, text_gt, self.weight_table)
+        loss = mse_loss + attention_loss * 10 + recognition_loss * 0.0005
+        return {"mse_loss": mse_loss, "attention_loss": attention_loss, "loss": loss}

ppocr/losses/vqa_token_layoutlm_loss.py

@@ -0,0 +1,61 @@
+# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from paddle import nn
+from ppocr.losses.basic_loss import DMLLoss
+
+
+class VQASerTokenLayoutLMLoss(nn.Layer):
+    def __init__(self, num_classes, key=None):
+        super().__init__()
+        self.loss_class = nn.CrossEntropyLoss()
+        self.num_classes = num_classes
+        self.ignore_index = self.loss_class.ignore_index
+        self.key = key
+
+    def forward(self, predicts, batch):
+        if isinstance(predicts, dict) and self.key is not None:
+            predicts = predicts[self.key]
+        labels = batch[5]
+        attention_mask = batch[2]
+        if attention_mask is not None:
+            active_loss = (
+                attention_mask.reshape(
+                    [
+                        -1,
+                    ]
+                )
+                == 1
+            )
+            active_output = predicts.reshape([-1, self.num_classes])[active_loss]
+            active_label = labels.reshape(
+                [
+                    -1,
+                ]
+            )[active_loss]
+            loss = self.loss_class(active_output, active_label)
+        else:
+            loss = self.loss_class(
+                predicts.reshape([-1, self.num_classes]),
+                labels.reshape(
+                    [
+                        -1,
+                    ]
+                ),
+            )
+        return {"loss": loss}

ppocr/metrics/__init__.py

@@ -0,0 +1,62 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import copy
+
+__all__ = ["build_metric"]
+
+from .det_metric import DetMetric, DetFCEMetric
+from .rec_metric import RecMetric, CNTMetric, CANMetric, LaTeXOCRMetric
+from .cls_metric import ClsMetric
+from .e2e_metric import E2EMetric
+from .distillation_metric import DistillationMetric
+from .table_metric import TableMetric
+from .kie_metric import KIEMetric
+from .vqa_token_ser_metric import VQASerTokenMetric
+from .vqa_token_re_metric import VQAReTokenMetric
+from .sr_metric import SRMetric
+from .ct_metric import CTMetric
+
+
+def build_metric(config):
+    support_dict = [
+        "DetMetric",
+        "DetFCEMetric",
+        "RecMetric",
+        "ClsMetric",
+        "E2EMetric",
+        "DistillationMetric",
+        "TableMetric",
+        "KIEMetric",
+        "VQASerTokenMetric",
+        "VQAReTokenMetric",
+        "SRMetric",
+        "CTMetric",
+        "CNTMetric",
+        "CANMetric",
+        "LaTeXOCRMetric",
+    ]
+
+    config = copy.deepcopy(config)
+    module_name = config.pop("name")
+    assert module_name in support_dict, Exception(
+        "metric only support {}".format(support_dict)
+    )
+    module_class = eval(module_name)(**config)
+    return module_class

ppocr/metrics/bleu.py

@@ -0,0 +1,243 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+"""
+This code is refer from:
+https://github.com/tensorflow/nmt/blob/master/nmt/scripts/bleu.py
+"""
+
+import re
+import math
+import collections
+from functools import lru_cache
+
+
+def _get_ngrams(segment, max_order):
+    """Extracts all n-grams upto a given maximum order from an input segment.
+
+    Args:
+      segment: text segment from which n-grams will be extracted.
+      max_order: maximum length in tokens of the n-grams returned by this
+          methods.
+
+    Returns:
+      The Counter containing all n-grams upto max_order in segment
+      with a count of how many times each n-gram occurred.
+    """
+    ngram_counts = collections.Counter()
+    for order in range(1, max_order + 1):
+        for i in range(0, len(segment) - order + 1):
+            ngram = tuple(segment[i : i + order])
+            ngram_counts[ngram] += 1
+    return ngram_counts
+
+
+def compute_bleu(reference_corpus, translation_corpus, max_order=4, smooth=False):
+    """Computes BLEU score of translated segments against one or more references.
+
+    Args:
+      reference_corpus: list of lists of references for each translation. Each
+          reference should be tokenized into a list of tokens.
+      translation_corpus: list of translations to score. Each translation
+          should be tokenized into a list of tokens.
+      max_order: Maximum n-gram order to use when computing BLEU score.
+      smooth: Whether or not to apply Lin et al. 2004 smoothing.
+
+    Returns:
+      3-Tuple with the BLEU score, n-gram precisions, geometric mean of n-gram
+      precisions and brevity penalty.
+    """
+    matches_by_order = [0] * max_order
+    possible_matches_by_order = [0] * max_order
+    reference_length = 0
+    translation_length = 0
+    for references, translation in zip(reference_corpus, translation_corpus):
+        reference_length += min(len(r) for r in references)
+        translation_length += len(translation)
+
+        merged_ref_ngram_counts = collections.Counter()
+        for reference in references:
+            merged_ref_ngram_counts |= _get_ngrams(reference, max_order)
+        translation_ngram_counts = _get_ngrams(translation, max_order)
+        overlap = translation_ngram_counts & merged_ref_ngram_counts
+        for ngram in overlap:
+            matches_by_order[len(ngram) - 1] += overlap[ngram]
+        for order in range(1, max_order + 1):
+            possible_matches = len(translation) - order + 1
+            if possible_matches > 0:
+                possible_matches_by_order[order - 1] += possible_matches
+
+    precisions = [0] * max_order
+    for i in range(0, max_order):
+        if smooth:
+            precisions[i] = (matches_by_order[i] + 1.0) / (
+                possible_matches_by_order[i] + 1.0
+            )
+        else:
+            if possible_matches_by_order[i] > 0:
+                precisions[i] = (
+                    float(matches_by_order[i]) / possible_matches_by_order[i]
+                )
+            else:
+                precisions[i] = 0.0
+
+    if min(precisions) > 0:
+        p_log_sum = sum((1.0 / max_order) * math.log(p) for p in precisions)
+        geo_mean = math.exp(p_log_sum)
+    else:
+        geo_mean = 0
+
+    if float(translation_length) == 0 or float(reference_length) == 0:
+        ratio = 1e-5
+    else:
+        ratio = float(translation_length) / reference_length
+
+    if ratio > 1.0:
+        bp = 1.0
+    else:
+        bp = math.exp(1 - 1.0 / ratio)
+
+    bleu = geo_mean * bp
+
+    return (bleu, precisions, bp, ratio, translation_length, reference_length)
+
+
+class BaseTokenizer:
+    """A base dummy tokenizer to derive from."""
+
+    def signature(self):
+        """
+        Returns a signature for the tokenizer.
+        :return: signature string
+        """
+        return "none"
+
+    def __call__(self, line):
+        """
+        Tokenizes an input line with the tokenizer.
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+        return line
+
+
+class TokenizerRegexp(BaseTokenizer):
+    def signature(self):
+        return "re"
+
+    def __init__(self):
+        self._re = [
+            # language-dependent part (assuming Western languages)
+            (re.compile(r"([\{-\~\[-\` -\&\(-\+\:-\@\/])"), r" \1 "),
+            # tokenize period and comma unless preceded by a digit
+            (re.compile(r"([^0-9])([\.,])"), r"\1 \2 "),
+            # tokenize period and comma unless followed by a digit
+            (re.compile(r"([\.,])([^0-9])"), r" \1 \2"),
+            # tokenize dash when preceded by a digit
+            (re.compile(r"([0-9])(-)"), r"\1 \2 "),
+            # one space only between words
+            # NOTE: Doing this in Python (below) is faster
+            # (re.compile(r'\s+'), r' '),
+        ]
+
+    @lru_cache(maxsize=2**16)
+    def __call__(self, line):
+        """Common post-processing tokenizer for `13a` and `zh` tokenizers.
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+        for _re, repl in self._re:
+            line = _re.sub(repl, line)
+
+        # no leading or trailing spaces, single space within words
+        # return ' '.join(line.split())
+        # This line is changed with regards to the original tokenizer (seen above) to return individual words
+        return line.split()
+
+
+class Tokenizer13a(BaseTokenizer):
+    def signature(self):
+        return "13a"
+
+    def __init__(self):
+        self._post_tokenizer = TokenizerRegexp()
+
+    @lru_cache(maxsize=2**16)
+    def __call__(self, line):
+        """Tokenizes an input line using a relatively minimal tokenization
+        that is however equivalent to mteval-v13a, used by WMT.
+
+        :param line: a segment to tokenize
+        :return: the tokenized line
+        """
+
+        # language-independent part:
+        line = line.replace("<skipped>", "")
+        line = line.replace("-\n", "")
+        line = line.replace("\n", " ")
+
+        if "&" in line:
+            line = line.replace("&quot;", '"')
+            line = line.replace("&amp;", "&")
+            line = line.replace("&lt;", "<")
+            line = line.replace("&gt;", ">")
+
+        return self._post_tokenizer(f" {line} ")
+
+
+def compute_bleu_score(
+    predictions, references, tokenizer=Tokenizer13a(), max_order=4, smooth=False
+):
+    # if only one reference is provided make sure we still use list of lists
+    if isinstance(references[0], str):
+        references = [[ref] for ref in references]
+
+    references = [[tokenizer(r) for r in ref] for ref in references]
+    predictions = [tokenizer(p) for p in predictions]
+    score = compute_bleu(
+        reference_corpus=references,
+        translation_corpus=predictions,
+        max_order=max_order,
+        smooth=smooth,
+    )
+    (bleu, precisions, bp, ratio, translation_length, reference_length) = score
+    return bleu
+
+
+def cal_distance(word1, word2):
+    m = len(word1)
+    n = len(word2)
+    if m * n == 0:
+        return m + n
+    dp = [[0] * (n + 1) for _ in range(m + 1)]
+    for i in range(m + 1):
+        dp[i][0] = i
+    for j in range(n + 1):
+        dp[0][j] = j
+    for i in range(1, m + 1):
+        for j in range(1, n + 1):
+            a = dp[i - 1][j] + 1
+            b = dp[i][j - 1] + 1
+            c = dp[i - 1][j - 1]
+            if word1[i - 1] != word2[j - 1]:
+                c += 1
+            dp[i][j] = min(a, b, c)
+    return dp[m][n]
+
+
+def compute_edit_distance(prediction, label):
+    prediction = prediction.strip().split(" ")
+    label = label.strip().split(" ")
+    distance = cal_distance(prediction, label)
+    return distance

ppocr/metrics/cls_metric.py

@@ -0,0 +1,48 @@
+# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
+#
+# 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
+#
+#    http://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.
+
+
+class ClsMetric(object):
+    def __init__(self, main_indicator="acc", **kwargs):
+        self.main_indicator = main_indicator
+        self.eps = 1e-5
+        self.reset()
+
+    def __call__(self, pred_label, *args, **kwargs):
+        preds, labels = pred_label
+        correct_num = 0
+        all_num = 0
+        for (pred, pred_conf), (target, _) in zip(preds, labels):
+            if pred == target:
+                correct_num += 1
+            all_num += 1
+        self.correct_num += correct_num
+        self.all_num += all_num
+        return {
+            "acc": correct_num / (all_num + self.eps),
+        }
+
+    def get_metric(self):
+        """
+        return metrics {
+                 'acc': 0
+            }
+        """
+        acc = self.correct_num / (self.all_num + self.eps)
+        self.reset()
+        return {"acc": acc}
+
+    def reset(self):
+        self.correct_num = 0
+        self.all_num = 0

ppocr/metrics/ct_metric.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+from scipy import io
+import numpy as np
+
+from ppocr.utils.e2e_metric.Deteval import combine_results, get_score_C
+
+
+class CTMetric(object):
+    def __init__(self, main_indicator, delimiter="\t", **kwargs):
+        self.delimiter = delimiter
+        self.main_indicator = main_indicator
+        self.reset()
+
+    def reset(self):
+        self.results = []  # clear results
+
+    def __call__(self, preds, batch, **kwargs):
+        # NOTE: only support bs=1 now, as the label length of different sample is Unequal
+        assert len(preds) == 1, "CentripetalText test now only support batch_size=1."
+        label = batch[2]
+        text = batch[3]
+        pred = preds[0]["points"]
+        result = get_score_C(label, text, pred)
+
+        self.results.append(result)
+
+    def get_metric(self):
+        """
+        Input format: y0,x0, ..... yn,xn. Each detection is separated by the end of line token ('\n')'
+        """
+        metrics = combine_results(self.results, rec_flag=False)
+        self.reset()
+        return metrics

ppocr/metrics/det_metric.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+#
+#     http://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 __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+__all__ = ["DetMetric", "DetFCEMetric"]
+
+from .eval_det_iou import DetectionIoUEvaluator
+
+
+class DetMetric(object):
+    def __init__(self, main_indicator="hmean", **kwargs):
+        self.evaluator = DetectionIoUEvaluator()
+        self.main_indicator = main_indicator
+        self.reset()
+
+    def __call__(self, preds, batch, **kwargs):
+        """
+        batch: a list produced by dataloaders.
+            image: np.ndarray  of shape (N, C, H, W).
+            ratio_list: np.ndarray  of shape(N,2)
+            polygons: np.ndarray  of shape (N, K, 4, 2), the polygons of objective regions.
+            ignore_tags: np.ndarray  of shape (N, K), indicates whether a region is ignorable or not.
+        preds: a list of dict produced by post process
+             points: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions.
+        """
+        gt_polyons_batch = batch[2]
+        ignore_tags_batch = batch[3]
+        for pred, gt_polyons, ignore_tags in zip(
+            preds, gt_polyons_batch, ignore_tags_batch
+        ):
+            # prepare gt
+            gt_info_list = [
+                {"points": gt_polyon, "text": "", "ignore": ignore_tag}
+                for gt_polyon, ignore_tag in zip(gt_polyons, ignore_tags)
+            ]
+            # prepare det
+            det_info_list = [
+                {"points": det_polyon, "text": ""} for det_polyon in pred["points"]
+            ]
+            result = self.evaluator.evaluate_image(gt_info_list, det_info_list)
+            self.results.append(result)
+
+    def get_metric(self):
+        """
+        return metrics {
+                 'precision': 0,
+                 'recall': 0,
+                 'hmean': 0
+            }
+        """
+
+        metrics = self.evaluator.combine_results(self.results)
+        self.reset()
+        return metrics
+
+    def reset(self):
+        self.results = []  # clear results
+
+
+class DetFCEMetric(object):
+    def __init__(self, main_indicator="hmean", **kwargs):
+        self.evaluator = DetectionIoUEvaluator()
+        self.main_indicator = main_indicator
+        self.reset()
+
+    def __call__(self, preds, batch, **kwargs):
+        """
+        batch: a list produced by dataloaders.
+            image: np.ndarray  of shape (N, C, H, W).
+            ratio_list: np.ndarray  of shape(N,2)
+            polygons: np.ndarray  of shape (N, K, 4, 2), the polygons of objective regions.
+            ignore_tags: np.ndarray  of shape (N, K), indicates whether a region is ignorable or not.
+        preds: a list of dict produced by post process
+             points: np.ndarray of shape (N, K, 4, 2), the polygons of objective regions.
+        """
+        gt_polyons_batch = batch[2]
+        ignore_tags_batch = batch[3]
+
+        for pred, gt_polyons, ignore_tags in zip(
+            preds, gt_polyons_batch, ignore_tags_batch
+        ):
+            # prepare gt
+            gt_info_list = [
+                {"points": gt_polyon, "text": "", "ignore": ignore_tag}
+                for gt_polyon, ignore_tag in zip(gt_polyons, ignore_tags)
+            ]
+            # prepare det
+            det_info_list = [
+                {"points": det_polyon, "text": "", "score": score}
+                for det_polyon, score in zip(pred["points"], pred["scores"])
+            ]
+
+            for score_thr in self.results.keys():
+                det_info_list_thr = [
+                    det_info
+                    for det_info in det_info_list
+                    if det_info["score"] >= score_thr
+                ]
+                result = self.evaluator.evaluate_image(gt_info_list, det_info_list_thr)
+                self.results[score_thr].append(result)
+
+    def get_metric(self):
+        """
+        return metrics {'heman':0,
+            'thr 0.3':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.4':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.5':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.6':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.7':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.8':'precision: 0 recall: 0 hmean: 0',
+            'thr 0.9':'precision: 0 recall: 0 hmean: 0',
+            }
+        """
+        metrics = {}
+        hmean = 0
+        for score_thr in self.results.keys():
+            metric = self.evaluator.combine_results(self.results[score_thr])
+            # for key, value in metric.items():
+            #     metrics['{}_{}'.format(key, score_thr)] = value
+            metric_str = "precision:{:.5f} recall:{:.5f} hmean:{:.5f}".format(
+                metric["precision"], metric["recall"], metric["hmean"]
+            )
+            metrics["thr {}".format(score_thr)] = metric_str
+            hmean = max(hmean, metric["hmean"])
+        metrics["hmean"] = hmean
+
+        self.reset()
+        return metrics
+
+    def reset(self):
+        self.results = {
+            0.3: [],
+            0.4: [],
+            0.5: [],
+            0.6: [],
+            0.7: [],
+            0.8: [],
+            0.9: [],
+        }  # clear results