pgnet.md 10.6 KB
Newer Older
Jethong's avatar
Jethong committed
1
2
3
4
# 端对端OCR算法-PGNet
- [一、简介](#简介)
- [二、环境配置](#环境配置)
- [三、快速使用](#快速使用)
Jethong's avatar
Jethong committed
5
- [四、模型训练、评估、推理](#模型训练、评估、推理)
Jethong's avatar
Jethong committed
6
7

<a name="简介"></a>
Double_V's avatar
Double_V committed
8
## 一、简介
Jethong's avatar
Jethong committed
9
10
11
12
13
14
15
16
17
18
OCR算法可以分为两阶段算法和端对端的算法。二阶段OCR算法一般分为两个部分,文本检测和文本识别算法,文件检测算法从图像中得到文本行的检测框,然后识别算法去识别文本框中的内容。而端对端OCR算法可以在一个算法中完成文字检测和文字识别,其基本思想是设计一个同时具有检测单元和识别模块的模型,共享其中两者的CNN特征,并联合训练。由于一个算法即可完成文字识别,端对端模型更小,速度更快。

### PGNet算法介绍
近些年来,端对端OCR算法得到了良好的发展,包括MaskTextSpotter系列、TextSnake、TextDragon、PGNet系列等算法。在这些算法中,PGNet算法具备其他算法不具备的优势,包括:
- 设计PGNet loss指导训练,不需要字符级别的标注
- 不需要NMS和ROI相关操作,加速预测
- 提出预测文本行内的阅读顺序模块;
- 提出基于图的修正模块(GRM)来进一步提高模型识别性能
- 精度更高,预测速度更快

Jethong's avatar
Jethong committed
19
PGNet算法细节详见[论文](https://www.aaai.org/AAAI21Papers/AAAI-2885.WangP.pdf) ,算法原理图如下所示:
Jethong's avatar
Jethong committed
20
21
22
![](../pgnet_framework.png)
输入图像经过特征提取送入四个分支,分别是:文本边缘偏移量预测TBO模块,文本中心线预测TCL模块,文本方向偏移量预测TDO模块,以及文本字符分类图预测TCC模块。
其中TBO以及TCL的输出经过后处理后可以得到文本的检测结果,TCL、TDO、TCC负责文本识别。
Jethong's avatar
Jethong committed
23

Jethong's avatar
Jethong committed
24
其检测识别效果图如下:
Jethong's avatar
Jethong committed
25

Jethong's avatar
Jethong committed
26
27
28
![](../imgs_results/e2e_res_img293_pgnet.png)
![](../imgs_results/e2e_res_img295_pgnet.png)

Jethong's avatar
Jethong committed
29
### 性能指标
Jethong's avatar
Jethong committed
30
31
32
####测试集: Total Text
####测试环境: NVIDIA Tesla V100-SXM2-16GB
| |det_precision|det_recall|det_f_score|e2e_precision|e2e_recall|e2e_f_score|FPS|下载|
Jethong's avatar
Jethong committed
33
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
Jethong's avatar
Jethong committed
34
35
|Paper|85.30|86.80|86.1|-|-|61.7|38.20 (size=640)|-|
|Ours|87.03|82.48|84.69|61.71|58.43|60.03|48.73 (size=768)|[下载链接](https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/en_server_pgnetA.tar)|
Jethong's avatar
Jethong committed
36
37
38

*note:PaddleOCR里的PGNet实现针对预测速度做了优化,在精度下降可接受范围内,可以显著提升端对端预测速度*

Jethong's avatar
Jethong committed
39
40


Jethong's avatar
Jethong committed
41
<a name="环境配置"></a>
Double_V's avatar
Double_V committed
42
## 二、环境配置
Jethong's avatar
Jethong committed
43
44
45
请先参考[快速安装](./installation.md)配置PaddleOCR运行环境。

<a name="快速使用"></a>
Double_V's avatar
Double_V committed
46
## 三、快速使用
Jethong's avatar
Jethong committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
### inference模型下载
本节以训练好的端到端模型为例,快速使用模型预测,首先下载训练好的端到端inference模型[下载地址](https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/e2e_server_pgnetA_infer.tar)
```
mkdir inference && cd inference
# 下载英文端到端模型并解压
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/e2e_server_pgnetA_infer.tar && tar xf e2e_server_pgnetA_infer.tar
```
* windows 环境下如果没有安装wget,下载模型时可将链接复制到浏览器中下载,并解压放置在相应目录下

解压完毕后应有如下文件结构:
```
├── e2e_server_pgnetA_infer
│   ├── inference.pdiparams
│   ├── inference.pdiparams.info
│   └── inference.pdmodel
```
### 单张图像或者图像集合预测
```bash
# 预测image_dir指定的单张图像
Jethong's avatar
Jethong committed
66
python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/img623.jpg" --e2e_model_dir="./inference/e2e_server_pgnetA_infer/" --e2e_pgnet_polygon=True
Jethong's avatar
Jethong committed
67
68

# 预测image_dir指定的图像集合
Jethong's avatar
Jethong committed
69
python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/" --e2e_model_dir="./inference/e2e_server_pgnetA_infer/" --e2e_pgnet_polygon=True
Jethong's avatar
Jethong committed
70
71

# 如果想使用CPU进行预测,需设置use_gpu参数为False
Jethong's avatar
Jethong committed
72
python3 tools/infer/predict_e2e.py --e2e_algorithm="PGNet" --image_dir="./doc/imgs_en/img623.jpg" --e2e_model_dir="./inference/e2e_server_pgnetA_infer/" --e2e_pgnet_polygon=True --use_gpu=False
Jethong's avatar
Jethong committed
73
```
Jethong's avatar
Jethong committed
74
### 可视化结果
Jethong's avatar
Jethong committed
75
76
可视化文本检测结果默认保存到./inference_results文件夹里面,结果文件的名称前缀为'e2e_res'。结果示例如下:
![](../imgs_results/e2e_res_img623_pgnet.jpg)
Double_V's avatar
Double_V committed
77

Jethong's avatar
Jethong committed
78
<a name="模型训练、评估、推理"></a>
79
## 四、模型训练、评估、推理
Jethong's avatar
Jethong committed
80
本节以totaltext数据集为例,介绍PaddleOCR中端到端模型的训练、评估与测试。
Double_V's avatar
Double_V committed
81
82

###  准备数据
Jethong's avatar
Jethong committed
83
下载解压[totaltext](https://github.com/cs-chan/Total-Text-Dataset/blob/master/Dataset/README.md) 数据集到PaddleOCR/train_data/目录,数据集组织结构:
Jethong's avatar
Jethong committed
84
85
```
/PaddleOCR/train_data/total_text/train/
Double_V's avatar
Double_V committed
86
  |- rgb/            # total_text数据集的训练数据
Jethong's avatar
Jethong committed
87
88
      |- gt_0.png
      | ...  
Double_V's avatar
Double_V committed
89
  |- total_text.txt  # total_text数据集的训练标注
Jethong's avatar
Jethong committed
90
91
```

Double_V's avatar
Double_V committed
92
total_text.txt标注文件格式如下,文件名和标注信息中间用"\t"分隔:
Jethong's avatar
Jethong committed
93
94
95
96
97
98
99
100
```
" 图像文件名                    json.dumps编码的图像标注信息"
rgb/gt_0.png    [{"transcription": "EST", "points": [[1004.0,689.0],[1019.0,698.0],[1034.0,708.0],[1049.0,718.0],[1064.0,728.0],[1079.0,738.0],[1095.0,748.0],[1094.0,774.0],[1079.0,765.0],[1065.0,756.0],[1050.0,747.0],[1036.0,738.0],[1021.0,729.0],[1007.0,721.0]]}, {...}]
```
json.dumps编码前的图像标注信息是包含多个字典的list,字典中的 `points` 表示文本框的四个点的坐标(x, y),从左上角的点开始顺时针排列。
`transcription` 表示当前文本框的文字,**当其内容为“###”时,表示该文本框无效,在训练时会跳过。**
如果您想在其他数据集上训练,可以按照上述形式构建标注文件。

Double_V's avatar
Double_V committed
101
### 启动训练
Jethong's avatar
Jethong committed
102

Double_V's avatar
Double_V committed
103
PGNet训练分为两个步骤:step1: 在合成数据上训练,得到预训练模型,此时模型精度依然较低;step2: 加载预训练模型,在totaltext数据集上训练;为快速训练,我们直接提供了step1的预训练模型。
Jethong's avatar
Jethong committed
104
105
```shell
cd PaddleOCR/
Double_V's avatar
Double_V committed
106
下载step1 预训练模型
Jethong's avatar
Jethong committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
wget -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/train_step1.tar
可以得到以下的文件格式
./pretrain_models/train_step1/
  └─ best_accuracy.pdopt
  └─ best_accuracy.states
  └─ best_accuracy.pdparams
```
*如果您安装的是cpu版本,请将配置文件中的 `use_gpu` 字段修改为false*

```shell
# 单机单卡训练 e2e 模型
python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.pretrained_model=./pretrain_models/train_step1/best_accuracy Global.load_static_weights=False
# 单机多卡训练,通过 --gpus 参数设置使用的GPU ID
python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.pretrained_model=./pretrain_models/train_step1/best_accuracy  Global.load_static_weights=False
```

上述指令中,通过-c 选择训练使用configs/e2e/e2e_r50_vd_pg.yml配置文件。
有关配置文件的详细解释,请参考[链接](./config.md)

您也可以通过-o参数在不需要修改yml文件的情况下,改变训练的参数,比如,调整训练的学习率为0.0001
```shell
python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Optimizer.base_lr=0.0001
```

Jethong's avatar
Jethong committed
131
132
133
134
135
#### 断点训练
如果训练程序中断,如果希望加载训练中断的模型从而恢复训练,可以通过指定Global.checkpoints指定要加载的模型路径:
```shell
python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.checkpoints=./your/trained/model
```
Jethong's avatar
Jethong committed
136

Jethong's avatar
Jethong committed
137
**注意**`Global.checkpoints`的优先级高于`Global.pretrain_weights`的优先级,即同时指定两个参数时,优先加载`Global.checkpoints`指定的模型,如果`Global.checkpoints`指定的模型路径有误,会加载`Global.pretrain_weights`指定的模型。
Jethong's avatar
Jethong committed
138
139
140
141
142
143
144
145
146
147
148

PaddleOCR计算三个OCR端到端相关的指标,分别是:Precision、Recall、Hmean。

运行如下代码,根据配置文件`e2e_r50_vd_pg.yml``save_res_path`指定的测试集检测结果文件,计算评估指标。

评估时设置后处理参数`max_side_len=768`,使用不同数据集、不同模型训练,可调整参数进行优化
训练中模型参数默认保存在`Global.save_model_dir`目录下。在评估指标时,需要设置`Global.checkpoints`指向保存的参数文件。
```shell
python3 tools/eval.py -c configs/e2e/e2e_r50_vd_pg.yml  -o Global.checkpoints="{path/to/weights}/best_accuracy"
```

Double_V's avatar
Double_V committed
149
### 模型预测
Jethong's avatar
Jethong committed
150
151
测试单张图像的端到端识别效果
```shell
Jethong's avatar
Jethong committed
152
python3 tools/infer_e2e.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.infer_img="./doc/imgs_en/img_10.jpg" Global.pretrained_model="./output/e2e_pgnet/best_accuracy" Global.load_static_weights=false
Jethong's avatar
Jethong committed
153
154
155
156
```

测试文件夹下所有图像的端到端识别效果
```shell
Jethong's avatar
Jethong committed
157
python3 tools/infer_e2e.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.infer_img="./doc/imgs_en/" Global.pretrained_model="./output/e2e_pgnet/best_accuracy" Global.load_static_weights=false
Jethong's avatar
Jethong committed
158
159
```

Jethong's avatar
Jethong committed
160
### 预测推理
Jethong's avatar
Jethong committed
161
#### (1). 四边形文本检测模型(ICDAR2015)  
Jethong's avatar
Jethong committed
162
163
164
首先将PGNet端到端训练过程中保存的模型,转换成inference model。以基于Resnet50_vd骨干网络,以英文数据集训练的模型为例[模型下载地址](https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/en_server_pgnetA.tar) ,可以使用如下命令进行转换:
```
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/en_server_pgnetA.tar && tar xf en_server_pgnetA.tar
Jethong's avatar
Jethong committed
165
python3 tools/export_model.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.pretrained_model=./en_server_pgnetA/best_accuracy Global.load_static_weights=False Global.save_inference_dir=./inference/e2e
Jethong's avatar
Jethong committed
166
167
168
169
170
171
172
173
174
```
**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)

Jethong's avatar
Jethong committed
175
#### (2). 弯曲文本检测模型(Total-Text)
Jethong's avatar
Jethong committed
176
177
178
179
180
181
182
183
184
对于弯曲文本样例

**PGNet端到端模型推理,需要设置参数`--e2e_algorithm="PGNet"`,同时,还需要增加参数`--e2e_pgnet_polygon=True`,**可以执行如下命令:
```
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
```
可视化文本端到端结果默认保存到`./inference_results`文件夹里面,结果文件的名称前缀为'e2e_res'。结果示例如下:

![](../imgs_results/e2e_res_img623_pgnet.jpg)