Merge branch 'dygraph' into pgnet-readme

configs/det/det_r50_vd_sast_icdar15.yml

@@ -20,6 +20,7 @@ Global:
   infer_img:
   save_res_path: ./output/sast_r50_vd_ic15/predicts_sast.txt
 
+
 Architecture:
   model_type: det
   algorithm: SAST

configs/e2e/e2e_r50_vd_pg.yml

+Global:
+  use_gpu: True
+  epoch_num: 600
+  log_smooth_window: 20
+  print_batch_step: 10
+  save_model_dir: ./output/pgnet_r50_vd_totaltext/
+  save_epoch_step: 10
+  # evaluation is run every 0 iterationss after the 1000th iteration
+  eval_batch_step: [ 0, 1000 ]
+  # 1. If pretrained_model is saved in static mode, such as classification pretrained model
+  #    from static branch, load_static_weights must be set as True.
+  # 2. If you want to finetune the pretrained models we provide in the docs,
+  #    you should set load_static_weights as False.
+  load_static_weights: False
+  cal_metric_during_train: False
+  pretrained_model:
+  checkpoints:
+  save_inference_dir:
+  use_visualdl: False
+  infer_img:
+  valid_set: totaltext # two mode: totaltext valid curved words, partvgg valid non-curved words
+  save_res_path: ./output/pgnet_r50_vd_totaltext/predicts_pgnet.txt
+  character_dict_path: ppocr/utils/ic15_dict.txt
+  character_type: EN
+  max_text_length: 50 # the max length in seq
+  max_text_nums: 30 # the max seq nums in a pic
+  tcl_len: 64
+
+Architecture:
+  model_type: e2e
+  algorithm: PGNet
+  Transform:
+  Backbone:
+    name: ResNet
+    layers: 50
+  Neck:
+    name: PGFPN
+  Head:
+    name: PGHead
+
+Loss:
+  name: PGLoss
+  tcl_bs: 64
+  max_text_length: 50 # the same as Global: max_text_length
+  max_text_nums: 30 # the same as Global：max_text_nums
+  pad_num: 36 # the length of dict for pad
+
+Optimizer:
+  name: Adam
+  beta1: 0.9
+  beta2: 0.999
+  lr:
+    learning_rate: 0.001
+  regularizer:
+    name: 'L2'
+    factor: 0
+
+
+PostProcess:
+  name: PGPostProcess
+  score_thresh: 0.5
+Metric:
+  name: E2EMetric
+  character_dict_path: ppocr/utils/ic15_dict.txt
+  main_indicator: f_score_e2e
+
+Train:
+  dataset:
+    name: PGDataSet
+    label_file_list: [.././train_data/total_text/train/]
+    ratio_list: [1.0]
+    data_format: icdar #two data format: icdar/textnet
+    transforms:
+      - DecodeImage: # load image
+          img_mode: BGR
+          channel_first: False
+      - PGProcessTrain:
+          batch_size: 14  # same as loader: batch_size_per_card
+          min_crop_size: 24
+          min_text_size: 4
+          max_text_size: 512
+      - KeepKeys:
+          keep_keys: [ 'images', 'tcl_maps', 'tcl_label_maps', 'border_maps','direction_maps', 'training_masks', 'label_list', 'pos_list', 'pos_mask' ] # dataloader will return list in this order
+  loader:
+    shuffle: True
+    drop_last: True
+    batch_size_per_card: 14
+    num_workers: 16
+
+Eval:
+  dataset:
+    name: PGDataSet
+    data_dir: ./train_data/
+    label_file_list: [./train_data/total_text/test/]
+    transforms:
+      - DecodeImage: # load image
+          img_mode: RGB
+          channel_first: False
+      - E2ELabelEncode:
+      - E2EResizeForTest:
+          max_side_len: 768
+      - NormalizeImage:
+          scale: 1./255.
+          mean: [ 0.485, 0.456, 0.406 ]
+          std: [ 0.229, 0.224, 0.225 ]
+          order: 'hwc'
+      - ToCHWImage:
+      - KeepKeys:
+          keep_keys: [ 'image', 'shape', 'polys', 'strs', 'tags' ]
+  loader:
+    shuffle: False
+    drop_last: False
+    batch_size_per_card: 1 # must be 1
+    num_workers: 2
\ No newline at end of file

doc/doc_ch/inference.md

@@ -13,6 +13,7 @@ inference 模型（`paddle.jit.save`保存的模型）
     - [检测模型转inference模型](#检测模型转inference模型)
     - [识别模型转inference模型](#识别模型转inference模型)  
     - [方向分类模型转inference模型](#方向分类模型转inference模型)
+    - [端到端模型转inference模型](#端到端模型转inference模型)
 
 - [二、文本检测模型推理](#文本检测模型推理)
     - [1. 超轻量中文检测模型推理](#超轻量中文检测模型推理)
@@ -27,10 +28,13 @@ inference 模型（`paddle.jit.save`保存的模型）
     - [4. 自定义文本识别字典的推理](#自定义文本识别字典的推理)
     - [5. 多语言模型的推理](#多语言模型的推理)
 
-- [四、方向分类模型推理](#方向识别模型推理)
+- [四、端到端模型推理](#端到端模型推理)
+    - [1. PGNet端到端模型推理](#PGNet端到端模型推理)
+
+- [五、方向分类模型推理](#方向识别模型推理)
     - [1. 方向分类模型推理](#方向分类模型推理)
 
-- [五、文本检测、方向分类和文字识别串联推理](#文本检测、方向分类和文字识别串联推理)
+- [六、文本检测、方向分类和文字识别串联推理](#文本检测、方向分类和文字识别串联推理)
     - [1. 超轻量中文OCR模型推理](#超轻量中文OCR模型推理)
     - [2. 其他模型推理](#其他模型推理)
 
@@ -118,6 +122,32 @@ python3 tools/export_model.py -c configs/cls/cls_mv3.yml -o Global.pretrained_mo
     ├── inference.pdiparams.info    # 分类inference模型的参数信息，可忽略
     └── inference.pdmodel           # 分类inference模型的program文件
 ```
+<a name="端到端模型转inference模型"></a>
+### 端到端模型转inference模型
+
+下载端到端模型：
+```
+wget -P ./ch_lite/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_train.tar && tar xf ./ch_lite/ch_ppocr_mobile_v2.0_cls_train.tar -C ./ch_lite/
+```
+
+端到端模型转inference模型与检测的方式相同，如下：
+```
+# -c 后面设置训练算法的yml配置文件
+# -o 配置可选参数
+# Global.pretrained_model 参数设置待转换的训练模型地址，不用添加文件后缀 .pdmodel，.pdopt或.pdparams。
+# Global.load_static_weights 参数需要设置为 False。
+# Global.save_inference_dir参数设置转换的模型将保存的地址。
+
+python3 tools/export_model.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.pretrained_model=./ch_lite/ch_ppocr_mobile_v2.0_cls_train/best_accuracy Global.load_static_weights=False Global.save_inference_dir=./inference/e2e/
+```
+
+转换成功后，在目录下有三个文件：
+```
+/inference/e2e/
+    ├── inference.pdiparams         # 分类inference模型的参数文件
+    ├── inference.pdiparams.info    # 分类inference模型的参数信息，可忽略
+    └── inference.pdmodel           # 分类inference模型的program文件
+```
 
 <a name="文本检测模型推理"></a>
 ## 二、文本检测模型推理
@@ -332,8 +362,38 @@ python3 tools/infer/predict_rec.py --image_dir="./doc/imgs_words/korean/1.jpg" -
 Predicts of ./doc/imgs_words/korean/1.jpg:('바탕으로', 0.9948904)
 ```
 
+<a name="端到端模型推理"></a>
+## 四、端到端模型推理
+
+端到端模型推理，默认使用PGNet模型的配置参数。当不使用PGNet模型时，在推理时，需要通过传入相应的参数进行算法适配，细节参考下文。
+<a name="PGNet端到端模型推理"></a>
+### 1. PGNet端到端模型推理
+#### (1). 四边形文本检测模型（ICDAR2015）  
+首先将PGNet端到端训练过程中保存的模型，转换成inference model。以基于Resnet50_vd骨干网络，在ICDAR2015英文数据集训练的模型为例([模型下载地址](https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/en_server_pgnetA.tar))，可以使用如下命令进行转换：
+```
+python3 tools/export_model.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.pretrained_model=./en_server_pgnetA/iter_epoch_450 Global.load_static_weights=False Global.save_inference_dir=./inference/e2e
+```
+**PGNet端到端模型推理，需要设置参数`--e2e_algorithm="PGNet"`**，可以执行如下命令：
+```
+python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/img_10.jpg" --e2e_model_dir="./inference/e2e/"  --e2e_pgnet_polygon=False
+```
+可视化文本检测结果默认保存到`./inference_results`文件夹里面，结果文件的名称前缀为'e2e_res'。结果示例如下：
+
+![](../imgs_results/e2e_res_img_10_pgnet.jpg)
+
+#### (2). 弯曲文本检测模型（Total-Text）  
+和四边形文本检测模型共用一个推理模型
+**PGNet端到端模型推理，需要设置参数`--e2e_algorithm="PGNet"`，同时，还需要增加参数`--e2e_pgnet_polygon=True`，**可以执行如下命令：
+```
+python3.7 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/img623.jpg" --e2e_model_dir="./inference/e2e/" --e2e_pgnet_polygon=True
+```
+可视化文本端到端结果默认保存到`./inference_results`文件夹里面，结果文件的名称前缀为'e2e_res'。结果示例如下：
+
+![](../imgs_results/e2e_res_img623_pgnet.jpg)
+
+
 <a name="方向分类模型推理"></a>
-## 四、方向分类模型推理
+## 五、方向分类模型推理
 
 下面将介绍方向分类模型推理。
 
@@ -358,7 +418,7 @@ Predicts of ./doc/imgs_words/ch/word_4.jpg:['0', 0.9999982]
 ```
 
 <a name="文本检测、方向分类和文字识别串联推理"></a>
-## 五、文本检测、方向分类和文字识别串联推理
+## 六、文本检测、方向分类和文字识别串联推理
 <a name="超轻量中文OCR模型推理"></a>
 ### 1. 超轻量中文OCR模型推理
 

doc/doc_ch/multi_languages.md

+# 多语言模型
+
+**近期更新**
+
+- 2021.4.9 支持**80种**语言的检测和识别
+- 2021.4.9 支持**轻量高精度**英文模型检测识别
+
+- [1 安装](#安装)
+    - [1.1 paddle 安装](#paddle安装)
+    - [1.2 paddleocr package 安装](#paddleocr_package_安装)  
+
+- [2 快速使用](#快速使用)
+    - [2.1 命令行运行](#命令行运行)
+     - [2.1.1 整图预测](#bash_检测+识别)
+     - [2.1.2 识别预测](#bash_识别)
+     - [2.1.3 检测预测](#bash_检测)
+    - [2.2 python 脚本运行](#python_脚本运行)
+     - [2.2.1 整图预测](#python_检测+识别)
+     - [2.2.2 识别预测](#python_识别)
+     - [2.2.3 检测预测](#python_检测)
+- [3 自定义训练](#自定义训练)
+- [4 支持语种及缩写](#语种缩写)
+
+<a name="安装"></a>
+## 1 安装
+
+<a name="paddle安装"></a>
+### 1.1 paddle 安装
+```
+# cpu
+pip install paddlepaddle
+
+# gpu
+pip instll paddlepaddle-gpu
+```
+
+<a name="paddleocr_package_安装"></a>
+### 1.2 paddleocr package 安装
+
+
+pip 安装
+```
+pip install "paddleocr>=2.0.4" # 推荐使用2.0.4版本
+```
+本地构建并安装
+```
+python3 setup.py bdist_wheel
+pip3 install dist/paddleocr-x.x.x-py3-none-any.whl # x.x.x是paddleocr的版本号
+```
+
+<a name="快速使用"></a>
+## 2 快速使用
+
+<a name="命令行运行"></a>
+### 2.1 命令行运行
+
+查看帮助信息
+
+```
+paddleocr -h
+```
+
+* 整图预测（检测+识别）
+
+Paddleocr目前支持80个语种，可以通过修改--lang参数进行切换，具体支持的[语种](#语种缩写)可查看表格。
+
+``` bash
+
+paddleocr --image_dir doc/imgs/japan_2.jpg --lang=japan
+```
+![](https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.0/doc/imgs/japan_2.jpg)
+
+结果是一个list，每个item包含了文本框，文字和识别置信度
+```text
+[[[671.0, 60.0], [847.0, 63.0], [847.0, 104.0], [671.0, 102.0]], ('もちもち', 0.9993342)]
+[[[394.0, 82.0], [536.0, 77.0], [538.0, 127.0], [396.0, 132.0]], ('天然の', 0.9919842)]
+[[[880.0, 89.0], [1014.0, 93.0], [1013.0, 127.0], [879.0, 124.0]], ('とろっと', 0.9976762)]
+[[[1067.0, 101.0], [1294.0, 101.0], [1294.0, 138.0], [1067.0, 138.0]], ('後味のよい', 0.9988712)]
+......
+```
+
+* 识别预测
+
+```bash
+paddleocr --image_dir doc/imgs_words/japan/1.jpg   --det false --lang=japan
+```
+
+![](https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.0/doc/imgs_words/japan/1.jpg)
+
+结果是一个tuple，返回识别结果和识别置信度
+
+```text
+('したがって', 0.99965394)
+```
+
+* 检测预测
+
+```
+paddleocr --image_dir PaddleOCR/doc/imgs/11.jpg --rec false
+```
+
+结果是一个list，每个item只包含文本框
+
+```
+[[26.0, 457.0], [137.0, 457.0], [137.0, 477.0], [26.0, 477.0]]
+[[25.0, 425.0], [372.0, 425.0], [372.0, 448.0], [25.0, 448.0]]
+[[128.0, 397.0], [273.0, 397.0], [273.0, 414.0], [128.0, 414.0]]
+......
+```
+
+<a name="python_脚本运行"></a>
+### 2.2 python 脚本运行
+
+ppocr 也支持在python脚本中运行，便于嵌入到您自己的代码中：
+
+* 整图预测（检测+识别）
+
+```
+from paddleocr import PaddleOCR, draw_ocr
+
+# 同样也是通过修改 lang 参数切换语种
+ocr = PaddleOCR(lang="korean") # 首次执行会自动下载模型文件
+img_path = 'doc/imgs/korean_1.jpg '
+result = ocr.ocr(img_path)
+# 打印检测框和识别结果
+for line in result:
+    print(line)
+
+# 可视化
+from PIL import Image
+image = Image.open(img_path).convert('RGB')
+boxes = [line[0] for line in result]
+txts = [line[1][0] for line in result]
+scores = [line[1][1] for line in result]
+im_show = draw_ocr(image, boxes, txts, scores, font_path='/path/to/PaddleOCR/doc/korean.ttf')
+im_show = Image.fromarray(im_show)
+im_show.save('result.jpg')
+```
+
+结果可视化:
+![](https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.0/doc/imgs_results/korean.jpg)
+
+
+* 识别预测
+
+```
+from paddleocr import PaddleOCR
+ocr = PaddleOCR(lang="german")
+img_path = 'PaddleOCR/doc/imgs_words/german/1.jpg'
+result = ocr.ocr(img_path, det=False, cls=True)
+for line in result:
+    print(line)
+```
+
+![](https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.0/doc/imgs_words/german/1.jpg)
+
+结果是一个tuple，只包含识别结果和识别置信度
+
+```
+('leider auch jetzt', 0.97538936)
+```
+
+* 检测预测
+
+```python
+from paddleocr import PaddleOCR, draw_ocr
+ocr = PaddleOCR() # need to run only once to download and load model into memory
+img_path = 'PaddleOCR/doc/imgs_en/img_12.jpg'
+result = ocr.ocr(img_path, rec=False)
+for line in result:
+    print(line)
+
+# 显示结果
+from PIL import Image
+
+image = Image.open(img_path).convert('RGB')
+im_show = draw_ocr(image, result, txts=None, scores=None, font_path='/path/to/PaddleOCR/doc/fonts/simfang.ttf')
+im_show = Image.fromarray(im_show)
+im_show.save('result.jpg')
+```
+结果是一个list，每个item只包含文本框
+```bash
+[[26.0, 457.0], [137.0, 457.0], [137.0, 477.0], [26.0, 477.0]]
+[[25.0, 425.0], [372.0, 425.0], [372.0, 448.0], [25.0, 448.0]]
+[[128.0, 397.0], [273.0, 397.0], [273.0, 414.0], [128.0, 414.0]]
+......
+```
+
+结果可视化 ：
+![](https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/release/2.0/doc/imgs_results/whl/12_det.jpg)
+
+ppocr 还支持方向分类， 更多使用方式请参考：[whl包使用说明](https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.0/doc/doc_ch/whl.md)。
+
+<a name="自定义训练"></a>
+## 3 自定义训练
+
+ppocr 支持使用自己的数据进行自定义训练或finetune, 其中识别模型可以参考 [法语配置文件](../../configs/rec/multi_language/rec_french_lite_train.yml)
+修改训练数据路径、字典等参数。
+
+具体数据准备、训练过程可参考：[文本检测](../doc_ch/detection.md)、[文本识别](../doc_ch/recognition.md)，更多功能如预测部署、
+数据标注等功能可以阅读完整的[文档教程](../../README_ch.md)。
+
+<a name="语种缩写"></a>
+## 4 支持语种及缩写
+
+| 语种 | 描述 | 缩写 |
+| --- | --- | --- |
+|中文|chinese and english|ch|
+|英文|english|en|
+|法文|french|fr|
+|德文|german|german|
+|日文|japan|japan|
+|韩文|korean|korean|
+|中文繁体|chinese traditional |ch_tra|
+|意大利文| Italian |it|
+|西班牙文|Spanish |es|
+|葡萄牙文| Portuguese|pt|
+|俄罗斯文|Russia|ru|
+|阿拉伯文|Arabic|ar|
+|印地文|Hindi|hi|
+|维吾尔|Uyghur|ug|
+|波斯文|Persian|fa|
+|乌尔都文|Urdu|ur|
+|塞尔维亚文（latin)| Serbian(latin) |rs_latin|
+|欧西坦文|Occitan |oc|
+|马拉地文|Marathi|mr|
+|尼泊尔文|Nepali|ne|
+|塞尔维亚文（cyrillic)|Serbian(cyrillic)|rs_cyrillic|
+|保加利亚文|Bulgarian |bg|
+|乌克兰文|Ukranian|uk|
+|白俄罗斯文|Belarusian|be|
+|泰卢固文|Telugu |te|
+|卡纳达文|Kannada |kn|
+|泰米尔文|Tamil |ta|
+|南非荷兰文 |Afrikaans |af|
+|阿塞拜疆文 |Azerbaijani	|az|
+|波斯尼亚文|Bosnian|bs|
+|捷克文|Czech|cs|
+|威尔士文 |Welsh |cy|
+|丹麦文 |Danish|da|
+|爱沙尼亚文 |Estonian |et|
+|爱尔兰文 |Irish |ga|
+|克罗地亚文|Croatian |hr|
+|匈牙利文|Hungarian |hu|
+|印尼文|Indonesian|id|
+|冰岛文 |Icelandic|is|
+|库尔德文 |Kurdish|ku|
+|立陶宛文|Lithuanian |lt|
+|拉脱维亚文 |Latvian |lv|
+|毛利文|Maori|mi|
+|马来文 |Malay|ms|
+|马耳他文 |Maltese |mt|
+|荷兰文 |Dutch |nl|
+|挪威文 |Norwegian |no|
+|波兰文|Polish |pl|
+| 罗马尼亚文|Romanian |ro|
+| 斯洛伐克文|Slovak |sk|
+| 斯洛文尼亚文|Slovenian |sl|
+| 阿尔巴尼亚文|Albanian |sq|
+| 瑞典文|Swedish |sv|
+| 西瓦希里文|Swahili |sw|
+| 塔加洛文|Tagalog |tl|
+| 土耳其文|Turkish |tr|
+| 乌兹别克文|Uzbek |uz|
+| 越南文|Vietnamese |vi|
+| 蒙古文|Mongolian |mn|
+| 阿巴扎文|Abaza |abq|
+| 阿迪赫文|Adyghe |ady|
+| 卡巴丹文|Kabardian |kbd|
+| 阿瓦尔文|Avar |ava|
+| 达尔瓦文|Dargwa |dar|
+| 因古什文|Ingush |inh|
+| 拉克文|Lak |lbe|
+| 莱兹甘文|Lezghian |lez|
+|塔巴萨兰文 |Tabassaran |tab|
+| 比尔哈文|Bihari |bh|
+| 迈蒂利文|Maithili |mai|
+| 昂加文|Angika |ang|
+| 孟加拉文|Bhojpuri |bho|
+| 摩揭陀文 |Magahi |mah|
+| 那格浦尔文|Nagpur |sck|
+| 尼瓦尔文|Newari |new|
+| 保加利亚文 |Goan Konkani|gom|
+| 沙特阿拉伯文|Saudi Arabia|sa|

doc/doc_ch/pgnet.md

-
-# 端对端OCR算法-PGNet
-
-----
 # 端对端OCR算法-PGNet
 - [一、简介](#简介)
 - [二、环境配置](#环境配置)
 - [三、快速使用](#快速使用)
 - [四、模型训练、评估、推理](#快速训练)
 
-
 <a name="简介"></a>
 ## 一、简介
 OCR算法可以分为两阶段算法和端对端的算法。二阶段OCR算法一般分为两个部分，文本检测和文本识别算法，文件检测算法从图像中得到文本行的检测框，然后识别算法去识别文本框中的内容。而端对端OCR算法可以在一个算法中完成文字检测和文字识别，其基本思想是设计一个同时具有检测单元和识别模块的模型，共享其中两者的CNN特征，并联合训练。由于一个算法即可完成文字识别，端对端模型更小，速度更快。
@@ -62,13 +57,12 @@ python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/im
 # 如果想使用CPU进行预测，需设置use_gpu参数为False
 python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/img623.jpg" --e2e_model_dir="./inference/e2e/" --e2e_pgnet_polygon=True --use_gpu=False
 ```
-
 ### 可视化结果
 可视化文本检测结果默认保存到./inference_results文件夹里面，结果文件的名称前缀为'e2e_res'。结果示例如下：
 ![](../imgs_results/e2e_res_img623_pgnet.jpg)
 
 <a name="快速训练"></a>
-## 四、快速训练
+## 四、模型训练、评估、推理
 本节以totaltext数据集为例，介绍PaddleOCR中端到端模型的训练、评估与测试。
 
 ###  准备数据
@@ -103,7 +97,6 @@ wget -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/tr
   └─ best_accuracy.states
   └─ best_accuracy.pdparams
 ```
-
 *如果您安装的是cpu版本，请将配置文件中的 `use_gpu` 字段修改为false*
 
 ```shell
@@ -121,8 +114,13 @@ python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/
 python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Optimizer.base_lr=0.0001
 ```
 
+#### 断点训练
+如果训练程序中断，如果希望加载训练中断的模型从而恢复训练，可以通过指定Global.checkpoints指定要加载的模型路径：
+```shell
+python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.checkpoints=./your/trained/model
+```
 
-### 模型评估
+**注意**：`Global.checkpoints`的优先级高于`Global.pretrain_weights`的优先级，即同时指定两个参数时，优先加载`Global.checkpoints`指定的模型，如果`Global.checkpoints`指定的模型路径有误，会加载`Global.pretrain_weights`指定的模型。
 
 PaddleOCR计算三个OCR端到端相关的指标，分别是：Precision、Recall、Hmean。
 

paddleocr.py

@@ -66,6 +66,46 @@ model_urls = {
             'url':
             'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/japan_mobile_v2.0_rec_infer.tar',
             'dict_path': './ppocr/utils/dict/japan_dict.txt'
+        },
+        'chinese_cht': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/chinese_cht_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/chinese_cht_dict.txt'
+        },
+        'ta': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ta_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/ta_dict.txt'
+        },
+        'te': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/te_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/te_dict.txt'
+        },
+        'ka': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/ka_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/ka_dict.txt'
+        },
+        'latin': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/latin_ppocr_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/latin_dict.txt'
+        },
+        'arabic': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/arabic_ppocr_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/arabic_dict.txt'
+        },
+        'cyrillic': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/cyrillic_ppocr_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/cyrillic_dict.txt'
+        },
+        'devanagari': {
+            'url':
+            'https://paddleocr.bj.bcebos.com/dygraph_v2.0/multilingual/devanagari_ppocr_mobile_v2.0_rec_infer.tar',
+            'dict_path': './ppocr/utils/dict/devanagari_dict.txt'
         }
     },
     'cls':
@@ -233,6 +273,29 @@ class PaddleOCR(predict_system.TextSystem):
         postprocess_params.__dict__.update(**kwargs)
         self.use_angle_cls = postprocess_params.use_angle_cls
         lang = postprocess_params.lang
+        latin_lang = [
+            'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'en', 'es', 'et', 'fr',
+            'ga', 'hr', 'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi',
+            'ms', 'mt', 'nl', 'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'rs_latin',
+            'sk', 'sl', 'sq', 'sv', 'sw', 'tl', 'tr', 'uz', 'vi'
+        ]
+        arabic_lang = ['ar', 'fa', 'ug', 'ur']
+        cyrillic_lang = [
+            'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd',
+            'ava', 'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
+        ]
+        devanagari_lang = [
+            'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new',
+            'gom', 'sa', 'bgc'
+        ]
+        if lang in latin_lang:
+            lang = "latin"
+        elif lang in arabic_lang:
+            lang = "arabic"
+        elif lang in cyrillic_lang:
+            lang = "cyrillic"
+        elif lang in devanagari_lang:
+            lang = "devanagari"
         assert lang in model_urls[
             'rec'], 'param lang must in {}, but got {}'.format(
                 model_urls['rec'].keys(), lang)

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 """
 

ppocr/data/imaug/operators.py

@@ -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/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/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/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

ppocr/metrics/eval_det_iou.py

@@ -200,7 +200,8 @@ class DetectionIoUEvaluator(object):
         methodPrecision = 0 if numGlobalCareDet == 0 else float(
             matchedSum) / numGlobalCareDet
         methodHmean = 0 if methodRecall + methodPrecision == 0 else 2 * \
-            methodRecall * methodPrecision / (methodRecall + methodPrecision)
+                                                                    methodRecall * methodPrecision / (
+                                                                            methodRecall + methodPrecision)
         # print(methodRecall, methodPrecision, methodHmean)
         # sys.exit(-1)
         methodMetrics = {

ppocr/modeling/backbones/__init__.py

@@ -26,6 +26,9 @@ def build_backbone(config, model_type):
         from .rec_resnet_vd import ResNet
         from .rec_resnet_fpn import ResNetFPN
         support_dict = ['MobileNetV3', 'ResNet', 'ResNetFPN']
+    elif model_type == 'e2e':
+        from .e2e_resnet_vd_pg import ResNet
+        support_dict = ['ResNet']
     else:
         raise NotImplementedError
 

ppocr/modeling/backbones/e2e_resnet_vd_pg.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
+
+import paddle
+from paddle import ParamAttr
+import paddle.nn as nn
+import paddle.nn.functional as F
+
+__all__ = ["ResNet"]
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=1,
+            groups=1,
+            is_vd_mode=False,
+            act=None,
+            name=None, ):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True)
+        self._conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False)
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = nn.BatchNorm(
+            out_channels,
+            act=act,
+            param_attr=ParamAttr(name=bn_name + '_scale'),
+            bias_attr=ParamAttr(bn_name + '_offset'),
+            moving_mean_name=bn_name + '_mean',
+            moving_variance_name=bn_name + '_variance')
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class BottleneckBlock(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 shortcut=True,
+                 if_first=False,
+                 name=None):
+        super(BottleneckBlock, self).__init__()
+
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            act='relu',
+            name=name + "_branch2a")
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act='relu',
+            name=name + "_branch2b")
+        self.conv2 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels * 4,
+            kernel_size=1,
+            act=None,
+            name=name + "_branch2c")
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels * 4,
+                kernel_size=1,
+                stride=stride,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1")
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+        conv2 = self.conv2(conv1)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv2)
+        y = F.relu(y)
+        return y
+
+
+class BasicBlock(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 stride,
+                 shortcut=True,
+                 if_first=False,
+                 name=None):
+        super(BasicBlock, self).__init__()
+        self.stride = stride
+        self.conv0 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            stride=stride,
+            act='relu',
+            name=name + "_branch2a")
+        self.conv1 = ConvBNLayer(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            act=None,
+            name=name + "_branch2b")
+
+        if not shortcut:
+            self.short = ConvBNLayer(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                is_vd_mode=False if if_first else True,
+                name=name + "_branch1")
+
+        self.shortcut = shortcut
+
+    def forward(self, inputs):
+        y = self.conv0(inputs)
+        conv1 = self.conv1(y)
+
+        if self.shortcut:
+            short = inputs
+        else:
+            short = self.short(inputs)
+        y = paddle.add(x=short, y=conv1)
+        y = F.relu(y)
+        return y
+
+
+class ResNet(nn.Layer):
+    def __init__(self, in_channels=3, layers=50, **kwargs):
+        super(ResNet, self).__init__()
+
+        self.layers = layers
+        supported_layers = [18, 34, 50, 101, 152, 200]
+        assert layers in supported_layers, \
+            "supported layers are {} but input layer is {}".format(
+                supported_layers, layers)
+
+        if layers == 18:
+            depth = [2, 2, 2, 2]
+        elif layers == 34 or layers == 50:
+            # depth = [3, 4, 6, 3]
+            depth = [3, 4, 6, 3, 3]
+        elif layers == 101:
+            depth = [3, 4, 23, 3]
+        elif layers == 152:
+            depth = [3, 8, 36, 3]
+        elif layers == 200:
+            depth = [3, 12, 48, 3]
+        num_channels = [64, 256, 512, 1024,
+                        2048] if layers >= 50 else [64, 64, 128, 256]
+        num_filters = [64, 128, 256, 512, 512]
+
+        self.conv1_1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=7,
+            stride=2,
+            act='relu',
+            name="conv1_1")
+        self.pool2d_max = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
+
+        self.stages = []
+        self.out_channels = [3, 64]
+        # num_filters = [64, 128, 256, 512, 512]
+        if layers >= 50:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    if layers in [101, 152] and block == 2:
+                        if i == 0:
+                            conv_name = "res" + str(block + 2) + "a"
+                        else:
+                            conv_name = "res" + str(block + 2) + "b" + str(i)
+                    else:
+                        conv_name = "res" + str(block + 2) + chr(97 + i)
+                    bottleneck_block = self.add_sublayer(
+                        'bb_%d_%d' % (block, i),
+                        BottleneckBlock(
+                            in_channels=num_channels[block]
+                            if i == 0 else num_filters[block] * 4,
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name))
+                    shortcut = True
+                    block_list.append(bottleneck_block)
+                self.out_channels.append(num_filters[block] * 4)
+                self.stages.append(nn.Sequential(*block_list))
+        else:
+            for block in range(len(depth)):
+                block_list = []
+                shortcut = False
+                for i in range(depth[block]):
+                    conv_name = "res" + str(block + 2) + chr(97 + i)
+                    basic_block = self.add_sublayer(
+                        'bb_%d_%d' % (block, i),
+                        BasicBlock(
+                            in_channels=num_channels[block]
+                            if i == 0 else num_filters[block],
+                            out_channels=num_filters[block],
+                            stride=2 if i == 0 and block != 0 else 1,
+                            shortcut=shortcut,
+                            if_first=block == i == 0,
+                            name=conv_name))
+                    shortcut = True
+                    block_list.append(basic_block)
+                self.out_channels.append(num_filters[block])
+                self.stages.append(nn.Sequential(*block_list))
+
+    def forward(self, inputs):
+        out = [inputs]
+        y = self.conv1_1(inputs)
+        out.append(y)
+        y = self.pool2d_max(y)
+        for block in self.stages:
+            y = block(y)
+            out.append(y)
+        return out

ppocr/modeling/heads/__init__.py

@@ -20,6 +20,7 @@ def build_head(config):
     from .det_db_head import DBHead
     from .det_east_head import EASTHead
     from .det_sast_head import SASTHead
+    from .e2e_pg_head import PGHead
 
     # rec head
     from .rec_ctc_head import CTCHead
@@ -30,8 +31,8 @@ def build_head(config):
     from .cls_head import ClsHead
     support_dict = [
         'DBHead', 'EASTHead', 'SASTHead', 'CTCHead', 'ClsHead', 'AttentionHead',
-        'SRNHead'
-    ]
+        'SRNHead', 'PGHead']
+
 
     module_name = config.pop('name')
     assert module_name in support_dict, Exception('head only support {}'.format(

ppocr/modeling/heads/e2e_pg_head.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
+
+import math
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 groups=1,
+                 if_act=True,
+                 act=None,
+                 name=None):
+        super(ConvBNLayer, self).__init__()
+        self.if_act = if_act
+        self.act = act
+        self.conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + '_weights'),
+            bias_attr=False)
+
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+            use_global_stats=False)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return x
+
+
+class PGHead(nn.Layer):
+    """
+    """
+
+    def __init__(self, in_channels, **kwargs):
+        super(PGHead, self).__init__()
+        self.conv_f_score1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_score{}".format(1))
+        self.conv_f_score2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act='relu',
+            name="conv_f_score{}".format(2))
+        self.conv_f_score3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_score{}".format(3))
+
+        self.conv1 = nn.Conv2D(
+            in_channels=128,
+            out_channels=1,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_score{}".format(4)),
+            bias_attr=False)
+
+        self.conv_f_boder1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_boder{}".format(1))
+        self.conv_f_boder2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act='relu',
+            name="conv_f_boder{}".format(2))
+        self.conv_f_boder3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_boder{}".format(3))
+        self.conv2 = nn.Conv2D(
+            in_channels=128,
+            out_channels=4,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_boder{}".format(4)),
+            bias_attr=False)
+        self.conv_f_char1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_char{}".format(1))
+        self.conv_f_char2 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act='relu',
+            name="conv_f_char{}".format(2))
+        self.conv_f_char3 = ConvBNLayer(
+            in_channels=128,
+            out_channels=256,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_char{}".format(3))
+        self.conv_f_char4 = ConvBNLayer(
+            in_channels=256,
+            out_channels=256,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act='relu',
+            name="conv_f_char{}".format(4))
+        self.conv_f_char5 = ConvBNLayer(
+            in_channels=256,
+            out_channels=256,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_char{}".format(5))
+        self.conv3 = nn.Conv2D(
+            in_channels=256,
+            out_channels=37,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_char{}".format(6)),
+            bias_attr=False)
+
+        self.conv_f_direc1 = ConvBNLayer(
+            in_channels=in_channels,
+            out_channels=64,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_direc{}".format(1))
+        self.conv_f_direc2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            act='relu',
+            name="conv_f_direc{}".format(2))
+        self.conv_f_direc3 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            act='relu',
+            name="conv_f_direc{}".format(3))
+        self.conv4 = nn.Conv2D(
+            in_channels=128,
+            out_channels=2,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=1,
+            weight_attr=ParamAttr(name="conv_f_direc{}".format(4)),
+            bias_attr=False)
+
+    def forward(self, x):
+        f_score = self.conv_f_score1(x)
+        f_score = self.conv_f_score2(f_score)
+        f_score = self.conv_f_score3(f_score)
+        f_score = self.conv1(f_score)
+        f_score = F.sigmoid(f_score)
+
+        # f_border
+        f_border = self.conv_f_boder1(x)
+        f_border = self.conv_f_boder2(f_border)
+        f_border = self.conv_f_boder3(f_border)
+        f_border = self.conv2(f_border)
+
+        f_char = self.conv_f_char1(x)
+        f_char = self.conv_f_char2(f_char)
+        f_char = self.conv_f_char3(f_char)
+        f_char = self.conv_f_char4(f_char)
+        f_char = self.conv_f_char5(f_char)
+        f_char = self.conv3(f_char)
+
+        f_direction = self.conv_f_direc1(x)
+        f_direction = self.conv_f_direc2(f_direction)
+        f_direction = self.conv_f_direc3(f_direction)
+        f_direction = self.conv4(f_direction)
+
+        predicts = {}
+        predicts['f_score'] = f_score
+        predicts['f_border'] = f_border
+        predicts['f_char'] = f_char
+        predicts['f_direction'] = f_direction
+        return predicts

ppocr/modeling/necks/__init__.py

@@ -14,12 +14,14 @@
 
 __all__ = ['build_neck']
 
+
 def build_neck(config):
     from .db_fpn import DBFPN
     from .east_fpn import EASTFPN
     from .sast_fpn import SASTFPN
     from .rnn import SequenceEncoder
-    support_dict = ['DBFPN', 'EASTFPN', 'SASTFPN', 'SequenceEncoder']
+    from .pg_fpn import PGFPN
+    support_dict = ['DBFPN', 'EASTFPN', 'SASTFPN', 'SequenceEncoder', 'PGFPN']
 
     module_name = config.pop('name')
     assert module_name in support_dict, Exception('neck only support {}'.format(

ppocr/modeling/necks/pg_fpn.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
+
+import paddle
+from paddle import nn
+import paddle.nn.functional as F
+from paddle import ParamAttr
+
+
+class ConvBNLayer(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 groups=1,
+                 is_vd_mode=False,
+                 act=None,
+                 name=None):
+        super(ConvBNLayer, self).__init__()
+
+        self.is_vd_mode = is_vd_mode
+        self._pool2d_avg = nn.AvgPool2D(
+            kernel_size=2, stride=2, padding=0, ceil_mode=True)
+        self._conv = nn.Conv2D(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + "_weights"),
+            bias_attr=False)
+        if name == "conv1":
+            bn_name = "bn_" + name
+        else:
+            bn_name = "bn" + name[3:]
+        self._batch_norm = nn.BatchNorm(
+            out_channels,
+            act=act,
+            param_attr=ParamAttr(name=bn_name + '_scale'),
+            bias_attr=ParamAttr(bn_name + '_offset'),
+            moving_mean_name=bn_name + '_mean',
+            moving_variance_name=bn_name + '_variance',
+            use_global_stats=False)
+
+    def forward(self, inputs):
+        y = self._conv(inputs)
+        y = self._batch_norm(y)
+        return y
+
+
+class DeConvBNLayer(nn.Layer):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=4,
+                 stride=2,
+                 padding=1,
+                 groups=1,
+                 if_act=True,
+                 act=None,
+                 name=None):
+        super(DeConvBNLayer, self).__init__()
+
+        self.if_act = if_act
+        self.act = act
+        self.deconv = nn.Conv2DTranspose(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            groups=groups,
+            weight_attr=ParamAttr(name=name + '_weights'),
+            bias_attr=False)
+        self.bn = nn.BatchNorm(
+            num_channels=out_channels,
+            act=act,
+            param_attr=ParamAttr(name="bn_" + name + "_scale"),
+            bias_attr=ParamAttr(name="bn_" + name + "_offset"),
+            moving_mean_name="bn_" + name + "_mean",
+            moving_variance_name="bn_" + name + "_variance",
+            use_global_stats=False)
+
+    def forward(self, x):
+        x = self.deconv(x)
+        x = self.bn(x)
+        return x
+
+
+class PGFPN(nn.Layer):
+    def __init__(self, in_channels, **kwargs):
+        super(PGFPN, self).__init__()
+        num_inputs = [2048, 2048, 1024, 512, 256]
+        num_outputs = [256, 256, 192, 192, 128]
+        self.out_channels = 128
+        self.conv_bn_layer_1 = ConvBNLayer(
+            in_channels=3,
+            out_channels=32,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name='FPN_d1')
+        self.conv_bn_layer_2 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name='FPN_d2')
+        self.conv_bn_layer_3 = ConvBNLayer(
+            in_channels=256,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            act=None,
+            name='FPN_d3')
+        self.conv_bn_layer_4 = ConvBNLayer(
+            in_channels=32,
+            out_channels=64,
+            kernel_size=3,
+            stride=2,
+            act=None,
+            name='FPN_d4')
+        self.conv_bn_layer_5 = ConvBNLayer(
+            in_channels=64,
+            out_channels=64,
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name='FPN_d5')
+        self.conv_bn_layer_6 = ConvBNLayer(
+            in_channels=64,
+            out_channels=128,
+            kernel_size=3,
+            stride=2,
+            act=None,
+            name='FPN_d6')
+        self.conv_bn_layer_7 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name='FPN_d7')
+        self.conv_bn_layer_8 = ConvBNLayer(
+            in_channels=128,
+            out_channels=128,
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name='FPN_d8')
+
+        self.conv_h0 = ConvBNLayer(
+            in_channels=num_inputs[0],
+            out_channels=num_outputs[0],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(0))
+        self.conv_h1 = ConvBNLayer(
+            in_channels=num_inputs[1],
+            out_channels=num_outputs[1],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(1))
+        self.conv_h2 = ConvBNLayer(
+            in_channels=num_inputs[2],
+            out_channels=num_outputs[2],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(2))
+        self.conv_h3 = ConvBNLayer(
+            in_channels=num_inputs[3],
+            out_channels=num_outputs[3],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(3))
+        self.conv_h4 = ConvBNLayer(
+            in_channels=num_inputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_h{}".format(4))
+
+        self.dconv0 = DeConvBNLayer(
+            in_channels=num_outputs[0],
+            out_channels=num_outputs[0 + 1],
+            name="dconv_{}".format(0))
+        self.dconv1 = DeConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[1 + 1],
+            act=None,
+            name="dconv_{}".format(1))
+        self.dconv2 = DeConvBNLayer(
+            in_channels=num_outputs[2],
+            out_channels=num_outputs[2 + 1],
+            act=None,
+            name="dconv_{}".format(2))
+        self.dconv3 = DeConvBNLayer(
+            in_channels=num_outputs[3],
+            out_channels=num_outputs[3 + 1],
+            act=None,
+            name="dconv_{}".format(3))
+        self.conv_g1 = ConvBNLayer(
+            in_channels=num_outputs[1],
+            out_channels=num_outputs[1],
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name="conv_g{}".format(1))
+        self.conv_g2 = ConvBNLayer(
+            in_channels=num_outputs[2],
+            out_channels=num_outputs[2],
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name="conv_g{}".format(2))
+        self.conv_g3 = ConvBNLayer(
+            in_channels=num_outputs[3],
+            out_channels=num_outputs[3],
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name="conv_g{}".format(3))
+        self.conv_g4 = ConvBNLayer(
+            in_channels=num_outputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=3,
+            stride=1,
+            act='relu',
+            name="conv_g{}".format(4))
+        self.convf = ConvBNLayer(
+            in_channels=num_outputs[4],
+            out_channels=num_outputs[4],
+            kernel_size=1,
+            stride=1,
+            act=None,
+            name="conv_f{}".format(4))
+
+    def forward(self, x):
+        c0, c1, c2, c3, c4, c5, c6 = x
+        # FPN_Down_Fusion
+        f = [c0, c1, c2]
+        g = [None, None, None]
+        h = [None, None, None]
+        h[0] = self.conv_bn_layer_1(f[0])
+        h[1] = self.conv_bn_layer_2(f[1])
+        h[2] = self.conv_bn_layer_3(f[2])
+
+        g[0] = self.conv_bn_layer_4(h[0])
+        g[1] = paddle.add(g[0], h[1])
+        g[1] = F.relu(g[1])
+        g[1] = self.conv_bn_layer_5(g[1])
+        g[1] = self.conv_bn_layer_6(g[1])
+
+        g[2] = paddle.add(g[1], h[2])
+        g[2] = F.relu(g[2])
+        g[2] = self.conv_bn_layer_7(g[2])
+        f_down = self.conv_bn_layer_8(g[2])
+
+        # FPN UP Fusion
+        f1 = [c6, c5, c4, c3, c2]
+        g = [None, None, None, None, None]
+        h = [None, None, None, None, None]
+        h[0] = self.conv_h0(f1[0])
+        h[1] = self.conv_h1(f1[1])
+        h[2] = self.conv_h2(f1[2])
+        h[3] = self.conv_h3(f1[3])
+        h[4] = self.conv_h4(f1[4])
+
+        g[0] = self.dconv0(h[0])
+        g[1] = paddle.add(g[0], h[1])
+        g[1] = F.relu(g[1])
+        g[1] = self.conv_g1(g[1])
+        g[1] = self.dconv1(g[1])
+
+        g[2] = paddle.add(g[1], h[2])
+        g[2] = F.relu(g[2])
+        g[2] = self.conv_g2(g[2])
+        g[2] = self.dconv2(g[2])
+
+        g[3] = paddle.add(g[2], h[3])
+        g[3] = F.relu(g[3])
+        g[3] = self.conv_g3(g[3])
+        g[3] = self.dconv3(g[3])
+
+        g[4] = paddle.add(x=g[3], y=h[4])
+        g[4] = F.relu(g[4])
+        g[4] = self.conv_g4(g[4])
+        f_up = self.convf(g[4])
+        f_common = paddle.add(f_down, f_up)
+        f_common = F.relu(f_common)
+        return f_common

ppocr/postprocess/__init__.py

@@ -28,10 +28,11 @@ def build_post_process(config, global_config=None):
     from .sast_postprocess import SASTPostProcess
     from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode
     from .cls_postprocess import ClsPostProcess
+    from .pg_postprocess import PGPostProcess
 
     support_dict = [
         'DBPostProcess', 'EASTPostProcess', 'SASTPostProcess', 'CTCLabelDecode',
-        'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode'
+        'AttnLabelDecode', 'ClsPostProcess', 'SRNLabelDecode', 'PGPostProcess'
     ]
 
     config = copy.deepcopy(config)

ppocr/postprocess/pg_postprocess.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
+
+import os
+import sys
+
+__dir__ = os.path.dirname(__file__)
+sys.path.append(__dir__)
+sys.path.append(os.path.join(__dir__, '..'))
+
+from ppocr.utils.e2e_utils.extract_textpoint import *
+from ppocr.utils.e2e_utils.visual import *
+import paddle
+
+
+class PGPostProcess(object):
+    """
+    The post process for PGNet.
+    """
+
+    def __init__(self, character_dict_path, valid_set, score_thresh, **kwargs):
+        self.Lexicon_Table = get_dict(character_dict_path)
+        self.valid_set = valid_set
+        self.score_thresh = score_thresh
+
+        # c++ la-nms is faster, but only support python 3.5
+        self.is_python35 = False
+        if sys.version_info.major == 3 and sys.version_info.minor == 5:
+            self.is_python35 = True
+
+    def __call__(self, outs_dict, shape_list):
+        p_score = outs_dict['f_score']
+        p_border = outs_dict['f_border']
+        p_char = outs_dict['f_char']
+        p_direction = outs_dict['f_direction']
+        if isinstance(p_score, paddle.Tensor):
+            p_score = p_score[0].numpy()
+            p_border = p_border[0].numpy()
+            p_direction = p_direction[0].numpy()
+            p_char = p_char[0].numpy()
+        else:
+            p_score = p_score[0]
+            p_border = p_border[0]
+            p_direction = p_direction[0]
+            p_char = p_char[0]
+        src_h, src_w, ratio_h, ratio_w = shape_list[0]
+        is_curved = self.valid_set == "totaltext"
+        instance_yxs_list = generate_pivot_list(
+            p_score,
+            p_char,
+            p_direction,
+            score_thresh=self.score_thresh,
+            is_backbone=True,
+            is_curved=is_curved)
+        p_char = paddle.to_tensor(np.expand_dims(p_char, axis=0))
+        char_seq_idx_set = []
+        for i in range(len(instance_yxs_list)):
+            gather_info_lod = paddle.to_tensor(instance_yxs_list[i])
+            f_char_map = paddle.transpose(p_char, [0, 2, 3, 1])
+            feature_seq = paddle.gather_nd(f_char_map, gather_info_lod)
+            feature_seq = np.expand_dims(feature_seq.numpy(), axis=0)
+            feature_len = [len(feature_seq[0])]
+            featyre_seq = paddle.to_tensor(feature_seq)
+            feature_len = np.array([feature_len]).astype(np.int64)
+            length = paddle.to_tensor(feature_len)
+            seq_pred = paddle.fluid.layers.ctc_greedy_decoder(
+                input=featyre_seq, blank=36, input_length=length)
+            seq_pred_str = seq_pred[0].numpy().tolist()[0]
+            seq_len = seq_pred[1].numpy()[0][0]
+            temp_t = []
+            for c in seq_pred_str[:seq_len]:
+                temp_t.append(c)
+            char_seq_idx_set.append(temp_t)
+        seq_strs = []
+        for char_idx_set in char_seq_idx_set:
+            pr_str = ''.join([self.Lexicon_Table[pos] for pos in char_idx_set])
+            seq_strs.append(pr_str)
+        poly_list = []
+        keep_str_list = []
+        all_point_list = []
+        all_point_pair_list = []
+        for yx_center_line, keep_str in zip(instance_yxs_list, seq_strs):
+            if len(yx_center_line) == 1:
+                yx_center_line.append(yx_center_line[-1])
+
+            offset_expand = 1.0
+            if self.valid_set == 'totaltext':
+                offset_expand = 1.2
+
+            point_pair_list = []
+            for batch_id, y, x in yx_center_line:
+                offset = p_border[:, y, x].reshape(2, 2)
+                if offset_expand != 1.0:
+                    offset_length = np.linalg.norm(
+                        offset, axis=1, keepdims=True)
+                    expand_length = np.clip(
+                        offset_length * (offset_expand - 1),
+                        a_min=0.5,
+                        a_max=3.0)
+                    offset_detal = offset / offset_length * expand_length
+                    offset = offset + offset_detal
+                ori_yx = np.array([y, x], dtype=np.float32)
+                point_pair = (ori_yx + offset)[:, ::-1] * 4.0 / np.array(
+                    [ratio_w, ratio_h]).reshape(-1, 2)
+                point_pair_list.append(point_pair)
+
+                all_point_list.append([
+                    int(round(x * 4.0 / ratio_w)),
+                    int(round(y * 4.0 / ratio_h))
+                ])
+                all_point_pair_list.append(point_pair.round().astype(np.int32)
+                                           .tolist())
+
+            detected_poly, pair_length_info = point_pair2poly(point_pair_list)
+            detected_poly = expand_poly_along_width(
+                detected_poly, shrink_ratio_of_width=0.2)
+            detected_poly[:, 0] = np.clip(
+                detected_poly[:, 0], a_min=0, a_max=src_w)
+            detected_poly[:, 1] = np.clip(
+                detected_poly[:, 1], a_min=0, a_max=src_h)
+
+            if len(keep_str) < 2:
+                continue
+
+            keep_str_list.append(keep_str)
+            if self.valid_set == 'partvgg':
+                middle_point = len(detected_poly) // 2
+                detected_poly = detected_poly[
+                    [0, middle_point - 1, middle_point, -1], :]
+                poly_list.append(detected_poly)
+            elif self.valid_set == 'totaltext':
+                poly_list.append(detected_poly)
+            else:
+                print('--> Not supported format.')
+                exit(-1)
+        data = {
+            'points': poly_list,
+            'strs': keep_str_list,
+        }
+        return data

ppocr/postprocess/sast_postprocess.py

@@ -18,6 +18,7 @@ from __future__ import print_function
 
 import os
 import sys
+
 __dir__ = os.path.dirname(__file__)
 sys.path.append(__dir__)
 sys.path.append(os.path.join(__dir__, '..'))
@@ -67,11 +68,15 @@ class SASTPostProcess(object):
             point_list[point_num - 1 - idx] = point_pair[1]
         return np.array(point_list).reshape(-1, 2)
 
-    def shrink_quad_along_width(self, quad, begin_width_ratio=0., end_width_ratio=1.):
+    def shrink_quad_along_width(self,
+                                quad,
+                                begin_width_ratio=0.,
+                                end_width_ratio=1.):
         """ 
         Generate shrink_quad_along_width.
         """
-        ratio_pair = np.array([[begin_width_ratio], [end_width_ratio]], dtype=np.float32)
+        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]])
@@ -81,16 +86,23 @@ class SASTPostProcess(object):
         expand poly along width.
         """
         point_num = poly.shape[0]
-        left_quad = np.array([poly[0], poly[1], poly[-2], poly[-1]], dtype=np.float32)
+        left_quad = np.array(
+            [poly[0], poly[1], poly[-2], poly[-1]], dtype=np.float32)
         left_ratio = -shrink_ratio_of_width * np.linalg.norm(left_quad[0] - left_quad[3]) / \
                      (np.linalg.norm(left_quad[0] - left_quad[1]) + 1e-6)
-        left_quad_expand = self.shrink_quad_along_width(left_quad, left_ratio, 1.0)
-        right_quad = np.array([poly[point_num // 2 - 2], poly[point_num // 2 - 1],
-                            poly[point_num // 2], poly[point_num // 2 + 1]], dtype=np.float32)
+        left_quad_expand = self.shrink_quad_along_width(left_quad, left_ratio,
+                                                        1.0)
+        right_quad = np.array(
+            [
+                poly[point_num // 2 - 2], poly[point_num // 2 - 1],
+                poly[point_num // 2], poly[point_num // 2 + 1]
+            ],
+            dtype=np.float32)
         right_ratio = 1.0 + \
                       shrink_ratio_of_width * np.linalg.norm(right_quad[0] - right_quad[3]) / \
                       (np.linalg.norm(right_quad[0] - right_quad[1]) + 1e-6)
-        right_quad_expand = self.shrink_quad_along_width(right_quad, 0.0, right_ratio)
+        right_quad_expand = self.shrink_quad_along_width(right_quad, 0.0,
+                                                         right_ratio)
         poly[0] = left_quad_expand[0]
         poly[-1] = left_quad_expand[-1]
         poly[point_num // 2 - 1] = right_quad_expand[1]
@@ -121,12 +133,10 @@ class SASTPostProcess(object):
         """
         compute area of a quad.
         """
-        edge = [
-            (quad[1][0] - quad[0][0]) * (quad[1][1] + quad[0][1]),
+        edge = [(quad[1][0] - quad[0][0]) * (quad[1][1] + quad[0][1]),
                 (quad[2][0] - quad[1][0]) * (quad[2][1] + quad[1][1]),
                 (quad[3][0] - quad[2][0]) * (quad[3][1] + quad[2][1]),
-            (quad[0][0] - quad[3][0]) * (quad[0][1] + quad[3][1])
-        ]
+                (quad[0][0] - quad[3][0]) * (quad[0][1] + quad[3][1])]
         return np.sum(edge) / 2.
 
     def nms(self, dets):
@@ -157,7 +167,8 @@ class SASTPostProcess(object):
         m = quads.shape[0]
         gt_tc = np.mean(quads, axis=1)  # (m, 2)
 
-        pred_tc_tile = np.tile(pred_tc[:, np.newaxis, :], (1, m, 1)) # (n, m, 2)
+        pred_tc_tile = np.tile(pred_tc[:, np.newaxis, :],
+                               (1, m, 1))  # (n, m, 2)
         gt_tc_tile = np.tile(gt_tc[np.newaxis, :, :], (n, 1, 1))  # (n, m, 2)
         dist_mat = np.linalg.norm(pred_tc_tile - gt_tc_tile, axis=2)  # (n, m)
         xy_text_assign = np.argmin(dist_mat, axis=1) + 1  # (n,)
@@ -169,26 +180,47 @@ class SASTPostProcess(object):
         """
         Estimate sample points number.
         """
-        eh = (np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[1] - quad[2])) / 2.0
-        ew = (np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3])) / 2.0
+        eh = (np.linalg.norm(quad[0] - quad[3]) +
+              np.linalg.norm(quad[1] - quad[2])) / 2.0
+        ew = (np.linalg.norm(quad[0] - quad[1]) +
+              np.linalg.norm(quad[2] - quad[3])) / 2.0
 
         dense_sample_pts_num = max(2, int(ew))
-        dense_xy_center_line = xy_text[np.linspace(0, xy_text.shape[0] - 1, dense_sample_pts_num,
-                                                endpoint=True, dtype=np.float32).astype(np.int32)]
-
-        dense_xy_center_line_diff = dense_xy_center_line[1:] - dense_xy_center_line[:-1]
-        estimate_arc_len = np.sum(np.linalg.norm(dense_xy_center_line_diff, axis=1))
+        dense_xy_center_line = xy_text[np.linspace(
+            0,
+            xy_text.shape[0] - 1,
+            dense_sample_pts_num,
+            endpoint=True,
+            dtype=np.float32).astype(np.int32)]
+
+        dense_xy_center_line_diff = dense_xy_center_line[
+            1:] - dense_xy_center_line[:-1]
+        estimate_arc_len = np.sum(
+            np.linalg.norm(
+                dense_xy_center_line_diff, axis=1))
 
         sample_pts_num = max(2, int(estimate_arc_len / eh))
         return sample_pts_num
 
-    def detect_sast(self, tcl_map, tvo_map, tbo_map, tco_map, ratio_w, ratio_h, src_w, src_h, 
-                shrink_ratio_of_width=0.3, tcl_map_thresh=0.5, offset_expand=1.0, out_strid=4.0):
+    def detect_sast(self,
+                    tcl_map,
+                    tvo_map,
+                    tbo_map,
+                    tco_map,
+                    ratio_w,
+                    ratio_h,
+                    src_w,
+                    src_h,
+                    shrink_ratio_of_width=0.3,
+                    tcl_map_thresh=0.5,
+                    offset_expand=1.0,
+                    out_strid=4.0):
         """
         first resize the tcl_map, tvo_map and tbo_map to the input_size, then restore the polys
         """
         # restore quad
-        scores, quads, xy_text = self.restore_quad(tcl_map, tcl_map_thresh, tvo_map)
+        scores, quads, xy_text = self.restore_quad(tcl_map, tcl_map_thresh,
+                                                   tvo_map)
         dets = np.hstack((quads, scores)).astype(np.float32, copy=False)
         dets = self.nms(dets)
         if dets.shape[0] == 0:
@@ -202,7 +234,8 @@ class SASTPostProcess(object):
 
         # instance segmentation
         # instance_count, instance_label_map = cv2.connectedComponents(tcl_map.astype(np.uint8), connectivity=8)
-        instance_count, instance_label_map = self.cluster_by_quads_tco(tcl_map, tcl_map_thresh, quads, tco_map)
+        instance_count, instance_label_map = self.cluster_by_quads_tco(
+            tcl_map, tcl_map_thresh, quads, tco_map)
 
         # restore single poly with tcl instance.
         poly_list = []
@@ -214,8 +247,8 @@ class SASTPostProcess(object):
                 continue
 
             #
-            len1 = float(np.linalg.norm(quad[0] -quad[1]))
-            len2 = float(np.linalg.norm(quad[1] -quad[2]))
+            len1 = float(np.linalg.norm(quad[0] - quad[1]))
+            len2 = float(np.linalg.norm(quad[1] - quad[2]))
             min_len = min(len1, len2)
             if min_len < 3:
                 continue
@@ -231,9 +264,11 @@ class SASTPostProcess(object):
                 continue
 
             # sort xy_text
-            left_center_pt = np.array([[(quad[0, 0] + quad[-1, 0]) / 2.0,
+            left_center_pt = np.array(
+                [[(quad[0, 0] + quad[-1, 0]) / 2.0,
                   (quad[0, 1] + quad[-1, 1]) / 2.0]])  # (1, 2)
-            right_center_pt = np.array([[(quad[1, 0] + quad[2, 0]) / 2.0,
+            right_center_pt = np.array(
+                [[(quad[1, 0] + quad[2, 0]) / 2.0,
                   (quad[1, 1] + quad[2, 1]) / 2.0]])  # (1, 2)
             proj_unit_vec = (right_center_pt - left_center_pt) / \
                             (np.linalg.norm(right_center_pt - left_center_pt) + 1e-6)
@@ -245,28 +280,40 @@ class SASTPostProcess(object):
                 sample_pts_num = self.estimate_sample_pts_num(quad, xy_text)
             else:
                 sample_pts_num = self.sample_pts_num
-            xy_center_line = xy_text[np.linspace(0, xy_text.shape[0] - 1, sample_pts_num,
-                                                endpoint=True, dtype=np.float32).astype(np.int32)]
+            xy_center_line = xy_text[np.linspace(
+                0,
+                xy_text.shape[0] - 1,
+                sample_pts_num,
+                endpoint=True,
+                dtype=np.float32).astype(np.int32)]
 
             point_pair_list = []
             for x, y in xy_center_line:
                 # get corresponding offset
                 offset = tbo_map[y, x, :].reshape(2, 2)
                 if offset_expand != 1.0:
-                    offset_length = np.linalg.norm(offset, axis=1, keepdims=True)
-                    expand_length = np.clip(offset_length * (offset_expand - 1), a_min=0.5, a_max=3.0)
+                    offset_length = np.linalg.norm(
+                        offset, axis=1, keepdims=True)
+                    expand_length = np.clip(
+                        offset_length * (offset_expand - 1),
+                        a_min=0.5,
+                        a_max=3.0)
                     offset_detal = offset / offset_length * expand_length
                     offset = offset + offset_detal
                     # original point
                 ori_yx = np.array([y, x], dtype=np.float32)
-                point_pair = (ori_yx +  offset)[:, ::-1]* out_strid / np.array([ratio_w, ratio_h]).reshape(-1, 2) 
+                point_pair = (ori_yx + offset)[:, ::-1] * out_strid / np.array(
+                    [ratio_w, ratio_h]).reshape(-1, 2)
                 point_pair_list.append(point_pair)
 
             # ndarry: (x, 2), expand poly along width
             detected_poly = self.point_pair2poly(point_pair_list)
-            detected_poly = self.expand_poly_along_width(detected_poly, shrink_ratio_of_width)
-            detected_poly[:, 0] = np.clip(detected_poly[:, 0], a_min=0, a_max=src_w)
-            detected_poly[:, 1] = np.clip(detected_poly[:, 1], a_min=0, a_max=src_h)
+            detected_poly = self.expand_poly_along_width(detected_poly,
+                                                         shrink_ratio_of_width)
+            detected_poly[:, 0] = np.clip(
+                detected_poly[:, 0], a_min=0, a_max=src_w)
+            detected_poly[:, 1] = np.clip(
+                detected_poly[:, 1], a_min=0, a_max=src_h)
             poly_list.append(detected_poly)
 
         return poly_list
@@ -285,16 +332,24 @@ class SASTPostProcess(object):
         img_num = len(shape_list)
         poly_lists = []
         for ino in range(img_num):
-            p_score = score_list[ino].transpose((1,2,0))
-            p_border = border_list[ino].transpose((1,2,0))
-            p_tvo = tvo_list[ino].transpose((1,2,0))
-            p_tco = tco_list[ino].transpose((1,2,0))
+            p_score = score_list[ino].transpose((1, 2, 0))
+            p_border = border_list[ino].transpose((1, 2, 0))
+            p_tvo = tvo_list[ino].transpose((1, 2, 0))
+            p_tco = tco_list[ino].transpose((1, 2, 0))
             src_h, src_w, ratio_h, ratio_w = shape_list[ino]
 
-            poly_list = self.detect_sast(p_score, p_tvo, p_border, p_tco, ratio_w, ratio_h, src_w, src_h, 
+            poly_list = self.detect_sast(
+                p_score,
+                p_tvo,
+                p_border,
+                p_tco,
+                ratio_w,
+                ratio_h,
+                src_w,
+                src_h,
                 shrink_ratio_of_width=self.shrink_ratio_of_width,
-                                         tcl_map_thresh=self.tcl_map_thresh, offset_expand=self.expand_scale)
+                tcl_map_thresh=self.tcl_map_thresh,
+                offset_expand=self.expand_scale)
             poly_lists.append({'points': np.array(poly_list)})
 
         return poly_lists
-

ppocr/utils/dict/arabic_dict.txt

+ 
+!
+#
+$
+%
+&
+'
+(
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+_
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+É
+é
+ء
+آ
+أ
+ؤ
+إ
+ئ
+ا
+ب
+ة
+ت
+ث
+ج
+ح
+خ
+د
+ذ
+ر
+ز
+س
+ش
+ص
+ض
+ط
+ظ
+ع
+غ
+ف
+ق
+ك
+ل
+م
+ن
+ه
+و
+ى
+ي
+ً
+ٌ
+ٍ
+َ
+ُ
+ِ
+ّ
+ْ
+ٓ
+ٔ
+ٰ
+ٱ
+ٹ
+پ
+چ
+ڈ
+ڑ
+ژ
+ک
+ڭ
+گ
+ں
+ھ
+ۀ
+ہ
+ۂ
+ۃ
+ۆ
+ۇ
+ۈ
+ۋ
+ی
+ې
+ے
+ۓ
+ە
+١
+٢
+٣
+٤
+٥
+٦
+٧
+٨
+٩

ppocr/utils/dict/cyrillic_dict.txt

+ 
+!
+#
+$
+%
+&
+'
+(
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+_
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+É
+é
+Ё
+Є
+І
+Ј
+Љ
+Ў
+А
+Б
+В
+Г
+Д
+Е
+Ж
+З
+И
+Й
+К
+Л
+М
+Н
+О
+П
+Р
+С
+Т
+У
+Ф
+Х
+Ц
+Ч
+Ш
+Щ
+Ъ
+Ы
+Ь
+Э
+Ю
+Я
+а
+б
+в
+г
+д
+е
+ж
+з
+и
+й
+к
+л
+м
+н
+о
+п
+р
+с
+т
+у
+ф
+х
+ц
+ч
+ш
+щ
+ъ
+ы
+ь
+э
+ю
+я
+ё
+ђ
+є
+і
+ј
+љ
+њ
+ћ
+ў
+џ
+Ґ
+ґ

ppocr/utils/dict/devanagari_dict.txt

+ 
+!
+#
+$
+%
+&
+'
+(
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+_
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+É
+é
+ँ
+ं
+ः
+अ
+आ
+इ
+ई
+उ
+ऊ
+ऋ
+ए
+ऐ
+ऑ
+ओ
+औ
+क
+ख
+ग
+घ
+ङ
+च
+छ
+ज
+झ
+ञ
+ट
+ठ
+ड
+ढ
+ण
+त
+थ
+द
+ध
+न
+ऩ
+प
+फ
+ब
+भ
+म
+य
+र
+ऱ
+ल
+ळ
+व
+श
+ष
+स
+ह
+़
+ा
+ि
+ी
+ु
+ू
+ृ
+ॅ
+े
+ै
+ॉ
+ो
+ौ
+्
+॒
+क़
+ख़
+ग़
+ज़
+ड़
+ढ़
+फ़
+ॠ
+।
+०
+१
+२
+३
+४
+५
+६
+७
+८
+९
+॰

ppocr/utils/dict/latin_dict.txt

+ 
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+<
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+]
+_
+`
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+{
+}
+¡
+£
+§
+ª
+«
+­
+°
+²
+³
+´
+µ
+·
+º
+»
+¿
+À
+Á
+Â
+Ä
+Å
+Ç
+È
+É
+Ê
+Ë
+Ì
+Í
+Î
+Ï
+Ò
+Ó
+Ô
+Õ
+Ö
+Ú
+Ü
+Ý
+ß
+à
+á
+â
+ã
+ä
+å
+æ
+ç
+è
+é
+ê
+ë
+ì
+í
+î
+ï
+ñ
+ò
+ó
+ô
+õ
+ö
+ø
+ù
+ú
+û
+ü
+ý
+ą
+Ć
+ć
+Č
+č
+Đ
+đ
+ę
+ı
+Ł
+ł
+ō
+Œ
+œ
+Š
+š
+Ÿ
+Ž
+ž
+ʒ
+β
+δ
+ε
+з
+Ṡ
+‘
+€
+™

ppocr/utils/e2e_metric/polygon_fast.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.
+import numpy as np
+from shapely.geometry import Polygon
+"""
+:param det_x: [1, N] Xs of detection's vertices 
+:param det_y: [1, N] Ys of detection's vertices
+:param gt_x: [1, N] Xs of groundtruth's vertices
+:param gt_y: [1, N] Ys of groundtruth's vertices
+
+##############
+All the calculation of 'AREA' in this script is handled by:
+1) First generating a binary mask with the polygon area filled up with 1's
+2) Summing up all the 1's
+"""
+
+
+def area(x, y):
+    polygon = Polygon(np.stack([x, y], axis=1))
+    return float(polygon.area)
+
+
+def approx_area_of_intersection(det_x, det_y, gt_x, gt_y):
+    """
+    This helper determine if both polygons are intersecting with each others with an approximation method.
+    Area of intersection represented by the minimum bounding rectangular [xmin, ymin, xmax, ymax]
+    """
+    det_ymax = np.max(det_y)
+    det_xmax = np.max(det_x)
+    det_ymin = np.min(det_y)
+    det_xmin = np.min(det_x)
+
+    gt_ymax = np.max(gt_y)
+    gt_xmax = np.max(gt_x)
+    gt_ymin = np.min(gt_y)
+    gt_xmin = np.min(gt_x)
+
+    all_min_ymax = np.minimum(det_ymax, gt_ymax)
+    all_max_ymin = np.maximum(det_ymin, gt_ymin)
+
+    intersect_heights = np.maximum(0.0, (all_min_ymax - all_max_ymin))
+
+    all_min_xmax = np.minimum(det_xmax, gt_xmax)
+    all_max_xmin = np.maximum(det_xmin, gt_xmin)
+    intersect_widths = np.maximum(0.0, (all_min_xmax - all_max_xmin))
+
+    return intersect_heights * intersect_widths
+
+
+def area_of_intersection(det_x, det_y, gt_x, gt_y):
+    p1 = Polygon(np.stack([det_x, det_y], axis=1)).buffer(0)
+    p2 = Polygon(np.stack([gt_x, gt_y], axis=1)).buffer(0)
+    return float(p1.intersection(p2).area)
+
+
+def area_of_union(det_x, det_y, gt_x, gt_y):
+    p1 = Polygon(np.stack([det_x, det_y], axis=1)).buffer(0)
+    p2 = Polygon(np.stack([gt_x, gt_y], axis=1)).buffer(0)
+    return float(p1.union(p2).area)
+
+
+def iou(det_x, det_y, gt_x, gt_y):
+    return area_of_intersection(det_x, det_y, gt_x, gt_y) / (
+        area_of_union(det_x, det_y, gt_x, gt_y) + 1.0)
+
+
+def iod(det_x, det_y, gt_x, gt_y):
+    """
+    This helper determine the fraction of intersection area over detection area
+    """
+    return area_of_intersection(det_x, det_y, gt_x, gt_y) / (
+        area(det_x, det_y) + 1.0)

ppocr/utils/e2e_utils/extract_batchsize.py

+import paddle
+import numpy as np
+import copy
+
+
+def org_tcl_rois(batch_size, pos_lists, pos_masks, label_lists, tcl_bs):
+    """
+    """
+    pos_lists_, pos_masks_, label_lists_ = [], [], []
+    img_bs = batch_size
+    ngpu = int(batch_size / img_bs)
+    img_ids = np.array(pos_lists, dtype=np.int32)[:, 0, 0].copy()
+    pos_lists_split, pos_masks_split, label_lists_split = [], [], []
+    for i in range(ngpu):
+        pos_lists_split.append([])
+        pos_masks_split.append([])
+        label_lists_split.append([])
+
+    for i in range(img_ids.shape[0]):
+        img_id = img_ids[i]
+        gpu_id = int(img_id / img_bs)
+        img_id = img_id % img_bs
+        pos_list = pos_lists[i].copy()
+        pos_list[:, 0] = img_id
+        pos_lists_split[gpu_id].append(pos_list)
+        pos_masks_split[gpu_id].append(pos_masks[i].copy())
+        label_lists_split[gpu_id].append(copy.deepcopy(label_lists[i]))
+    # repeat or delete
+    for i in range(ngpu):
+        vp_len = len(pos_lists_split[i])
+        if vp_len <= tcl_bs:
+            for j in range(0, tcl_bs - vp_len):
+                pos_list = pos_lists_split[i][j].copy()
+                pos_lists_split[i].append(pos_list)
+                pos_mask = pos_masks_split[i][j].copy()
+                pos_masks_split[i].append(pos_mask)
+                label_list = copy.deepcopy(label_lists_split[i][j])
+                label_lists_split[i].append(label_list)
+        else:
+            for j in range(0, vp_len - tcl_bs):
+                c_len = len(pos_lists_split[i])
+                pop_id = np.random.permutation(c_len)[0]
+                pos_lists_split[i].pop(pop_id)
+                pos_masks_split[i].pop(pop_id)
+                label_lists_split[i].pop(pop_id)
+    # merge
+    for i in range(ngpu):
+        pos_lists_.extend(pos_lists_split[i])
+        pos_masks_.extend(pos_masks_split[i])
+        label_lists_.extend(label_lists_split[i])
+    return pos_lists_, pos_masks_, label_lists_
+
+
+def pre_process(label_list, pos_list, pos_mask, max_text_length, max_text_nums,
+                pad_num, tcl_bs):
+    label_list = label_list.numpy()
+    batch, _, _, _ = label_list.shape
+    pos_list = pos_list.numpy()
+    pos_mask = pos_mask.numpy()
+    pos_list_t = []
+    pos_mask_t = []
+    label_list_t = []
+    for i in range(batch):
+        for j in range(max_text_nums):
+            if pos_mask[i, j].any():
+                pos_list_t.append(pos_list[i][j])
+                pos_mask_t.append(pos_mask[i][j])
+                label_list_t.append(label_list[i][j])
+    pos_list, pos_mask, label_list = org_tcl_rois(batch, pos_list_t, pos_mask_t,
+                                                  label_list_t, tcl_bs)
+    label = []
+    tt = [l.tolist() for l in label_list]
+    for i in range(tcl_bs):
+        k = 0
+        for j in range(max_text_length):
+            if tt[i][j][0] != pad_num:
+                k += 1
+            else:
+                break
+        label.append(k)
+    label = paddle.to_tensor(label)
+    label = paddle.cast(label, dtype='int64')
+    pos_list = paddle.to_tensor(pos_list)
+    pos_mask = paddle.to_tensor(pos_mask)
+    label_list = paddle.squeeze(paddle.to_tensor(label_list), axis=2)
+    label_list = paddle.cast(label_list, dtype='int32')
+    return pos_list, pos_mask, label_list, label

ppocr/utils/e2e_utils/visual.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.
+import numpy as np
+import cv2
+import time
+
+
+def resize_image(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
+
+    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)
+
+
+def resize_image_min(im, max_side_len=512):
+    """
+    """
+    h, w, _ = im.shape
+
+    resize_w = w
+    resize_h = h
+
+    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)
+
+
+def resize_image_for_totaltext(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 point_pair2poly(point_pair_list):
+    """
+    Transfer vertical point_pairs into poly point in clockwise.
+    """
+    pair_length_list = []
+    for point_pair in point_pair_list:
+        pair_length = np.linalg.norm(point_pair[0] - point_pair[1])
+        pair_length_list.append(pair_length)
+    pair_length_list = np.array(pair_length_list)
+    pair_info = (pair_length_list.max(), pair_length_list.min(),
+                 pair_length_list.mean())
+
+    point_num = len(point_pair_list) * 2
+    point_list = [0] * point_num
+    for idx, point_pair in enumerate(point_pair_list):
+        point_list[idx] = point_pair[0]
+        point_list[point_num - 1 - idx] = point_pair[1]
+    return np.array(point_list).reshape(-1, 2), pair_info
+
+
+def shrink_quad_along_width(quad, begin_width_ratio=0., end_width_ratio=1.):
+    """
+    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 expand_poly_along_width(poly, shrink_ratio_of_width=0.3):
+    """
+    expand poly along width.
+    """
+    point_num = poly.shape[0]
+    left_quad = np.array(
+        [poly[0], poly[1], poly[-2], poly[-1]], dtype=np.float32)
+    left_ratio = -shrink_ratio_of_width * np.linalg.norm(left_quad[0] - left_quad[3]) / \
+                 (np.linalg.norm(left_quad[0] - left_quad[1]) + 1e-6)
+    left_quad_expand = shrink_quad_along_width(left_quad, left_ratio, 1.0)
+    right_quad = np.array(
+        [
+            poly[point_num // 2 - 2], poly[point_num // 2 - 1],
+            poly[point_num // 2], poly[point_num // 2 + 1]
+        ],
+        dtype=np.float32)
+    right_ratio = 1.0 + \
+                  shrink_ratio_of_width * np.linalg.norm(right_quad[0] - right_quad[3]) / \
+                  (np.linalg.norm(right_quad[0] - right_quad[1]) + 1e-6)
+    right_quad_expand = shrink_quad_along_width(right_quad, 0.0, right_ratio)
+    poly[0] = left_quad_expand[0]
+    poly[-1] = left_quad_expand[-1]
+    poly[point_num // 2 - 1] = right_quad_expand[1]
+    poly[point_num // 2] = right_quad_expand[2]
+    return poly
+
+
+def norm2(x, axis=None):
+    if axis:
+        return np.sqrt(np.sum(x**2, axis=axis))
+    return np.sqrt(np.sum(x**2))
+
+
+def cos(p1, p2):
+    return (p1 * p2).sum() / (norm2(p1) * norm2(p2))

requirements.txt

@@ -8,3 +8,4 @@ tqdm
 numpy
 visualdl
 python-Levenshtein
+opencv-contrib-python
\ No newline at end of file

setup.py

@@ -32,7 +32,7 @@ setup(
     package_dir={'paddleocr': ''},
     include_package_data=True,
     entry_points={"console_scripts": ["paddleocr= paddleocr.paddleocr:main"]},
-    version='2.0.3',
+    version='2.0.4',
     install_requires=requirements,
     license='Apache License 2.0',
     description='Awesome OCR toolkits based on PaddlePaddle （8.6M ultra-lightweight pre-trained model, support training and deployment among server, mobile, embeded and IoT devices',

tools/infer/utility.py

@@ -74,6 +74,19 @@ def parse_args():
         "--vis_font_path", type=str, default="./doc/fonts/simfang.ttf")
     parser.add_argument("--drop_score", type=float, default=0.5)
 
+    # params for e2e
+    parser.add_argument("--e2e_algorithm", type=str, default='PGNet')
+    parser.add_argument("--e2e_model_dir", type=str)
+    parser.add_argument("--e2e_limit_side_len", type=float, default=768)
+    parser.add_argument("--e2e_limit_type", type=str, default='max')
+
+    # PGNet parmas
+    parser.add_argument("--e2e_pgnet_score_thresh", type=float, default=0.5)
+    parser.add_argument(
+        "--e2e_char_dict_path", type=str, default="./ppocr/utils/ic15_dict.txt")
+    parser.add_argument("--e2e_pgnet_valid_set", type=str, default='totaltext')
+    parser.add_argument("--e2e_pgnet_polygon", type=bool, default=True)
+
     # params for text classifier
     parser.add_argument("--use_angle_cls", type=str2bool, default=False)
     parser.add_argument("--cls_model_dir", type=str)
@@ -93,8 +106,10 @@ def create_predictor(args, mode, logger):
         model_dir = args.det_model_dir
     elif mode == 'cls':
         model_dir = args.cls_model_dir
-    else:
+    elif mode == 'rec':
         model_dir = args.rec_model_dir
+    else:
+        model_dir = args.e2e_model_dir
 
     if model_dir is None:
         logger.info("not find {} model file path {}".format(mode, model_dir))
@@ -148,6 +163,22 @@ def create_predictor(args, mode, logger):
     return predictor, input_tensor, output_tensors
 
 
+def draw_e2e_res(dt_boxes, strs, img_path):
+    src_im = cv2.imread(img_path)
+    for box, str in zip(dt_boxes, strs):
+        box = box.astype(np.int32).reshape((-1, 1, 2))
+        cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
+        cv2.putText(
+            src_im,
+            str,
+            org=(int(box[0, 0, 0]), int(box[0, 0, 1])),
+            fontFace=cv2.FONT_HERSHEY_COMPLEX,
+            fontScale=0.7,
+            color=(0, 255, 0),
+            thickness=1)
+    return src_im
+
+
 def draw_text_det_res(dt_boxes, img_path):
     src_im = cv2.imread(img_path)
     for box in dt_boxes:

tools/program.py

@@ -375,7 +375,8 @@ def preprocess(is_train=False):
 
     alg = config['Architecture']['algorithm']
     assert alg in [
-        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS'
+        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
+        'CLS', 'PGNet'
     ]
 
     device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'