pgnet.md 10.4 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
29
![](../imgs_results/e2e_res_img293_pgnet.png)
![](../imgs_results/e2e_res_img295_pgnet.png)

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

<a name="快速使用"></a>
Double_V's avatar
Double_V committed
34
## 三、快速使用
Jethong's avatar
Jethong committed
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
### 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
54
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
55
56

# 预测image_dir指定的图像集合
Jethong's avatar
Jethong committed
57
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
58
59

# 如果想使用CPU进行预测,需设置use_gpu参数为False
Jethong's avatar
Jethong committed
60
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
61
```
Jethong's avatar
Jethong committed
62
### 可视化结果
Jethong's avatar
Jethong committed
63
64
可视化文本检测结果默认保存到./inference_results文件夹里面,结果文件的名称前缀为'e2e_res'。结果示例如下:
![](../imgs_results/e2e_res_img623_pgnet.jpg)
Double_V's avatar
Double_V committed
65

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

###  准备数据
Jethong's avatar
Jethong committed
71
下载解压[totaltext](https://github.com/cs-chan/Total-Text-Dataset/blob/master/Dataset/README.md) 数据集到PaddleOCR/train_data/目录,数据集组织结构:
Jethong's avatar
Jethong committed
72
73
```
/PaddleOCR/train_data/total_text/train/
Double_V's avatar
Double_V committed
74
  |- rgb/            # total_text数据集的训练数据
Jethong's avatar
Jethong committed
75
76
      |- gt_0.png
      | ...  
Double_V's avatar
Double_V committed
77
  |- total_text.txt  # total_text数据集的训练标注
Jethong's avatar
Jethong committed
78
79
```

Double_V's avatar
Double_V committed
80
total_text.txt标注文件格式如下,文件名和标注信息中间用"\t"分隔:
Jethong's avatar
Jethong committed
81
82
83
84
85
86
87
88
```
" 图像文件名                    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
89
### 启动训练
Jethong's avatar
Jethong committed
90

Double_V's avatar
Double_V committed
91
PGNet训练分为两个步骤:step1: 在合成数据上训练,得到预训练模型,此时模型精度依然较低;step2: 加载预训练模型,在totaltext数据集上训练;为快速训练,我们直接提供了step1的预训练模型。
Jethong's avatar
Jethong committed
92
93
```shell
cd PaddleOCR/
Double_V's avatar
Double_V committed
94
下载step1 预训练模型
Jethong's avatar
Jethong committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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
119
120
121
122
123
#### 断点训练
如果训练程序中断,如果希望加载训练中断的模型从而恢复训练,可以通过指定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
124

Jethong's avatar
Jethong committed
125
**注意**`Global.checkpoints`的优先级高于`Global.pretrain_weights`的优先级,即同时指定两个参数时,优先加载`Global.checkpoints`指定的模型,如果`Global.checkpoints`指定的模型路径有误,会加载`Global.pretrain_weights`指定的模型。
Jethong's avatar
Jethong committed
126
127
128
129
130
131
132
133
134
135
136

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
137
### 模型预测
Jethong's avatar
Jethong committed
138
139
测试单张图像的端到端识别效果
```shell
Jethong's avatar
Jethong committed
140
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
141
142
143
144
```

测试文件夹下所有图像的端到端识别效果
```shell
Jethong's avatar
Jethong committed
145
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
146
147
```

Jethong's avatar
Jethong committed
148
### 预测推理
Jethong's avatar
Jethong committed
149
#### (1). 四边形文本检测模型(ICDAR2015)  
Jethong's avatar
Jethong committed
150
151
152
首先将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
153
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
154
155
156
157
158
159
160
161
162
```
**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
163
#### (2). 弯曲文本检测模型(Total-Text)
Jethong's avatar
Jethong committed
164
165
166
167
168
169
170
171
172
对于弯曲文本样例

**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)
Jethong's avatar
Jethong committed
173

Jethong's avatar
Jethong committed
174
175
176
177
178
179
#### (3). 性能指标
| |det_precision|det_recall|det_f_score|e2e_precision|e2e_recall|e2e_f_score|FPS (size=640)|
| --- | --- | --- | --- | --- | --- | --- | --- |
|Ours|87.03|82.48|84.69|61.71|58.43|60.03|62.61|
|Paper|85.30|86.80|86.1|-|-|61.7|38.20|
*note:PaddleOCR里的PGNet实现针对预测速度做了优化,在精度下降可接受范围内,可以显著提升端对端预测速度*