Merge branch 'android_demo' of https://github.com/WenmuZhou/PaddleOCR into android_demo

paddleocr.py

@@ -146,7 +146,8 @@ def parse_args(mMain=True, add_help=True):
         # DB parmas
         parser.add_argument("--det_db_thresh", type=float, default=0.3)
         parser.add_argument("--det_db_box_thresh", type=float, default=0.5)
-        parser.add_argument("--det_db_unclip_ratio", type=float, default=2.0)
+        parser.add_argument("--det_db_unclip_ratio", type=float, default=1.6)
+        parser.add_argument("--use_dilation", type=bool, default=False)
 
         # EAST parmas
         parser.add_argument("--det_east_score_thresh", type=float, default=0.8)
@@ -193,7 +194,8 @@ def parse_args(mMain=True, add_help=True):
             det_limit_type='max',
             det_db_thresh=0.3,
             det_db_box_thresh=0.5,
-            det_db_unclip_ratio=2.0,
+            det_db_unclip_ratio=1.6,
+            use_dilation=False,
             det_east_score_thresh=0.8,
             det_east_cover_thresh=0.1,
             det_east_nms_thresh=0.2,
@@ -234,7 +236,9 @@ class PaddleOCR(predict_system.TextSystem):
         assert lang in model_urls[
             'rec'], 'param lang must in {}, but got {}'.format(
                 model_urls['rec'].keys(), lang)
+        use_inner_dict = False
         if postprocess_params.rec_char_dict_path is None:
+            use_inner_dict = True
             postprocess_params.rec_char_dict_path = model_urls['rec'][lang][
                 'dict_path']
 
@@ -261,7 +265,7 @@ class PaddleOCR(predict_system.TextSystem):
         if postprocess_params.rec_algorithm not in SUPPORT_REC_MODEL:
             logger.error('rec_algorithm must in {}'.format(SUPPORT_REC_MODEL))
             sys.exit(0)
-
+        if use_inner_dict:
             postprocess_params.rec_char_dict_path = str(
                 Path(__file__).parent / postprocess_params.rec_char_dict_path)
 
@@ -280,8 +284,13 @@ class PaddleOCR(predict_system.TextSystem):
         if isinstance(img, list) and det == True:
             logger.error('When input a list of images, det must be false')
             exit(0)
+        if cls == False:
+            self.use_angle_cls = False
+        elif cls == True and self.use_angle_cls == False:
+            logger.warning(
+                'Since the angle classifier is not initialized, the angle classifier will not be uesd during the forward process'
+            )
 
-        self.use_angle_cls = cls
         if isinstance(img, str):
             # download net image
             if img.startswith('http'):

ppocr/data/__init__.py

@@ -34,6 +34,7 @@ import paddle.distributed as dist
 from ppocr.data.imaug import transform, create_operators
 from ppocr.data.simple_dataset import SimpleDataSet
 from ppocr.data.lmdb_dataset import LMDBDataSet
+from ppocr.data.pgnet_dataset import PGDataSet
 
 __all__ = ['build_dataloader', 'transform', 'create_operators']
 
@@ -54,7 +55,7 @@ signal.signal(signal.SIGTERM, term_mp)
 def build_dataloader(config, mode, device, logger, seed=None):
     config = copy.deepcopy(config)
 
-    support_dict = ['SimpleDataSet', 'LMDBDataSet']
+    support_dict = ['SimpleDataSet', 'LMDBDataSet', 'PGDataSet']
     module_name = config[mode]['dataset']['name']
     assert module_name in support_dict, Exception(
         'DataSet only support {}'.format(support_dict))
@@ -72,14 +73,14 @@ def build_dataloader(config, mode, device, logger, seed=None):
     else:
         use_shared_memory = True
     if mode == "Train":
-        #Distribute data to multiple cards
+        # Distribute data to multiple cards
         batch_sampler = DistributedBatchSampler(
             dataset=dataset,
             batch_size=batch_size,
             shuffle=shuffle,
             drop_last=drop_last)
     else:
-        #Distribute data to single card
+        # Distribute data to single card
         batch_sampler = BatchSampler(
             dataset=dataset,
             batch_size=batch_size,

ppocr/data/imaug/__init__.py

@@ -28,6 +28,7 @@ from .label_ops import *
 
 from .east_process import *
 from .sast_process import *
+from .pg_process import *
 
 
 def transform(data, ops=None):

ppocr/data/imaug/label_ops.py

@@ -187,6 +187,34 @@ class CTCLabelEncode(BaseRecLabelEncode):
         return dict_character
 
 
+class E2ELabelEncode(BaseRecLabelEncode):
+    def __init__(self,
+                 max_text_length,
+                 character_dict_path=None,
+                 character_type='EN',
+                 use_space_char=False,
+                 **kwargs):
+        super(E2ELabelEncode,
+              self).__init__(max_text_length, character_dict_path,
+                             character_type, use_space_char)
+        self.pad_num = len(self.dict)  # the length to pad
+
+    def __call__(self, data):
+        text_label_index_list, temp_text = [], []
+        texts = data['strs']
+        for text in texts:
+            text = text.lower()
+            temp_text = []
+            for c_ in text:
+                if c_ in self.dict:
+                    temp_text.append(self.dict[c_])
+            temp_text = temp_text + [self.pad_num] * (self.max_text_len -
+                                                      len(temp_text))
+            text_label_index_list.append(temp_text)
+        data['strs'] = np.array(text_label_index_list)
+        return data
+
+
 class AttnLabelEncode(BaseRecLabelEncode):
     """ Convert between text-label and text-index """
 
@@ -215,7 +243,7 @@ class AttnLabelEncode(BaseRecLabelEncode):
             return None
         data['length'] = np.array(len(text))
         text = [0] + text + [len(self.character) - 1] + [0] * (self.max_text_len
-                                                               - len(text) - 1)
+                                                               - len(text) - 2)
         data['label'] = np.array(text)
         return data
 
@@ -261,7 +289,7 @@ class SRNLabelEncode(BaseRecLabelEncode):
         if len(text) > self.max_text_len:
             return None
         data['length'] = np.array(len(text))
-        text = text + [char_num] * (self.max_text_len - len(text))
+        text = text + [char_num - 1] * (self.max_text_len - len(text))
         data['label'] = np.array(text)
         return data
 

ppocr/data/imaug/make_shrink_map.py

@@ -32,7 +32,6 @@ class MakeShrinkMap(object):
         text_polys, ignore_tags = self.validate_polygons(text_polys,
                                                          ignore_tags, h, w)
         gt = np.zeros((h, w), dtype=np.float32)
-        # gt = np.zeros((1, h, w), dtype=np.float32)
         mask = np.ones((h, w), dtype=np.float32)
         for i in range(len(text_polys)):
             polygon = text_polys[i]
@@ -44,21 +43,34 @@ class MakeShrinkMap(object):
                 ignore_tags[i] = True
             else:
                 polygon_shape = Polygon(polygon)
-                distance = polygon_shape.area * (
-                    1 - np.power(self.shrink_ratio, 2)) / polygon_shape.length
-                subject = [tuple(l) for l in text_polys[i]]
+                subject = [tuple(l) for l in polygon]
                 padding = pyclipper.PyclipperOffset()
                 padding.AddPath(subject, pyclipper.JT_ROUND,
                                 pyclipper.ET_CLOSEDPOLYGON)
+                shrinked = []
+
+                # 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
                     shrinked = padding.Execute(-distance)
+                    if len(shrinked) == 1:
+                        break
+
                 if shrinked == []:
                     cv2.fillPoly(mask,
                                  polygon.astype(np.int32)[np.newaxis, :, :], 0)
                     ignore_tags[i] = True
                     continue
-                shrinked = np.array(shrinked[0]).reshape(-1, 2)
-                cv2.fillPoly(gt, [shrinked.astype(np.int32)], 1)
-                # cv2.fillPoly(gt[0], [shrinked.astype(np.int32)], 1)
+
+                for each_shirnk in shrinked:
+                    shirnk = np.array(each_shirnk).reshape(-1, 2)
+                    cv2.fillPoly(gt, [shirnk.astype(np.int32)], 1)
 
         data['shrink_map'] = gt
         data['shrink_mask'] = mask
@@ -84,11 +96,12 @@ class MakeShrinkMap(object):
         return polygons, ignore_tags
 
     def polygon_area(self, polygon):
-        # return cv2.contourArea(polygon.astype(np.float32))
-        edge = 0
-        for i in range(polygon.shape[0]):
-            next_index = (i + 1) % polygon.shape[0]
-            edge += (polygon[next_index, 0] - polygon[i, 0]) * (
-                polygon[next_index, 1] - polygon[i, 1])
-
-        return edge / 2.
+        """
+        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

@@ -185,8 +185,8 @@ class DetResizeForTest(object):
         resize_h = int(h * ratio)
         resize_w = int(w * ratio)
 
-        resize_h = int(round(resize_h / 32) * 32)
-        resize_w = int(round(resize_w / 32) * 32)
+        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:
@@ -197,7 +197,6 @@ class DetResizeForTest(object):
             sys.exit(0)
         ratio_h = resize_h / float(h)
         ratio_w = resize_w / float(w)
-        # return img, np.array([h, w])
         return img, [ratio_h, ratio_w]
 
     def resize_image_type2(self, img):
@@ -206,7 +205,6 @@ class DetResizeForTest(object):
         resize_w = w
         resize_h = h
 
-        # Fix the longer side
         if resize_h > resize_w:
             ratio = float(self.resize_long) / resize_h
         else:
@@ -223,3 +221,72 @@ class DetResizeForTest(object):
         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)

ppocr/data/imaug/randaugment.py

@@ -117,13 +117,16 @@ class RawRandAugment(object):
 class RandAugment(RawRandAugment):
     """ RandAugment wrapper to auto fit different img types """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, prob=0.5, *args, **kwargs):
+        self.prob = prob
         if six.PY2:
             super(RandAugment, self).__init__(*args, **kwargs)
         else:
             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)

ppocr/data/pgnet_dataset.py

+# 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']
+
+        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)
+        self.data_format = dataset_config.get('data_format', 'icdar')
+        assert len(
+            ratio_list
+        ) == data_source_num, "The length of ratio_list should be the same as the file_list."
+        self.do_shuffle = loader_config['shuffle']
+
+        logger.info("Initialize indexs of datasets:%s" % label_file_list)
+        self.data_lines = self.get_image_info_list(label_file_list, ratio_list,
+                                                   self.data_format)
+        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)
+
+    def shuffle_data_random(self):
+        if self.do_shuffle:
+            random.seed(self.seed)
+            random.shuffle(self.data_lines)
+        return
+
+    def extract_polys(self, poly_txt_path):
+        """
+        Read text_polys, txt_tags, txts from give txt file.
+        """
+        text_polys, txt_tags, txts = [], [], []
+        with open(poly_txt_path) as f:
+            for line in f.readlines():
+                poly_str, txt = line.strip().split('\t')
+                poly = list(map(float, poly_str.split(',')))
+                if self.mode.lower() == "eval":
+                    while len(poly) < 100:
+                        poly.append(-1)
+                text_polys.append(
+                    np.array(
+                        poly, dtype=np.float32).reshape(-1, 2))
+                txts.append(txt)
+                txt_tags.append(txt == '###')
+
+        return np.array(list(map(np.array, text_polys))), \
+               np.array(txt_tags, dtype=np.bool), txts
+
+    def extract_info_textnet(self, im_fn, img_dir=''):
+        """
+        Extract information from line in textnet format.
+        """
+        info_list = im_fn.split('\t')
+        img_path = ''
+        for ext in [
+                'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif', 'JPG'
+        ]:
+            if os.path.exists(os.path.join(img_dir, info_list[0] + "." + ext)):
+                img_path = os.path.join(img_dir, info_list[0] + "." + ext)
+                break
+
+        if img_path == '':
+            print('Image {0} NOT found in {1}, and it will be ignored.'.format(
+                info_list[0], img_dir))
+
+        nBox = (len(info_list) - 1) // 9
+        wordBBs, txts, txt_tags = [], [], []
+        for n in range(0, nBox):
+            wordBB = list(map(float, info_list[n * 9 + 1:(n + 1) * 9]))
+            txt = info_list[(n + 1) * 9]
+            wordBBs.append([[wordBB[0], wordBB[1]], [wordBB[2], wordBB[3]],
+                            [wordBB[4], wordBB[5]], [wordBB[6], wordBB[7]]])
+            txts.append(txt)
+            if txt == '###':
+                txt_tags.append(True)
+            else:
+                txt_tags.append(False)
+        return img_path, np.array(wordBBs, dtype=np.float32), txt_tags, txts
+
+    def get_image_info_list(self, file_list, ratio_list, data_format='textnet'):
+        if isinstance(file_list, str):
+            file_list = [file_list]
+        data_lines = []
+        for idx, data_source in enumerate(file_list):
+            image_files = []
+            if data_format == 'icdar':
+                image_files = [(data_source, x) for x in
+                               os.listdir(os.path.join(data_source, 'rgb'))
+                               if x.split('.')[-1] in [
+                                   'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif',
+                                   'tiff', 'gif', 'JPG'
+                               ]]
+            elif data_format == 'textnet':
+                with open(data_source) as f:
+                    image_files = [(data_source, x.strip())
+                                   for x in f.readlines()]
+            else:
+                print("Unrecognized data format...")
+                exit(-1)
+            random.seed(self.seed)
+            image_files = random.sample(
+                image_files, round(len(image_files) * ratio_list[idx]))
+            data_lines.extend(image_files)
+        return data_lines
+
+    def __getitem__(self, idx):
+        file_idx = self.data_idx_order_list[idx]
+        data_path, data_line = self.data_lines[file_idx]
+        try:
+            if self.data_format == 'icdar':
+                im_path = os.path.join(data_path, 'rgb', data_line)
+                if self.mode.lower() == "eval":
+                    poly_path = os.path.join(data_path, 'poly_gt',
+                                             data_line.split('.')[0] + '.txt')
+                else:
+                    poly_path = os.path.join(data_path, 'poly',
+                                             data_line.split('.')[0] + '.txt')
+                text_polys, text_tags, text_strs = self.extract_polys(poly_path)
+            else:
+                image_dir = os.path.join(os.path.dirname(data_path), 'image')
+                im_path, text_polys, text_tags, text_strs = self.extract_info_textnet(
+                    data_line, image_dir)
+
+            data = {
+                'img_path': im_path,
+                'polys': text_polys,
+                'tags': text_tags,
+                'strs': text_strs
+            }
+            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/simple_dataset.py

@@ -23,6 +23,7 @@ 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']
@@ -45,7 +46,7 @@ class SimpleDataSet(Dataset):
         logger.info("Initialize indexs 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":
+        if self.mode == "train" and self.do_shuffle:
             self.shuffle_data_random()
         self.ops = create_operators(dataset_config['transforms'], global_config)
 
@@ -56,6 +57,7 @@ class SimpleDataSet(Dataset):
         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]))
@@ -63,7 +65,6 @@ class SimpleDataSet(Dataset):
         return data_lines
 
     def shuffle_data_random(self):
-        if self.do_shuffle:
         random.seed(self.seed)
         random.shuffle(self.data_lines)
         return
@@ -90,7 +91,10 @@ class SimpleDataSet(Dataset):
                     data_line, e))
             outs = None
         if outs is None:
-            return self.__getitem__(np.random.randint(self.__len__()))
+            # 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):

ppocr/losses/__init__.py

@@ -29,10 +29,11 @@ def build_loss(config):
     # cls loss
     from .cls_loss import ClsLoss
 
+    # e2e loss
+    from .e2e_pg_loss import PGLoss
     support_dict = [
         'DBLoss', 'EASTLoss', 'SASTLoss', 'CTCLoss', 'ClsLoss', 'AttentionLoss',
-        'SRNLoss'
-    ]
+        'SRNLoss', 'PGLoss']
 
     config = copy.deepcopy(config)
     module_name = config.pop('name')

ppocr/losses/det_basic_loss.py

@@ -200,6 +200,6 @@ def ohem_batch(scores, gt_texts, training_masks, ohem_ratio):
                 i, :, :], ohem_ratio))
 
     selected_masks = np.concatenate(selected_masks, 0)
-    selected_masks = paddle.to_variable(selected_masks)
+    selected_masks = paddle.to_tensor(selected_masks)
 
     return selected_masks

ppocr/losses/e2e_pg_loss.py

+# 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, 37])  # len(Lexicon_Table)+1
+        f_tcl_char_fg, f_tcl_char_bg = paddle.split(f_tcl_char, [36, 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, 36 * 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/metrics/__init__.py

@@ -26,8 +26,9 @@ def build_metric(config):
     from .det_metric import DetMetric
     from .rec_metric import RecMetric
     from .cls_metric import ClsMetric
+    from .e2e_metric import E2EMetric
 
-    support_dict = ['DetMetric', 'RecMetric', 'ClsMetric']
+    support_dict = ['DetMetric', 'RecMetric', 'ClsMetric', 'E2EMetric']
 
     config = copy.deepcopy(config)
     module_name = config.pop('name')

ppocr/metrics/e2e_metric.py

+# 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.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+__all__ = ['E2EMetric']
+
+from ppocr.utils.e2e_metric.Deteval import get_socre, combine_results
+from ppocr.utils.e2e_utils.extract_textpoint import get_dict
+
+
+class E2EMetric(object):
+    def __init__(self,
+                 character_dict_path,
+                 main_indicator='f_score_e2e',
+                 **kwargs):
+        self.label_list = get_dict(character_dict_path)
+        self.max_index = len(self.label_list)
+        self.main_indicator = main_indicator
+        self.reset()
+
+    def __call__(self, preds, batch, **kwargs):
+        temp_gt_polyons_batch = batch[2]
+        temp_gt_strs_batch = batch[3]
+        ignore_tags_batch = batch[4]
+        gt_polyons_batch = []
+        gt_strs_batch = []
+
+        temp_gt_polyons_batch = temp_gt_polyons_batch[0].tolist()
+        for temp_list in temp_gt_polyons_batch:
+            t = []
+            for index in temp_list:
+                if index[0] != -1 and index[1] != -1:
+                    t.append(index)
+            gt_polyons_batch.append(t)
+
+        temp_gt_strs_batch = temp_gt_strs_batch[0].tolist()
+        for temp_list in temp_gt_strs_batch:
+            t = ""
+            for index in temp_list:
+                if index < self.max_index:
+                    t += self.label_list[index]
+            gt_strs_batch.append(t)
+
+        for pred, gt_polyons, gt_strs, ignore_tags in zip(
+            [preds], [gt_polyons_batch], [gt_strs_batch], ignore_tags_batch):
+            # prepare gt
+            gt_info_list = [{
+                'points': gt_polyon,
+                'text': gt_str,
+                'ignore': ignore_tag
+            } for gt_polyon, gt_str, ignore_tag in
+                            zip(gt_polyons, gt_strs, ignore_tags)]
+            # prepare det
+            e2e_info_list = [{
+                'points': det_polyon,
+                'text': pred_str
+            } for det_polyon, pred_str in zip(pred['points'], pred['strs'])]
+            result = get_socre(gt_info_list, e2e_info_list)
+            self.results.append(result)
+
+    def get_metric(self):
+        metircs = combine_results(self.results)
+        self.reset()
+        return metircs
+
+    def reset(self):
+        self.results = []  # clear results