基于PaddleOCR复现文本识别算法SATRN

https://aistudio.baidu.com/aistudio/projectdetail/5298068?forkThirdPart=1

AI Studio

697人浏览 · 2022-12-21 15:25:32

AI Studio · 2022-12-21 15:25:32 发布

请点击此处查看本环境基本用法.

Please click here for more detailed instructions.

论文名称

On Recognizing Texts of Arbitrary Shapes with 2D Self-Attention

Junyeop Lee, Sungrae Park, Jeonghun Baek, Seong Joon Oh, Seonghyeon Kim, Hwalsuk Lee

1. 简介

在文本识别任务中，任意形状的文本识别依然具有很大的挑战性，文章中提出了SATRN网络结构，通过全局的注意力机制提高了模型在对任意形状的文本识别性能。SATRN模型是基于Transformer结构的文本识别网络，利用自注意力机制来描述场景图像中字符与字符之间的二维空间依赖关系。
本项目基于PaddleOCR，参考了MMOCR开源的PyTorch代码，提高了复现论文的效率。

2. 数据集和复现精度

训练集

训练集为MJSynth和SynthText，具体可参考MMOCR
测试集

数据集数量
IIIT5K 3000
SVT 647
IC13_1015 1015
IC15_2077 2077
SVTP 645
CUTE80 288
训练集和测试集已经包含在项目当中，若需要本地训练模型可以从此处下载
复现精度

数据集	数量
IIIT5K	3000
SVT	647
IC13_1015	1015
IC15_2077	2077
SVTP	645
CUTE80	288

Methods	Regular Text			Irregular Text			Avg acc
Name	IIIT5K	SVT	IC13	IC15	SVTP	CUTE80	Avg
Satrn_small(PyTorch)	94.23	90.11	95.67	78.86	85.74	84.72	88.23
Satrn_small(Paddle)	94.53	91.04	94.68	78.24	83.72	86.11	88.05

3. 开始使用

3.1 准备环境

解压数据集：

cd ~/data/data166485/
unzip rfl_dataset2.zip
cd ~/data/data182867/
unzip evaluation_academic.zip

安装第三方库：

cd ~/work/PaddleOCR_satrn
pip install -r requirements.txt

3.2 快速开始

**模型训练:

python3 -m paddle.distributed.launch --log_dir=./debug/ --gpus '0,1,2,3'  tools/train.py -c configs/rec/rec_satrn.yml \
    -o Train.dataset.data_dir=~/data/data166485/rfl_dataset2/training \
    Eval.dataset.data_dir=~/data/data182867/evaluation_academic \
    Global.save_model_dir=./output/rec_satrn/

**模型验证:

cfg_file=configs/rec/rec_satrn.yml
model_file=~/work/rec_satrn/best_accuracy.pdparams
data_dir=~/data/data182867/evaluation_academic

# IIIT5k
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/IIIT5k_3000

# SVT
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/SVT

# IC13
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/IC13_1015

# IC15
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/IC15_2077

# SVTP
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/SVTP

# CE80
python3 tools/eval.py -c $cfg_file \
            -o Global.pretrained_model=$model_file \
            Eval.dataset.data_dir=$data_dir/CUTE80

4. 代码结构

├─configs
├─doc                        
├─output                           
├─ppocr      
├─tools                                            
│  infer/                     
│  eval.py                        
│  export_model.py                        
│  infer_cls.py                        
│  infer_det.py                        
│  infer_e2e.py                        
│  infer_rec.py                        
│  infer_table.py                        
│  program.py                     
│  test_hubserving.py                   
│  train.py  
├─train.sh 
├─requirements.txt
├─setup.py 
├─README.md

5. 复现心得

本项目是参加论文复现营活动中完成的，在本次的SATRN论文复现过程中，首先，为了更好的完成复现，对模型的网络结构、损失函数以及训练细节等都需要掌握的更加彻底，使我更加深入的学习到了这篇论文里面的知识；其次，需要对MMOCR和PaddleOCR这两个OCR套件比较熟悉才能完成复现，这驱动着我深入学习了解了这两个OCR套件；最后，在完成各个模块的对齐过程当中，我学习到了Paddle和PyTorch框架下不同算子的定义方式以及它们在功能实现上的异同。总之，通过本次论文复现，我学习到了很多之前仅仅通过阅读论文无法学习到的知识，感谢论文复现官方组织的这次活动。

此文章为搬运
原项目链接