场景文本识别算法-ViTSTR¶

1. 算法简介¶

论文信息：

Vision Transformer for Fast and Efficient Scene Text Recognition Rowel Atienza ICDAR, 2021

ViTSTR使用MJSynth和SynthText两个文字识别数据集训练，在IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE数据集上进行评估，算法复现效果如下：

模型	骨干网络	配置文件	Acc	下载链接
ViTSTR	ViTSTR	rec_vitstr_none_ce.yml	79.82%	训练模型

2. 环境配置¶

请先参考《运行环境准备》配置PaddleOCR运行环境，参考《项目克隆》克隆项目代码。

3. 模型训练、评估、预测¶

3.1 模型训练¶

请参考文本识别训练教程。PaddleOCR对代码进行了模块化，训练ViTSTR识别模型时需要更换配置文件为ViTSTR的配置文件。

启动训练¶

具体地，在完成数据准备后，便可以启动训练，训练命令如下：

#单卡训练（训练周期长，不建议）
python3 tools/train.py -c configs/rec/rec_vitstr_none_ce.yml

# 多卡训练，通过--gpus参数指定卡号
python3 -m paddle.distributed.launch --gpus '0,1,2,3'  tools/train.py -c configs/rec/rec_vitstr_none_ce.yml

3.2 评估¶

可下载已训练完成的模型文件，使用如下命令进行评估：

# 注意将pretrained_model的路径设置为本地路径。
python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr_none_ce.yml -o Global.pretrained_model=./rec_vitstr_none_ce_train/best_accuracy

3.3 预测¶

使用如下命令进行单张图片预测：

# 注意将pretrained_model的路径设置为本地路径。
python3 tools/infer_rec.py -c configs/rec/rec_vitstr_none_ce.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_vitstr_none_ce_train/best_accuracy
# 预测文件夹下所有图像时，可修改infer_img为文件夹，如 Global.infer_img='./doc/imgs_words_en/'。

4. 推理部署¶

4.1 Python推理¶

首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（模型下载地址 )，可以使用如下命令进行转换：

# 注意将pretrained_model的路径设置为本地路径。
python3 tools/export_model.py -c configs/rec/rec_vitstr_none_ce.yml -o Global.pretrained_model=./rec_vitstr_none_ce_train/best_accuracy Global.save_inference_dir=./inference/rec_vitstr/

注意：

如果您是在自己的数据集上训练的模型，并且调整了字典文件，请注意修改配置文件中的character_dict_path是否是所需要的字典文件。
如果您修改了训练时的输入大小，请修改tools/export_model.py文件中的对应ViTSTR的infer_shape。

转换成功后，在目录下有三个文件：

/inference/rec_vitstr/
    ├── inference.pdiparams         # 识别inference模型的参数文件
    ├── inference.pdiparams.info    # 识别inference模型的参数信息，可忽略
    └── inference.pdmodel           # 识别inference模型的program文件

执行如下命令进行模型推理：

python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' --rec_model_dir='./inference/rec_vitstr/' --rec_algorithm='ViTSTR' --rec_image_shape='1,224,224' --rec_char_dict_path='./ppocr/utils/EN_symbol_dict.txt'
# 预测文件夹下所有图像时，可修改image_dir为文件夹，如 --image_dir='./doc/imgs_words_en/'。

执行命令后，上面图像的预测结果（识别的文本和得分）会打印到屏幕上，示例如下：结果如下：

Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9998350143432617)

注意：

训练上述模型采用的图像分辨率是[1，224，224]，需要通过参数rec_image_shape设置为您训练时的识别图像形状。
在推理时需要设置参数rec_char_dict_path指定字典，如果您修改了字典，请修改该参数为您的字典文件。
如果您修改了预处理方法，需修改tools/infer/predict_rec.py中ViTSTR的预处理为您的预处理方法。

4.2 C++推理部署¶

由于C++预处理后处理还未支持ViTSTR，所以暂未支持

4.3 Serving服务化部署¶

暂不支持

4.4 更多推理部署¶