PaddleOCR文字识别模型训练
PaddleOCR文字识别模型训练
- 数据准备
-
- 中文场景文字识别技术创新大赛数据集
- 启动训练
- 预测
-
- 预测中文结果
本文主要介绍PaddleOCR中文字识别模型的训练、评估及测试。
数据准备
中文场景文字识别技术创新大赛数据集
数据集共212023张文字图片。数据集地址
- 训练集
将训练图片放入同一个文件夹（train_images），并用一个txt文件（rec_gt_train.txt）记录图片路径和标签。
注意： 默认请将图片路径和图片标签用 \t 分割，如用其他方式分割将造成训练报错
" 图像文件名 图像标注信息 "
img_001.jpg 简单可依赖
img_002.jpg 用科技让复杂的世界更简单
训练集txt生成代码如下代码：
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Auther: Andy
date: 11/25/20 1:36 PM
desc:
"""
train_list_path = "/data11/shaozhl/datasets/OCR/baidu_rec/train.list"
label_path = "/data11/shaozhl/datasets/OCR/baidu_rec/train_data/rec_data/rec_gt_train.txt"
with open(train_list_path, 'r') as fr:
with open(label_path, 'w', encoding="utf-8") as fw:
lines = fr.readlines()
for line in lines:
line = line.split("\t")
fw.writelines(line[-2] + "\t" + line[-1])
最终训练集应有如下文件结构：
|-train_data
|-rec_data
|- rec_gt_train.txt
|- train
|- img_001.png
|- img_002.jpg
|- img_003.jpg
| ...
测试集同理，
- 字典
因为是所用的数据集已经有标注（标注列表为train.list），所以只要将标注的数据转化为字典即可。
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Auther: Andy
date: 11/25/20 11:09 AM
desc: 将中文场景文字识别技术创新大赛数据集label转化为PaddleOCR文字识别label
"""
train_list_path = "/data11/shaozhl/datasets/OCR/baidu_rec/train.list"
dict_path = "../ppocr/utils/baidu_rec_dict.txt"
texts = ""
with open(train_list_path, 'r') as f:
lines = f.readlines()
for line in lines:
line = line.split("\t")
texts += line[-1].replace("\n", "")
texts = set(texts)
with open(dict_path, 'w') as f:
for text in texts:
f.writelines(text+"\n")
在 configs/rec/rec_baidu_train.yml 中添加 character_dict_path 字段, 指向您的字典路径。 并将 character_type 设置为 ch，因为该数据集label有的比较长，所以也要设置下max_text_length
character_dict_path: ./ppocr/utils/baidu_rec_dict.txt
character_type: ch
max_text_length: 100
rec_baidu_train.yml可参照rec_icdar15_train.yml重写。
下面是我自己的yml，仅供参考。rec_baidu_train.yml如下
Global:
algorithm: CRNN
use_gpu: true
epoch_num: 1000
log_smooth_window: 20
print_batch_step: 10
save_model_dir: ./output/rec_baidu_CRNN
save_epoch_step: 50
eval_batch_step: 500
train_batch_size_per_card: 256
test_batch_size_per_card: 256
image_shape: [3, 48, 140]
max_text_length: 100
character_type: ch
loss_type: ctc
distort: true
debug: false
reader_yml: ./configs/rec/rec_baidu_reader.yml
pretrain_weights: ./pretrain_models/rec_mv3_none_bilstm_ctc/best_accuracy
checkpoints:
save_inference_dir:
infer_img:
character_dict_path: ./ppocr/utils/baidu_rec_dict.txt
use_space_char: true
Architecture:
function: ppocr.modeling.architectures.rec_model,RecModel
Backbone:
function: ppocr.modeling.backbones.rec_mobilenet_v3,MobileNetV3
scale: 0.5
model_name: large
Head:
function: ppocr.modeling.heads.rec_ctc_head,CTCPredict
encoder_type: rnn
SeqRNN:
hidden_size: 96
Loss:
function: ppocr.modeling.losses.rec_ctc_loss,CTCLoss
Optimizer:
function: ppocr.optimizer,AdamDecay
base_lr: 0.0005
beta1: 0.9
beta2: 0.999
decay:
function: cosine_decay
step_each_epoch: 20
total_epoch: 1000
rec_baidu_reader.yml如下
TrainReader:
reader_function: ppocr.data.rec.dataset_traversal,SimpleReader
num_workers: 2
img_set_dir: /data11/shaozhl/datasets/OCR/baidu_rec/train_data/rec_data/train
label_file_path: /data11/shaozhl/datasets/OCR/baidu_rec/train_data/rec_data/rec_gt_train.txt
EvalReader:
reader_function: ppocr.data.rec.dataset_traversal,SimpleReader
img_set_dir: /data11/shaozhl/datasets/OCR/baidu_rec/train_data/rec_data/test
label_file_path: /data11/shaozhl/datasets/OCR/baidu_rec/train_data/rec_data/rec_gt_test.txt
TestReader:
reader_function: ppocr.data.rec.dataset_traversal,SimpleReader
- 添加空格类别
如果希望支持识别"空格"类别, 请将yml文件中的 use_space_char 字段设置为 true。
启动训练
PaddleOCR提供了训练脚本、评估脚本和预测脚本，本节将以 CRNN 识别模型为例：
首先下载pretrain model，您可以下载训练好的模型在 中文场景文字识别技术创新大赛数据集上进行finetune。
cd PaddleOCR/
# 下载MobileNetV3的预训练模型
wget -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/rec_mv3_none_bilstm_ctc.tar
# 解压模型参数
cd pretrain_models
tar -xf rec_mv3_none_bilstm_ctc.tar && rm -rf rec_mv3_none_bilstm_ctc.tar
开始训练:
如果安装的是cpu版本，请将配置文件中的 use_gpu 字段修改为false
# GPU训练 支持单卡，多卡训练，通过CUDA_VISIBLE_DEVICES指定卡号
export CUDA_VISIBLE_DEVICES=1,2
# 并将训练日志保存为 train_rec_baidu.log
python3 tools/train.py -c configs/rec/rec_baidu_train.yml 2>&1 | tee train_rec_baidu.log
PS
- 这个训练时间挺长的，一轮差不多要15分钟
- 怎么多利用GPU显存，现在还不清楚，我的一直是9000M，没有增加过，我的显卡是16G的
预测
预测使用的配置文件必须与训练一致，如通过 python3 tools/train.py -c configs/rec/rec_baidu_train.yml 完成了中文模型的训练，可以使用如下命令进行中文模型预测。
预测中文结果
python3 tools/infer_rec.py -c configs/rec/rec_baidu_train.yml -o Global.checkpoints=./output/rec_baidu_CRNN/best_accuracy Global.infer_img=doc/imgs_words/ch/word_1.jpg
预测图片：
得到输入图像的预测结果：
2020-12-08 10:25:45,504-INFO: Finish initing model from ./output/rec_baidu_CRNN/best_accuracy
2020-12-08 10:25:45,593-INFO: infer_img:doc/imgs_words/ch/word_1.jpg
2020-12-03 10:25:45,864-INFO: index: [3995 1072 1441 198]
2020-12-03 10:25:45,864-INFO: word : 韩国小馆
2020-12-03 10:25:45,864-INFO: score: 0.9923042058944702