端對端OCR算法-PGNet

一、簡介

OCR算法可以分為兩階段算法和端對端的算法。二階段OCR算法一般分為兩個部分，文本檢測和文本識別算法，文件檢測算法從圖像中得到文本行的檢測框，然后識別算法去識別文本框中的內容。而端對端OCR算法可以在一個算法中完成文字檢測和文字識別，其基本思想是設計一個同時具有檢測單元和識別模塊的模型，共享其中兩者的CNN特征，並聯合訓練。由於一個算法即可完成文字識別，端對端模型更小，速度更快。

PGNet算法介紹

近些年來，端對端OCR算法得到了良好的發展，包括MaskTextSpotter系列、TextSnake、TextDragon、PGNet系列等算法。在這些算法中，PGNet算法具備其他算法不具備的優勢，包括：

設計PGNet loss指導訓練，不需要字符級別的標注
不需要NMS和ROI相關操作，加速預測
提出預測文本行內的閱讀順序模塊；
提出基於圖的修正模塊（GRM）來進一步提高模型識別性能
精度更高，預測速度更快

PGNet算法細節詳見論文 ,算法原理圖如下所示：輸入圖像經過特征提取送入四個分支，分別是：文本邊緣偏移量預測TBO模塊，文本中心線預測TCL模塊，文本方向偏移量預測TDO模塊，以及文本字符分類圖預測TCC模塊。其中TBO以及TCL的輸出經過后處理后可以得到文本的檢測結果，TCL、TDO、TCC負責文本識別。

其檢測識別效果圖如下：

性能指標

測試集: Total Text

測試環境: NVIDIA Tesla V100-SXM2-16GB

PGNetA	det_precision	det_recall	det_f_score	e2e_precision	e2e_recall	e2e_f_score	FPS	下載
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)	下載鏈接

note：PaddleOCR里的PGNet實現針對預測速度做了優化，在精度下降可接受范圍內，可以顯著提升端對端預測速度

二、環境配置

請先參考快速安裝配置PaddleOCR運行環境。

三、快速使用

inference模型下載

本節以訓練好的端到端模型為例，快速使用模型預測，首先下載訓練好的端到端inference模型下載地址

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

單張圖像或者圖像集合預測

# 預測image_dir指定的單張圖像 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  # 預測image_dir指定的圖像集合 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  # 如果想使用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_server_pgnetA_infer/" --e2e_pgnet_polygon=True --use_gpu=False

可視化結果

可視化文本檢測結果默認保存到./inference_results文件夾里面，結果文件的名稱前綴為'e2e_res'。結果示例如下：

四、模型訓練、評估、推理

本節以totaltext數據集為例，介紹PaddleOCR中端到端模型的訓練、評估與測試。

准備數據

下載解壓totaltext 數據集到PaddleOCR/train_data/目錄，數據集組織結構：

/PaddleOCR/train_data/total_text/train/
 |- rgb/ # total_text數據集的訓練數據  |- img11.jpg  | ...  |- train.txt # total_text數據集的訓練標注

total_text.txt標注文件格式如下，文件名和標注信息中間用"\t"分隔：

" 圖像文件名                    json.dumps編碼的圖像標注信息"
rgb/img11.jpg [{"transcription": "ASRAMA", "points": [[214.0, 325.0], [235.0, 308.0], [259.0, 296.0], [286.0, 291.0], [313.0, 295.0], [338.0, 305.0], [362.0, 320.0], [349.0, 347.0], [330.0, 337.0], [310.0, 329.0], [290.0, 324.0], [269.0, 328.0], [249.0, 336.0], [231.0, 346.0]]}, {...}]

json.dumps編碼前的圖像標注信息是包含多個字典的list，字典中的 points 表示文本框的四個點的坐標(x, y)，從左上角的點開始順時針排列。 transcription 表示當前文本框的文字，當其內容為“###”時，表示該文本框無效，在訓練時會跳過。如果您想在其他數據集上訓練，可以按照上述形式構建標注文件。

啟動訓練

PGNet訓練分為兩個步驟：step1: 在合成數據上訓練，得到預訓練模型，此時模型精度依然較低；step2: 加載預訓練模型，在totaltext數據集上訓練；為快速訓練，我們直接提供了step1的預訓練模型。

cd PaddleOCR/ 下載step1 預訓練模型 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

# 單機單卡訓練 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配置文件。有關配置文件的詳細解釋，請參考鏈接。

您也可以通過-o參數在不需要修改yml文件的情況下，改變訓練的參數，比如，調整訓練的學習率為0.0001

python3 tools/train.py -c configs/e2e/e2e_r50_vd_pg.yml -o Optimizer.base_lr=0.0001

斷點訓練

如果訓練程序中斷，如果希望加載訓練中斷的模型從而恢復訓練，可以通過指定Global.checkpoints指定要加載的模型路徑：

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。

運行如下代碼，根據配置文件e2e_r50_vd_pg.yml中save_res_path指定的測試集檢測結果文件，計算評估指標。

評估時設置后處理參數max_side_len=768，使用不同數據集、不同模型訓練，可調整參數進行優化訓練中模型參數默認保存在Global.save_model_dir目錄下。在評估指標時，需要設置Global.checkpoints指向保存的參數文件。

python3 tools/eval.py -c configs/e2e/e2e_r50_vd_pg.yml -o Global.checkpoints="{path/to/weights}/best_accuracy"

模型預測

測試單張圖像的端到端識別效果

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

測試文件夾下所有圖像的端到端識別效果

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

預測推理

(1). 四邊形文本檢測模型（ICDAR2015）

首先將PGNet端到端訓練過程中保存的模型，轉換成inference model。以基於Resnet50_vd骨干網絡，以英文數據集訓練的模型為例模型下載地址，可以使用如下命令進行轉換：

wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/pgnet/en_server_pgnetA.tar && tar xf en_server_pgnetA.tar
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

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'。結果示例如下：

(2). 彎曲文本檢測模型（Total-Text）

對於彎曲文本樣例

**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'。結果示例如下：