使用paddleOCR训练自己的数据集到ONNX推理

news2024/9/27 19:15:29

一、环境安装

1、安装paddlepaddle；
https://www.paddlepaddle.org.cn/
在这里插入图片描述
这里安装2.6.1的话使用onnx会出现swish算子报错的问题

python -m pip install paddlepaddle-gpu==2.5.2 -i https://pypi.tuna.tsinghua.edu.cn/simple

验证是否成功安装

python
import paddle
paddle.utils.run_check()

2、安装paddleOCR
直接按照，失败，失败安装的包就手动安装

pip install paddleocr -i https://mirror.baidu.com/pypi/simple

或者

git clone https://gitee.com/PaddlePaddle/PaddleOCR
cd PaddleOCR
python -m pip install -r requirments.txt
pip install -r requirements.txt -i https://mirror.baidu.com/pypi/simple

3、安装onnx

pip install paddle2onnx -i https://mirror.baidu.com/pypi/simple
pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成
在这里插入图片描述
4、验证是否安装正常

cd D:\CodeGitDownload\PaddleOCR
paddleocr --image_dir ./doc/imgs/11.jpg --use_angle_cls true --use_gpu false

在这里插入图片描述

二、数据集标注

打开标注软件，在PaddleOCR代码中

python PPOCRLabel.py --lang ch

在这里插入图片描述
使用说明：https://blog.csdn.net/Castlehe/article/details/115699134
注意：如果要训练识别模型应该在文件-》导出识别结果，不然rec文件夹中没数据；

标注完成后提出软件,，对刚才标注的图片进行划分

python gen_ocr_train_val_test.py --trainValTestRatio 6:2:2 --datasetRootPath E:/Dataset--recRootPath E:/Dataset/rec --detRootPath E:/Dataset/det

三、训练

因为识别的产品字符只包含了数字和字母，而且经过测试发现检测推理模型是可以直接用，所有这里使用英文识别模型，不再训练推理模型。
模型下载
在这里插入图片描述

下载所需文件内容参考：
在这里插入图片描述
修改yaml文件参数，主要如下；

开始训练

python tools/train.py -c E:/ModelOCR/en_PP-OCRv4_rec.yml

训练结束后，将训练模型（.pdopt、/.pdparams）转为推理模型（.pdmodel、/.pdiparams）

python tools/export_model.py -c E:/ModelOCR/en_PP-OCRv4_rec.yml -o Global.checkpoints=E:/ModelOCR/train_rec_output/best_accuracy Global.save_inference_dir=E:/ModelOCR/infer_rec_output/

python tools/infer/predict_rec.py --image_dir="E:/Dataset/neimeng0707/rec/train/Image_20230818133344612_crop_0.jpg" --rec_model_dir="E:/ModelOCR/infer_rec_output/" --rec_image_shape="3, 48, 320" --rec_char_dict_path="E:/ModelOCR/en_dict.txt"

在这里插入图片描述

三、ONNX推理

paddle2onnx --model_dir E:/ModelOCR/infer_rec_output --model_filename inference.pdmodel --params_filename inference.pdiparams --save_file E:/ModelOCR/infer_rec_output/inference.onnx --opset_version 11 --enable_onnx_checker True

这里使用C++的ONNX进行推理，参考https://github.com/hpc203/PaddleOCR-v3-onnxrun-cpp-py
考虑兼容问题，版本用的是：onnxruntime-win-x64-1.12.0
使用训练后的模型推理效果如下，B错识别为8，收集数据集后重新迭代训练即可优化；
在这里插入图片描述