From 9816aebdb078ec14ca3141560f2431981c8948f5 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Fri, 27 May 2022 06:49:52 +0000 Subject: [PATCH 01/43] add rec vitstr algorithm. --- configs/rec/rec_mtb_nrtr.yml | 10 +- configs/rec/rec_svtrnet.yml | 4 +- configs/rec/rec_vitstr.yml | 100 ++++++++++++ doc/doc_ch/algorithm_overview.md | 3 +- doc/doc_ch/algorithm_rec_vitstr.md | 154 ++++++++++++++++++ doc/doc_en/algorithm_overview_en.md | 3 +- doc/doc_en/algorithm_rec_vitstr_en.md | 134 +++++++++++++++ ppocr/data/imaug/__init__.py | 2 +- ppocr/data/imaug/label_ops.py | 35 +++- ppocr/data/imaug/rec_img_aug.py | 19 ++- ppocr/losses/__init__.py | 7 +- ...rec_nrtr_loss.py => rec_ce_smooth_loss.py} | 14 +- ppocr/modeling/backbones/__init__.py | 3 +- ppocr/modeling/backbones/rec_svtrnet.py | 102 ++++++------ ppocr/modeling/backbones/rec_vitstr.py | 120 ++++++++++++++ ppocr/postprocess/__init__.py | 4 +- ppocr/postprocess/rec_postprocess.py | 34 +++- .../configs/rec_mtb_nrtr/rec_mtb_nrtr.yml | 6 +- test_tipc/configs/rec_svtrnet/rec_svtrnet.yml | 119 ++++++++++++++ .../rec_svtrnet/train_infer_python.txt | 53 ++++++ test_tipc/configs/rec_vitstr/rec_vitstr.yml | 101 ++++++++++++ .../configs/rec_vitstr/train_infer_python.txt | 53 ++++++ tools/export_model.py | 6 + tools/infer/predict_rec.py | 19 ++- tools/program.py | 32 ++-- 25 files changed, 1044 insertions(+), 93 deletions(-) create mode 100644 configs/rec/rec_vitstr.yml create mode 100644 doc/doc_ch/algorithm_rec_vitstr.md create mode 100644 doc/doc_en/algorithm_rec_vitstr_en.md rename ppocr/losses/{rec_nrtr_loss.py => rec_ce_smooth_loss.py} (73%) create mode 100644 ppocr/modeling/backbones/rec_vitstr.py create mode 100644 test_tipc/configs/rec_svtrnet/rec_svtrnet.yml create mode 100644 test_tipc/configs/rec_svtrnet/train_infer_python.txt create mode 100644 test_tipc/configs/rec_vitstr/rec_vitstr.yml create mode 100644 test_tipc/configs/rec_vitstr/train_infer_python.txt diff --git a/configs/rec/rec_mtb_nrtr.yml b/configs/rec/rec_mtb_nrtr.yml index 0426750085..dfe2cc9811 100644 --- a/configs/rec/rec_mtb_nrtr.yml +++ b/configs/rec/rec_mtb_nrtr.yml @@ -49,7 +49,7 @@ Architecture: Loss: - name: NRTRLoss + name: CESmoothingLoss smoothing: True PostProcess: @@ -68,8 +68,8 @@ Train: img_mode: BGR channel_first: False - NRTRLabelEncode: # Class handling label - - NRTRRecResizeImg: - image_shape: [100, 32] + - GrayRecResizeImg: + image_shape: [100, 32] # W H resize_type: PIL # PIL or OpenCV - KeepKeys: keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order @@ -88,8 +88,8 @@ Eval: img_mode: BGR channel_first: False - NRTRLabelEncode: # Class handling label - - NRTRRecResizeImg: - image_shape: [100, 32] + - GrayRecResizeImg: + image_shape: [100, 32] # W H resize_type: PIL # PIL or OpenCV - KeepKeys: keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order diff --git a/configs/rec/rec_svtrnet.yml b/configs/rec/rec_svtrnet.yml index 233d5e2765..a700e4bd92 100644 --- a/configs/rec/rec_svtrnet.yml +++ b/configs/rec/rec_svtrnet.yml @@ -77,7 +77,7 @@ Metric: Train: dataset: name: LMDBDataSet - data_dir: ./train_data/data_lmdb_release/training/ + data_dir: ./train_data/data_lmdb_release/training transforms: - DecodeImage: # load image img_mode: BGR @@ -98,7 +98,7 @@ Train: Eval: dataset: name: LMDBDataSet - data_dir: ./train_data/data_lmdb_release/validation/ + data_dir: ./train_data/data_lmdb_release/validation transforms: - DecodeImage: # load image img_mode: BGR diff --git a/configs/rec/rec_vitstr.yml b/configs/rec/rec_vitstr.yml new file mode 100644 index 0000000000..005db0184a --- /dev/null +++ b/configs/rec/rec_vitstr.yml @@ -0,0 +1,100 @@ +Global: + use_gpu: True + epoch_num: 20 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/vitstr/ + save_epoch_step: 1 + # evaluation is run every 2000 iterations after the 0th iteration# + eval_batch_step: [0, 50] + cal_metric_during_train: True + pretrained_model: + checkpoints: + save_inference_dir: + use_visualdl: False + infer_img: doc/imgs_words_en/word_10.png + # for data or label process + character_dict_path: ppocr/utils/EN_symbol_dict.txt + max_text_length: 25 + infer_mode: False + use_space_char: False + save_res_path: ./output/rec/predicts_vitstr.txt + + +Optimizer: + name: Adadelta + epsilon: 0.00000001 + rho: 0.95 + clip_norm: 5.0 + lr: + learning_rate: 1.0 + +Architecture: + model_type: rec + algorithm: ViTSTR + in_channels: 1 + Transform: + Backbone: + name: ViTSTR + scale: tiny + Neck: + name: SequenceEncoder + encoder_type: reshape + Head: + name: CTCHead + +Loss: + name: CESmoothingLoss + smoothing: False + with_all: True + +PostProcess: + name: ViTSTRLabelDecode + +Metric: + name: RecMetric + main_indicator: acc + +Train: + dataset: + name: LMDBDataSet + data_dir: ./train_data/data_lmdb_release/training + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - ViTSTRLabelEncode: # Class handling label + - GrayRecResizeImg: + image_shape: [224, 224] # W H + resize_type: PIL # PIL or OpenCV + inter_type: 'Image.BICUBIC' + scale: false + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: True + batch_size_per_card: 48 + drop_last: True + num_workers: 2 + +Eval: + dataset: + name: LMDBDataSet + data_dir: ./train_data/data_lmdb_release/validation + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - ViTSTRLabelEncode: # Class handling label + - GrayRecResizeImg: + image_shape: [224, 224] # W H + resize_type: PIL # PIL or OpenCV + inter_type: 'Image.BICUBIC' + scale: false + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: False + drop_last: False + batch_size_per_card: 256 + num_workers: 2 diff --git a/doc/doc_ch/algorithm_overview.md b/doc/doc_ch/algorithm_overview.md index 6227a21498..934ac08537 100755 --- a/doc/doc_ch/algorithm_overview.md +++ b/doc/doc_ch/algorithm_overview.md @@ -66,6 +66,7 @@ - [x] [SAR](./algorithm_rec_sar.md) - [x] [SEED](./algorithm_rec_seed.md) - [x] [SVTR](./algorithm_rec_svtr.md) +- [x] [ViTSTR](./algorithm_rec_vitstr.md) 参考[DTRB](https://arxiv.org/abs/1904.01906)[3]文字识别训练和评估流程，使用MJSynth和SynthText两个文字识别数据集训练，在IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE数据集上进行评估，算法效果如下： @@ -84,7 +85,7 @@ |SAR|Resnet31| 87.20% | rec_r31_sar | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | - +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | diff --git a/doc/doc_ch/algorithm_rec_vitstr.md b/doc/doc_ch/algorithm_rec_vitstr.md new file mode 100644 index 0000000000..bd83b8d9c2 --- /dev/null +++ b/doc/doc_ch/algorithm_rec_vitstr.md @@ -0,0 +1,154 @@ +# 场景文本识别算法-ViTSTR + +- [1. 算法简介](#1) +- [2. 环境配置](#2) +- [3. 模型训练、评估、预测](#3) + - [3.1 训练](#3-1) + - [3.2 评估](#3-2) + - [3.3 预测](#3-3) +- [4. 推理部署](#4) + - [4.1 Python推理](#4-1) + - [4.2 C++推理](#4-2) + - [4.3 Serving服务化部署](#4-3) + - [4.4 更多推理部署](#4-4) +- [5. FAQ](#5) + + +## 1. 算法简介 + +论文信息： +> [Vision Transformer for Fast and Efficient Scene Text Recognition](https://arxiv.org/abs/2105.08582) +> Rowel Atienza +> ICDAR, 2021 + + + +`ViTSTR`使用MJSynth和SynthText两个文字识别数据集训练，在IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE数据集上进行评估，算法复现效果如下： + +|模型|骨干网络|配置文件|Acc|下载链接| +| --- | --- | --- | --- | --- | +|ViTSTR|ViTSTR|[rec_vitstr.yml](../../configs/rec/rec_vitstr.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| + + +## 2. 环境配置 +请先参考[《运行环境准备》](./environment.md)配置PaddleOCR运行环境，参考[《项目克隆》](./clone.md)克隆项目代码。 + + + +## 3. 模型训练、评估、预测 + + +### 3.1 模型训练 + +请参考[文本识别训练教程](./recognition.md)。PaddleOCR对代码进行了模块化，训练`ViTSTR`识别模型时需要**更换配置文件**为`ViTSTR`的[配置文件](../../configs/rec/rec_ViTSTR.yml)。 + +#### 启动训练 + + +具体地，在完成数据准备后，便可以启动训练，训练命令如下： +```shell +#单卡训练（训练周期长，不建议） +python3 tools/train.py -c configs/rec/rec_vitstr.yml + +#多卡训练，通过--gpus参数指定卡号 +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_vitstr.yml +``` + + +### 3.2 评估 + +可下载已训练完成的[模型文件](#model)，使用如下命令进行评估： + +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy +``` + + +### 3.3 预测 + +使用如下命令进行单张图片预测： +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_vitstr_train/best_accuracy +# 预测文件夹下所有图像时，可修改infer_img为文件夹，如 Global.infer_img='./doc/imgs_words_en/'。 +``` + + + +## 4. 推理部署 + + +### 4.1 Python推理 +首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) )，可以使用如下命令进行转换： + +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 tools/export_model.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy Global.save_inference_dir=./inference/rec_vitstr/ +``` +**注意：** +- 如果您是在自己的数据集上训练的模型，并且调整了字典文件，请注意修改配置文件中的`character_dict_path`是否是所需要的字典文件。 +- 如果您修改了训练时的输入大小，请修改`tools/export_model.py`文件中的对应NRTR的`infer_shape`。 + +转换成功后，在目录下有三个文件： +``` +/inference/rec_vitstr/ + ├── inference.pdiparams # 识别inference模型的参数文件 + ├── inference.pdiparams.info # 识别inference模型的参数信息，可忽略 + └── inference.pdmodel # 识别inference模型的program文件 +``` + +执行如下命令进行模型推理： + +```shell +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/'。 +``` + +![](../imgs_words_en/word_10.png) + +执行命令后，上面图像的预测结果（识别的文本和得分）会打印到屏幕上，示例如下： +结果如下： +```shell +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +``` + +**注意**： + +- 训练上述模型采用的图像分辨率是[1，224，224]，需要通过参数`rec_image_shape`设置为您训练时的识别图像形状。 +- 在推理时需要设置参数`rec_char_dict_path`指定字典，如果您修改了字典，请修改该参数为您的字典文件。 +- 如果您修改了预处理方法，需修改`tools/infer/predict_rec.py`中NRTR的预处理为您的预处理方法。 + + + +### 4.2 C++推理部署 + +由于C++预处理后处理还未支持NRTR，所以暂未支持 + + +### 4.3 Serving服务化部署 + +暂不支持 + + +### 4.4 更多推理部署 + +暂不支持 + + +## 5. FAQ + +1. 在`ViTSTR`论文中，使用在ImageNet1k上的预训练权重进行初始化训练，我们在训练未采用预训练权重，最终精度没有变化甚至有所提高。 +2. 我们仅仅复现了`ViTSTR`中的tiny版本，如果有需要使用small、base版本，可直接使用源开源repo中的预训练权重转为Paddle权重即可使用。 + +## 引用 + +```bibtex +@article{Atienza2021ViTSTR, + title = {Vision Transformer for Fast and Efficient Scene Text Recognition}, + author = {Rowel Atienza}, + booktitle = {ICDAR}, + year = {2021}, + url = {https://arxiv.org/abs/2105.08582} +} +``` diff --git a/doc/doc_en/algorithm_overview_en.md b/doc/doc_en/algorithm_overview_en.md index 383cbe39bb..213d95807d 100755 --- a/doc/doc_en/algorithm_overview_en.md +++ b/doc/doc_en/algorithm_overview_en.md @@ -65,6 +65,7 @@ Supported text recognition algorithms (Click the link to get the tutorial): - [x] [SAR](./algorithm_rec_sar_en.md) - [x] [SEED](./algorithm_rec_seed_en.md) - [x] [SVTR](./algorithm_rec_svtr_en.md) +- [x] [ViTSTR](./algorithm_rec_vitstr_en.md) Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation result of these above text recognition (using MJSynth and SynthText for training, evaluate on IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE) is as follow: @@ -83,7 +84,7 @@ Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation r |SAR|Resnet31| 87.20% | rec_r31_sar | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | - +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | diff --git a/doc/doc_en/algorithm_rec_vitstr_en.md b/doc/doc_en/algorithm_rec_vitstr_en.md new file mode 100644 index 0000000000..666798c3e0 --- /dev/null +++ b/doc/doc_en/algorithm_rec_vitstr_en.md @@ -0,0 +1,134 @@ +# ViTSTR + +- [1. Introduction](#1) +- [2. Environment](#2) +- [3. Model Training / Evaluation / Prediction](#3) + - [3.1 Training](#3-1) + - [3.2 Evaluation](#3-2) + - [3.3 Prediction](#3-3) +- [4. Inference and Deployment](#4) + - [4.1 Python Inference](#4-1) + - [4.2 C++ Inference](#4-2) + - [4.3 Serving](#4-3) + - [4.4 More](#4-4) +- [5. FAQ](#5) + + +## 1. Introduction + +Paper: +> [Vision Transformer for Fast and Efficient Scene Text Recognition](https://arxiv.org/abs/2105.08582) +> Rowel Atienza +> ICDAR, 2021 + +Using MJSynth and SynthText two text recognition datasets for training, and evaluating on IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE datasets, the algorithm reproduction effect is as follows: + +|Model|Backbone|config|Acc|Download link| +| --- | --- | --- | --- | --- | +|ViTSTR|ViTSTR|[rec_vitstr.yml](../../configs/rec/rec_vitstr.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| + + +## 2. Environment +Please refer to ["Environment Preparation"](./environment_en.md) to configure the PaddleOCR environment, and refer to ["Project Clone"](./clone_en.md) to clone the project code. + + + +## 3. Model Training / Evaluation / Prediction + +Please refer to [Text Recognition Tutorial](./recognition_en.md). PaddleOCR modularizes the code, and training different recognition models only requires **changing the configuration file**. + +Training: + +Specifically, after the data preparation is completed, the training can be started. The training command is as follows: + +``` +#Single GPU training (long training period, not recommended) +python3 tools/train.py -c configs/rec/rec_vitstr.yml + +#Multi GPU training, specify the gpu number through the --gpus parameter +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_vitstr.yml +``` + +Evaluation: + +``` +# GPU evaluation +python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model={path/to/weights}/best_accuracy +``` + +Prediction: + +``` +# The configuration file used for prediction must match the training +python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_vitstr_train/best_accuracy +``` + + +## 4. Inference and Deployment + + +### 4.1 Python Inference +First, the model saved during the ViTSTR text recognition training process is converted into an inference model. ( [Model download link](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)) ), you can use the following command to convert: + +``` +python3 tools/export_model.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy Global.save_inference_dir=./inference/rec_vitstr +``` + +**Note:** +- If you are training the model on your own dataset and have modified the dictionary file, please pay attention to modify the `character_dict_path` in the configuration file to the modified dictionary file. +- If you modified the input size during training, please modify the `infer_shape` corresponding to ViTSTR in the `tools/export_model.py` file. + +After the conversion is successful, there are three files in the directory: +``` +/inference/rec_vitstr/ + ├── inference.pdiparams + ├── inference.pdiparams.info + └── inference.pdmodel +``` + + +For ViTSTR text recognition model inference, the following commands can be executed: + +``` +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' +``` + +![](../imgs_words_en/word_10.png) + +After executing the command, the prediction result (recognized text and score) of the image above is printed to the screen, an example is as follows: +The result is as follows: +```shell +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +``` + + +### 4.2 C++ Inference + +Not supported + + +### 4.3 Serving + +Not supported + + +### 4.4 More + +Not supported + + +## 5. FAQ + +1. In the `ViTSTR` paper, using pre-trained weights on ImageNet1k for initial training, we did not use pre-trained weights in training, and the final accuracy did not change or even improved. + +## Citation + +```bibtex +@article{Atienza2021ViTSTR, + title = {Vision Transformer for Fast and Efficient Scene Text Recognition}, + author = {Rowel Atienza}, + booktitle = {ICDAR}, + year = {2021}, + url = {https://arxiv.org/abs/2105.08582} +} +``` diff --git a/ppocr/data/imaug/__init__.py b/ppocr/data/imaug/__init__.py index 548832fb0d..2dbc92a703 100644 --- a/ppocr/data/imaug/__init__.py +++ b/ppocr/data/imaug/__init__.py @@ -23,7 +23,7 @@ from .random_crop_data import EastRandomCropData, RandomCropImgMask from .make_pse_gt import MakePseGt from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ - SRNRecResizeImg, NRTRRecResizeImg, SARRecResizeImg, PRENResizeImg + SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg from .ssl_img_aug import SSLRotateResize from .randaugment import RandAugment from .copy_paste import CopyPaste diff --git a/ppocr/data/imaug/label_ops.py b/ppocr/data/imaug/label_ops.py index 02a5187dad..0366e3f785 100644 --- a/ppocr/data/imaug/label_ops.py +++ b/ppocr/data/imaug/label_ops.py @@ -443,7 +443,9 @@ class KieLabelEncode(object): elif 'key_cls' in anno.keys(): labels.append(anno['key_cls']) else: - raise ValueError("Cannot found 'key_cls' in ann.keys(), please check your training annotation.") + raise ValueError( + "Cannot found 'key_cls' in ann.keys(), please check your training annotation." + ) edges.append(ann.get('edge', 0)) ann_infos = dict( image=data['image'], @@ -838,6 +840,37 @@ class PRENLabelEncode(BaseRecLabelEncode): return data +class ViTSTRLabelEncode(BaseRecLabelEncode): + """ Convert between text-label and text-index """ + + def __init__(self, + max_text_length, + character_dict_path=None, + use_space_char=False, + **kwargs): + + super(ViTSTRLabelEncode, self).__init__( + max_text_length, character_dict_path, use_space_char) + + def __call__(self, data): + text = data['label'] + text = self.encode(text) + if text is None: + return None + if len(text) >= self.max_text_len: + return None + data['length'] = np.array(len(text)) + text.insert(0, 0) + text.append(1) + text = text + [0] * (self.max_text_len + 2 - len(text)) + data['label'] = np.array(text) + return data + + def add_special_char(self, dict_character): + dict_character = ['', ''] + dict_character + return dict_character + + class VQATokenLabelEncode(object): """ Label encode for NLP VQA methods diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 7483dffe5b..0697baf436 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -87,11 +87,19 @@ class ClsResizeImg(object): return data -class NRTRRecResizeImg(object): - def __init__(self, image_shape, resize_type, padding=False, **kwargs): +class GrayRecResizeImg(object): + def __init__(self, + image_shape, + resize_type, + inter_type='Image.ANTIALIAS', + scale=True, + padding=False, + **kwargs): self.image_shape = image_shape self.resize_type = resize_type self.padding = padding + self.inter_type = eval(inter_type) + self.scale = scale def __call__(self, data): img = data['image'] @@ -117,13 +125,16 @@ class NRTRRecResizeImg(object): return data if self.resize_type == 'PIL': image_pil = Image.fromarray(np.uint8(img)) - img = image_pil.resize(self.image_shape, Image.ANTIALIAS) + img = image_pil.resize(self.image_shape, self.inter_type) img = np.array(img) if self.resize_type == 'OpenCV': img = cv2.resize(img, self.image_shape) norm_img = np.expand_dims(img, -1) norm_img = norm_img.transpose((2, 0, 1)) - data['image'] = norm_img.astype(np.float32) / 128. - 1. + if self.scale: + data['image'] = norm_img.astype(np.float32) / 128. - 1. + else: + data['image'] = norm_img.astype(np.float32) / 255. return data diff --git a/ppocr/losses/__init__.py b/ppocr/losses/__init__.py index de8419b7c1..6c4545eb21 100755 --- a/ppocr/losses/__init__.py +++ b/ppocr/losses/__init__.py @@ -30,7 +30,7 @@ from .det_fce_loss import FCELoss from .rec_ctc_loss import CTCLoss from .rec_att_loss import AttentionLoss from .rec_srn_loss import SRNLoss -from .rec_nrtr_loss import NRTRLoss +from .rec_ce_smooth_loss import CESmoothingLoss from .rec_sar_loss import SARLoss from .rec_aster_loss import AsterLoss from .rec_pren_loss import PRENLoss @@ -60,8 +60,9 @@ def build_loss(config): support_dict = [ 'DBLoss', 'PSELoss', 'EASTLoss', 'SASTLoss', 'FCELoss', 'CTCLoss', 'ClsLoss', 'AttentionLoss', 'SRNLoss', 'PGLoss', 'CombinedLoss', - 'NRTRLoss', 'TableAttentionLoss', 'SARLoss', 'AsterLoss', 'SDMGRLoss', - 'VQASerTokenLayoutLMLoss', 'LossFromOutput', 'PRENLoss', 'MultiLoss' + 'CESmoothingLoss', 'TableAttentionLoss', 'SARLoss', 'AsterLoss', + 'SDMGRLoss', 'VQASerTokenLayoutLMLoss', 'LossFromOutput', 'PRENLoss', + 'MultiLoss' ] config = copy.deepcopy(config) module_name = config.pop('name') diff --git a/ppocr/losses/rec_nrtr_loss.py b/ppocr/losses/rec_ce_smooth_loss.py similarity index 73% rename from ppocr/losses/rec_nrtr_loss.py rename to ppocr/losses/rec_ce_smooth_loss.py index 200a6d0486..22243ed41f 100644 --- a/ppocr/losses/rec_nrtr_loss.py +++ b/ppocr/losses/rec_ce_smooth_loss.py @@ -3,16 +3,20 @@ from paddle import nn import paddle.nn.functional as F -class NRTRLoss(nn.Layer): - def __init__(self, smoothing=True, **kwargs): - super(NRTRLoss, self).__init__() +class CESmoothingLoss(nn.Layer): + def __init__(self, smoothing=True, with_all=False, **kwargs): + super(CESmoothingLoss, self).__init__() self.loss_func = nn.CrossEntropyLoss(reduction='mean', ignore_index=0) self.smoothing = smoothing + self.with_all = with_all def forward(self, pred, batch): pred = pred.reshape([-1, pred.shape[2]]) - max_len = batch[2].max() - tgt = batch[1][:, 1:2 + max_len] + if self.with_all: + tgt = batch[1] + else: + max_len = batch[2].max() + tgt = batch[1][:, 1:2 + max_len] tgt = tgt.reshape([-1]) if self.smoothing: eps = 0.1 diff --git a/ppocr/modeling/backbones/__init__.py b/ppocr/modeling/backbones/__init__.py index 072d6e0f84..a368e74816 100755 --- a/ppocr/modeling/backbones/__init__.py +++ b/ppocr/modeling/backbones/__init__.py @@ -32,10 +32,11 @@ def build_backbone(config, model_type): from .rec_micronet import MicroNet from .rec_efficientb3_pren import EfficientNetb3_PREN from .rec_svtrnet import SVTRNet + from .rec_vitstr import ViTSTR support_dict = [ 'MobileNetV1Enhance', 'MobileNetV3', 'ResNet', 'ResNetFPN', 'MTB', "ResNet31", "ResNet_ASTER", 'MicroNet', 'EfficientNetb3_PREN', - 'SVTRNet' + 'SVTRNet', 'ViTSTR' ] elif model_type == "e2e": from .e2e_resnet_vd_pg import ResNet diff --git a/ppocr/modeling/backbones/rec_svtrnet.py b/ppocr/modeling/backbones/rec_svtrnet.py index c57bf46345..c2c07f4476 100644 --- a/ppocr/modeling/backbones/rec_svtrnet.py +++ b/ppocr/modeling/backbones/rec_svtrnet.py @@ -147,7 +147,7 @@ class Attention(nn.Layer): dim, num_heads=8, mixer='Global', - HW=[8, 25], + HW=None, local_k=[7, 11], qkv_bias=False, qk_scale=None, @@ -210,7 +210,7 @@ class Block(nn.Layer): num_heads, mixer='Global', local_mixer=[7, 11], - HW=[8, 25], + HW=None, mlp_ratio=4., qkv_bias=False, qk_scale=None, @@ -274,7 +274,9 @@ class PatchEmbed(nn.Layer): img_size=[32, 100], in_channels=3, embed_dim=768, - sub_num=2): + sub_num=2, + patch_size=[4, 4], + mode='pope'): super().__init__() num_patches = (img_size[1] // (2 ** sub_num)) * \ (img_size[0] // (2 ** sub_num)) @@ -282,50 +284,56 @@ class PatchEmbed(nn.Layer): self.num_patches = num_patches self.embed_dim = embed_dim self.norm = None - if sub_num == 2: - self.proj = nn.Sequential( - ConvBNLayer( - in_channels=in_channels, - out_channels=embed_dim // 2, - kernel_size=3, - stride=2, - padding=1, - act=nn.GELU, - bias_attr=None), - ConvBNLayer( - in_channels=embed_dim // 2, - out_channels=embed_dim, - kernel_size=3, - stride=2, - padding=1, - act=nn.GELU, - bias_attr=None)) - if sub_num == 3: - self.proj = nn.Sequential( - ConvBNLayer( - in_channels=in_channels, - out_channels=embed_dim // 4, - kernel_size=3, - stride=2, - padding=1, - act=nn.GELU, - bias_attr=None), - ConvBNLayer( - in_channels=embed_dim // 4, - out_channels=embed_dim // 2, - kernel_size=3, - stride=2, - padding=1, - act=nn.GELU, - bias_attr=None), - ConvBNLayer( - in_channels=embed_dim // 2, - out_channels=embed_dim, - kernel_size=3, - stride=2, - padding=1, - act=nn.GELU, - bias_attr=None)) + if mode == 'pope': + if sub_num == 2: + self.proj = nn.Sequential( + ConvBNLayer( + in_channels=in_channels, + out_channels=embed_dim // 2, + kernel_size=3, + stride=2, + padding=1, + act=nn.GELU, + bias_attr=None), + ConvBNLayer( + in_channels=embed_dim // 2, + out_channels=embed_dim, + kernel_size=3, + stride=2, + padding=1, + act=nn.GELU, + bias_attr=None)) + if sub_num == 3: + self.proj = nn.Sequential( + ConvBNLayer( + in_channels=in_channels, + out_channels=embed_dim // 4, + kernel_size=3, + stride=2, + padding=1, + act=nn.GELU, + bias_attr=None), + ConvBNLayer( + in_channels=embed_dim // 4, + out_channels=embed_dim // 2, + kernel_size=3, + stride=2, + padding=1, + act=nn.GELU, + bias_attr=None), + ConvBNLayer( + in_channels=embed_dim // 2, + out_channels=embed_dim, + kernel_size=3, + stride=2, + padding=1, + act=nn.GELU, + bias_attr=None)) + elif mode == 'linear': + self.proj = nn.Conv2D( + 1, embed_dim, kernel_size=patch_size, stride=patch_size) + self.num_patches = img_size[0] // patch_size[0] * img_size[ + 1] // patch_size[1] def forward(self, x): B, C, H, W = x.shape diff --git a/ppocr/modeling/backbones/rec_vitstr.py b/ppocr/modeling/backbones/rec_vitstr.py new file mode 100644 index 0000000000..d5d7d5148a --- /dev/null +++ b/ppocr/modeling/backbones/rec_vitstr.py @@ -0,0 +1,120 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This code is refer from: +https://github.com/roatienza/deep-text-recognition-benchmark/blob/master/modules/vitstr.py +""" + +import numpy as np +import paddle +import paddle.nn as nn +from ppocr.modeling.backbones.rec_svtrnet import Block, PatchEmbed, zeros_, trunc_normal_, ones_ + +scale_dim_heads = {'tiny': [192, 3], 'small': [384, 6], 'base': [768, 12]} + + +class ViTSTR(nn.Layer): + def __init__(self, + img_size=[224, 224], + in_channels=1, + scale='tiny', + seqlen=27, + patch_size=[16, 16], + embed_dim=None, + depth=12, + num_heads=None, + mlp_ratio=4, + qkv_bias=True, + qk_scale=None, + drop_path_rate=0., + drop_rate=0., + attn_drop_rate=0., + norm_layer='nn.LayerNorm', + act_layer='nn.GELU', + epsilon=1e-6, + out_channels=None, + **kwargs): + super().__init__() + self.seqlen = seqlen + embed_dim = embed_dim if embed_dim is not None else scale_dim_heads[ + scale][0] + num_heads = num_heads if num_heads is not None else scale_dim_heads[ + scale][1] + out_channels = out_channels if out_channels is not None else embed_dim + self.patch_embed = PatchEmbed( + img_size=img_size, + in_channels=in_channels, + embed_dim=embed_dim, + patch_size=patch_size, + mode='linear') + num_patches = self.patch_embed.num_patches + + self.pos_embed = self.create_parameter( + shape=[1, num_patches + 1, embed_dim], default_initializer=zeros_) + self.add_parameter("pos_embed", self.pos_embed) + self.cls_token = self.create_parameter( + shape=[1, 1, embed_dim], default_initializer=zeros_) + self.add_parameter("cls_token", self.cls_token) + + self.pos_drop = nn.Dropout(p=drop_rate) + + dpr = np.linspace(0, drop_path_rate, depth) + self.blocks = nn.LayerList([ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + norm_layer=norm_layer, + act_layer=eval(act_layer), + epsilon=epsilon, + prenorm=False) for i in range(depth) + ]) + self.norm = eval(norm_layer)(embed_dim, epsilon=epsilon) + + self.out_channels = out_channels + + trunc_normal_(self.pos_embed) + trunc_normal_(self.cls_token) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight) + if isinstance(m, nn.Linear) and m.bias is not None: + zeros_(m.bias) + elif isinstance(m, nn.LayerNorm): + zeros_(m.bias) + ones_(m.weight) + + def forward_features(self, x): + B = x.shape[0] + x = self.patch_embed(x) + cls_tokens = paddle.tile(self.cls_token, repeat_times=[B, 1, 1]) + x = paddle.concat((cls_tokens, x), axis=1) + x = x + self.pos_embed + x = self.pos_drop(x) + for blk in self.blocks: + x = blk(x) + x = self.norm(x) + return x + + def forward(self, x): + x = self.forward_features(x) + x = x[:, :self.seqlen] + return x.transpose([0, 2, 1]).unsqueeze(2) diff --git a/ppocr/postprocess/__init__.py b/ppocr/postprocess/__init__.py index f50b5f1c5f..4f900ee1fc 100644 --- a/ppocr/postprocess/__init__.py +++ b/ppocr/postprocess/__init__.py @@ -27,7 +27,7 @@ from .sast_postprocess import SASTPostProcess from .fce_postprocess import FCEPostProcess from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, \ DistillationCTCLabelDecode, TableLabelDecode, NRTRLabelDecode, SARLabelDecode, \ - SEEDLabelDecode, PRENLabelDecode + SEEDLabelDecode, PRENLabelDecode, ViTSTRLabelDecode from .cls_postprocess import ClsPostProcess from .pg_postprocess import PGPostProcess from .vqa_token_ser_layoutlm_postprocess import VQASerTokenLayoutLMPostProcess @@ -42,7 +42,7 @@ def build_post_process(config, global_config=None): 'DistillationDBPostProcess', 'NRTRLabelDecode', 'SARLabelDecode', 'SEEDLabelDecode', 'VQASerTokenLayoutLMPostProcess', 'VQAReTokenLayoutLMPostProcess', 'PRENLabelDecode', - 'DistillationSARLabelDecode' + 'DistillationSARLabelDecode', 'ViTSTRLabelDecode' ] if config['name'] == 'PSEPostProcess': diff --git a/ppocr/postprocess/rec_postprocess.py b/ppocr/postprocess/rec_postprocess.py index bf0fd890bf..df6203fada 100644 --- a/ppocr/postprocess/rec_postprocess.py +++ b/ppocr/postprocess/rec_postprocess.py @@ -188,13 +188,13 @@ class NRTRLabelDecode(BaseRecLabelDecode): char_list = [] conf_list = [] for idx in range(len(text_index[batch_idx])): - if text_index[batch_idx][idx] == 3: # end - break try: - char_list.append(self.character[int(text_index[batch_idx][ - idx])]) + char_idx = self.character[int(text_index[batch_idx][idx])] except: continue + if char_idx == '': # end + break + char_list.append(char_idx) if text_prob is not None: conf_list.append(text_prob[batch_idx][idx]) else: @@ -204,6 +204,32 @@ class NRTRLabelDecode(BaseRecLabelDecode): return result_list +class ViTSTRLabelDecode(NRTRLabelDecode): + """ Convert between text-label and text-index """ + + def __init__(self, character_dict_path=None, use_space_char=False, + **kwargs): + super(ViTSTRLabelDecode, self).__init__(character_dict_path, + use_space_char) + + def __call__(self, preds, label=None, *args, **kwargs): + if isinstance(preds, paddle.Tensor): + preds = preds[:, 1:].numpy() + else: + preds = preds[:, 1:] + preds_idx = preds.argmax(axis=2) + preds_prob = preds.max(axis=2) + text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) + if label is None: + return text + label = self.decode(label[:, 1:]) + return text, label + + def add_special_char(self, dict_character): + dict_character = ['', ''] + dict_character + return dict_character + + class AttnLabelDecode(BaseRecLabelDecode): """ Convert between text-label and text-index """ diff --git a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml index 15119bb2a9..3936ab58ad 100644 --- a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml +++ b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml @@ -49,7 +49,7 @@ Architecture: Loss: - name: NRTRLoss + name: CESmoothingLoss smoothing: True PostProcess: @@ -69,7 +69,7 @@ Train: img_mode: BGR channel_first: False - NRTRLabelEncode: # Class handling label - - NRTRRecResizeImg: + - GrayRecResizeImg: image_shape: [100, 32] resize_type: PIL # PIL or OpenCV - KeepKeys: @@ -90,7 +90,7 @@ Eval: img_mode: BGR channel_first: False - NRTRLabelEncode: # Class handling label - - NRTRRecResizeImg: + - GrayRecResizeImg: image_shape: [100, 32] resize_type: PIL # PIL or OpenCV - KeepKeys: diff --git a/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml new file mode 100644 index 0000000000..26facca34d --- /dev/null +++ b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml @@ -0,0 +1,119 @@ +Global: + use_gpu: True + epoch_num: 20 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/svtr/ + save_epoch_step: 1 + # evaluation is run every 2000 iterations after the 0th iteration + eval_batch_step: [0, 2000] + cal_metric_during_train: True + pretrained_model: + checkpoints: + save_inference_dir: + use_visualdl: False + infer_img: doc/imgs_words_en/word_10.png + # for data or label process + character_dict_path: + character_type: en + max_text_length: 25 + infer_mode: False + use_space_char: False + save_res_path: ./output/rec/predicts_svtr_tiny.txt + + +Optimizer: + name: AdamW + beta1: 0.9 + beta2: 0.99 + epsilon: 0.00000008 + weight_decay: 0.05 + no_weight_decay_name: norm pos_embed + one_dim_param_no_weight_decay: true + lr: + name: Cosine + learning_rate: 0.0005 + warmup_epoch: 2 + +Architecture: + model_type: rec + algorithm: SVTR + Transform: + name: STN_ON + tps_inputsize: [32, 64] + tps_outputsize: [32, 100] + num_control_points: 20 + tps_margins: [0.05,0.05] + stn_activation: none + Backbone: + name: SVTRNet + img_size: [32, 100] + out_char_num: 25 + out_channels: 192 + patch_merging: 'Conv' + embed_dim: [64, 128, 256] + depth: [3, 6, 3] + num_heads: [2, 4, 8] + mixer: ['Local','Local','Local','Local','Local','Local','Global','Global','Global','Global','Global','Global'] + local_mixer: [[7, 11], [7, 11], [7, 11]] + last_stage: True + prenorm: false + Neck: + name: SequenceEncoder + encoder_type: reshape + Head: + name: CTCHead + +Loss: + name: CTCLoss + +PostProcess: + name: CTCLabelDecode + +Metric: + name: RecMetric + main_indicator: acc + +Train: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data/ + label_file_list: ["./train_data/ic15_data/rec_gt_train.txt"] + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - CTCLabelEncode: # Class handling label + - RecResizeImg: + character_dict_path: + image_shape: [3, 64, 256] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: True + batch_size_per_card: 512 + drop_last: True + num_workers: 4 + +Eval: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data + label_file_list: ["./train_data/ic15_data/rec_gt_test.txt"] + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - CTCLabelEncode: # Class handling label + - RecResizeImg: + character_dict_path: + image_shape: [3, 64, 256] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: False + drop_last: False + batch_size_per_card: 256 + num_workers: 2 diff --git a/test_tipc/configs/rec_svtrnet/train_infer_python.txt b/test_tipc/configs/rec_svtrnet/train_infer_python.txt new file mode 100644 index 0000000000..72526063e0 --- /dev/null +++ b/test_tipc/configs/rec_svtrnet/train_infer_python.txt @@ -0,0 +1,53 @@ +===========================train_params=========================== +model_name:rec_svtrnet +python:python3.7 +gpu_list:0|0,1 +Global.use_gpu:True|True +Global.auto_cast:null +Global.epoch_num:lite_train_lite_infer=2|whole_train_whole_infer=300 +Global.save_model_dir:./output/ +Train.loader.batch_size_per_card:lite_train_lite_infer=16|whole_train_whole_infer=64 +Global.pretrained_model:null +train_model_name:latest +train_infer_img_dir:./inference/rec_inference +null:null +## +trainer:norm_train +norm_train:tools/train.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o +pact_train:null +fpgm_train:null +distill_train:null +null:null +null:null +## +===========================eval_params=========================== +eval:tools/eval.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o +null:null +## +===========================infer_params=========================== +Global.save_inference_dir:./output/ +Global.checkpoints: +norm_export:tools/export_model.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o +quant_export:null +fpgm_export:null +distill_export:null +export1:null +export2:null +## +train_model:./inference/rec_svtrnet_train/best_accuracy +infer_export:tools/export_model.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o +infer_quant:False +inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="SVTR" +--use_gpu:True|False +--enable_mkldnn:False +--cpu_threads:1|6 +--rec_batch_num:1|6 +--use_tensorrt:False +--precision:fp32|int8 +--rec_model_dir: +--image_dir:./inference/rec_inference +--save_log_path:./test/output/ +--benchmark:True +null:null +===========================infer_benchmark_params========================== +random_infer_input:[{float32,[1,224,224]}] diff --git a/test_tipc/configs/rec_vitstr/rec_vitstr.yml b/test_tipc/configs/rec_vitstr/rec_vitstr.yml new file mode 100644 index 0000000000..427bce4b5a --- /dev/null +++ b/test_tipc/configs/rec_vitstr/rec_vitstr.yml @@ -0,0 +1,101 @@ +Global: + use_gpu: True + epoch_num: 20 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/vitstr/ + save_epoch_step: 1 + # evaluation is run every 2000 iterations after the 0th iteration# + eval_batch_step: [0, 2000] + cal_metric_during_train: True + pretrained_model: + checkpoints: + save_inference_dir: + use_visualdl: False + infer_img: doc/imgs_words_en/word_10.png + # for data or label process + character_dict_path: ppocr/utils/EN_symbol_dict.txt + max_text_length: 25 + infer_mode: False + use_space_char: False + save_res_path: ./output/rec/predicts_vitstr.txt + + +Optimizer: + name: Adadelta + epsilon: 0.00000001 + rho: 0.95 + clip_norm: 5.0 + lr: + learning_rate: 1.0 + +Architecture: + model_type: rec + algorithm: ViTSTR + in_channels: 1 + Transform: + Backbone: + name: ViTSTR + Neck: + name: SequenceEncoder + encoder_type: reshape + Head: + name: CTCHead + +Loss: + name: CESmoothingLoss + smoothing: False + with_all: True + +PostProcess: + name: ViTSTRLabelDecode + +Metric: + name: RecMetric + main_indicator: acc + +Train: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data/ + label_file_list: ["./train_data/ic15_data/rec_gt_train.txt"] + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - ViTSTRLabelEncode: # Class handling label + - GrayRecResizeImg: + image_shape: [224, 224] # W H + resize_type: PIL # PIL or OpenCV + inter_type: 'Image.BICUBIC' + scale: false + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: True + batch_size_per_card: 48 + drop_last: True + num_workers: 2 + +Eval: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data + label_file_list: ["./train_data/ic15_data/rec_gt_test.txt"] + transforms: + - DecodeImage: # load image + img_mode: BGR + channel_first: False + - ViTSTRLabelEncode: # Class handling label + - GrayRecResizeImg: + image_shape: [224, 224] # W H + resize_type: PIL # PIL or OpenCV + inter_type: 'Image.BICUBIC' + scale: false + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: False + drop_last: False + batch_size_per_card: 256 + num_workers: 2 diff --git a/test_tipc/configs/rec_vitstr/train_infer_python.txt b/test_tipc/configs/rec_vitstr/train_infer_python.txt new file mode 100644 index 0000000000..6c7af1fb87 --- /dev/null +++ b/test_tipc/configs/rec_vitstr/train_infer_python.txt @@ -0,0 +1,53 @@ +===========================train_params=========================== +model_name:rec_vitstr +python:python3.7 +gpu_list:0|0,1 +Global.use_gpu:True|True +Global.auto_cast:null +Global.epoch_num:lite_train_lite_infer=2|whole_train_whole_infer=300 +Global.save_model_dir:./output/ +Train.loader.batch_size_per_card:lite_train_lite_infer=16|whole_train_whole_infer=64 +Global.pretrained_model:null +train_model_name:latest +train_infer_img_dir:./inference/rec_inference +null:null +## +trainer:norm_train +norm_train:tools/train.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +pact_train:null +fpgm_train:null +distill_train:null +null:null +null:null +## +===========================eval_params=========================== +eval:tools/eval.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +null:null +## +===========================infer_params=========================== +Global.save_inference_dir:./output/ +Global.checkpoints: +norm_export:tools/export_model.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +quant_export:null +fpgm_export:null +distill_export:null +export1:null +export2:null +## +train_model:./inference/rec_vitstr_train/best_accuracy +infer_export:tools/export_model.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +infer_quant:False +inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="ViTSTR" +--use_gpu:True|False +--enable_mkldnn:False +--cpu_threads:1|6 +--rec_batch_num:1|6 +--use_tensorrt:False +--precision:fp32|int8 +--rec_model_dir: +--image_dir:./inference/rec_inference +--save_log_path:./test/output/ +--benchmark:True +null:null +===========================infer_benchmark_params========================== +random_infer_input:[{float32,[1,224,224]}] diff --git a/tools/export_model.py b/tools/export_model.py index c0cbcd361c..6e003f2ffa 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -73,6 +73,12 @@ def export_single_model(model, arch_config, save_path, logger, quanter=None): shape=[None, 3, 64, 512], dtype="float32"), ] model = to_static(model, input_spec=other_shape) + elif arch_config["algorithm"] == "ViTSTR": + other_shape = [ + paddle.static.InputSpec( + shape=[None, 1, 224, 224], dtype="float32"), + ] + model = to_static(model, input_spec=other_shape) else: infer_shape = [3, -1, -1] if arch_config["model_type"] == "rec": diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index 3664ef2caf..1945667972 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -69,6 +69,12 @@ class TextRecognizer(object): "character_dict_path": args.rec_char_dict_path, "use_space_char": args.use_space_char } + elif self.rec_algorithm == 'ViTSTR': + postprocess_params = { + 'name': 'ViTSTRLabelDecode', + "character_dict_path": args.rec_char_dict_path, + "use_space_char": args.use_space_char + } self.postprocess_op = build_post_process(postprocess_params) self.predictor, self.input_tensor, self.output_tensors, self.config = \ utility.create_predictor(args, 'rec', logger) @@ -96,15 +102,22 @@ class TextRecognizer(object): def resize_norm_img(self, img, max_wh_ratio): imgC, imgH, imgW = self.rec_image_shape - if self.rec_algorithm == 'NRTR': + if self.rec_algorithm == 'NRTR' or self.rec_algorithm == 'ViTSTR': img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # return padding_im image_pil = Image.fromarray(np.uint8(img)) - img = image_pil.resize([100, 32], Image.ANTIALIAS) + if self.rec_algorithm == 'ViTSTR': + img = image_pil.resize([imgW, imgH], Image.BICUBIC) + else: + img = image_pil.resize([imgW, imgH], Image.ANTIALIAS) img = np.array(img) norm_img = np.expand_dims(img, -1) norm_img = norm_img.transpose((2, 0, 1)) - return norm_img.astype(np.float32) / 128. - 1. + if self.rec_algorithm == 'ViTSTR': + norm_img = norm_img.astype(np.float32) / 255. + else: + norm_img = norm_img.astype(np.float32) / 128. - 1. + return norm_img assert imgC == img.shape[2] imgW = int((imgH * max_wh_ratio)) diff --git a/tools/program.py b/tools/program.py index aa0d2698cf..745c28b872 100755 --- a/tools/program.py +++ b/tools/program.py @@ -307,7 +307,8 @@ def train(config, train_stats.update(stats) if log_writer is not None and dist.get_rank() == 0: - log_writer.log_metrics(metrics=train_stats.get(), prefix="TRAIN", step=global_step) + log_writer.log_metrics( + metrics=train_stats.get(), prefix="TRAIN", step=global_step) if dist.get_rank() == 0 and ( (global_step > 0 and global_step % print_batch_step == 0) or @@ -354,7 +355,8 @@ def train(config, # logger metric if log_writer is not None: - log_writer.log_metrics(metrics=cur_metric, prefix="EVAL", step=global_step) + log_writer.log_metrics( + metrics=cur_metric, prefix="EVAL", step=global_step) if cur_metric[main_indicator] >= best_model_dict[ main_indicator]: @@ -377,11 +379,18 @@ def train(config, logger.info(best_str) # logger best metric if log_writer is not None: - log_writer.log_metrics(metrics={ - "best_{}".format(main_indicator): best_model_dict[main_indicator] - }, prefix="EVAL", step=global_step) - - log_writer.log_model(is_best=True, prefix="best_accuracy", metadata=best_model_dict) + log_writer.log_metrics( + metrics={ + "best_{}".format(main_indicator): + best_model_dict[main_indicator] + }, + prefix="EVAL", + step=global_step) + + log_writer.log_model( + is_best=True, + prefix="best_accuracy", + metadata=best_model_dict) reader_start = time.time() if dist.get_rank() == 0: @@ -413,7 +422,8 @@ def train(config, epoch=epoch, global_step=global_step) if log_writer is not None: - log_writer.log_model(is_best=False, prefix='iter_epoch_{}'.format(epoch)) + log_writer.log_model( + is_best=False, prefix='iter_epoch_{}'.format(epoch)) best_str = 'best metric, {}'.format(', '.join( ['{}: {}'.format(k, v) for k, v in best_model_dict.items()])) @@ -564,7 +574,8 @@ def preprocess(is_train=False): assert alg in [ 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', - 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR' + 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', + 'ViTSTR' ] if use_xpu: @@ -585,7 +596,8 @@ def preprocess(is_train=False): vdl_writer_path = '{}/vdl/'.format(save_model_dir) log_writer = VDLLogger(save_model_dir) loggers.append(log_writer) - if ('use_wandb' in config['Global'] and config['Global']['use_wandb']) or 'wandb' in config: + if ('use_wandb' in config['Global'] and + config['Global']['use_wandb']) or 'wandb' in config: save_dir = config['Global']['save_model_dir'] wandb_writer_path = "{}/wandb".format(save_dir) if "wandb" in config: From c956f4a79c229c907ba60a7d4bf87024a4e6dc76 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Fri, 27 May 2022 07:50:48 +0000 Subject: [PATCH 02/43] fix cpu_thread and precision --- test_tipc/configs/rec_svtrnet/train_infer_python.txt | 4 ++-- test_tipc/configs/rec_vitstr/train_infer_python.txt | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test_tipc/configs/rec_svtrnet/train_infer_python.txt b/test_tipc/configs/rec_svtrnet/train_infer_python.txt index 72526063e0..df7638faad 100644 --- a/test_tipc/configs/rec_svtrnet/train_infer_python.txt +++ b/test_tipc/configs/rec_svtrnet/train_infer_python.txt @@ -40,10 +40,10 @@ infer_quant:False inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="SVTR" --use_gpu:True|False --enable_mkldnn:False ---cpu_threads:1|6 +--cpu_threads:6 --rec_batch_num:1|6 --use_tensorrt:False ---precision:fp32|int8 +--precision:fp32 --rec_model_dir: --image_dir:./inference/rec_inference --save_log_path:./test/output/ diff --git a/test_tipc/configs/rec_vitstr/train_infer_python.txt b/test_tipc/configs/rec_vitstr/train_infer_python.txt index 6c7af1fb87..8b58e220f7 100644 --- a/test_tipc/configs/rec_vitstr/train_infer_python.txt +++ b/test_tipc/configs/rec_vitstr/train_infer_python.txt @@ -40,10 +40,10 @@ infer_quant:False inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="ViTSTR" --use_gpu:True|False --enable_mkldnn:False ---cpu_threads:1|6 +--cpu_threads:6 --rec_batch_num:1|6 --use_tensorrt:False ---precision:fp32|int8 +--precision:fp32 --rec_model_dir: --image_dir:./inference/rec_inference --save_log_path:./test/output/ From bc5522b6f38dc1119bd3af06f985ee80bc9bda9d Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Fri, 27 May 2022 08:29:32 +0000 Subject: [PATCH 03/43] fix svtr tipc --- test_tipc/configs/rec_svtrnet/train_infer_python.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test_tipc/configs/rec_svtrnet/train_infer_python.txt b/test_tipc/configs/rec_svtrnet/train_infer_python.txt index df7638faad..60fd25ad4f 100644 --- a/test_tipc/configs/rec_svtrnet/train_infer_python.txt +++ b/test_tipc/configs/rec_svtrnet/train_infer_python.txt @@ -37,7 +37,7 @@ export2:null train_model:./inference/rec_svtrnet_train/best_accuracy infer_export:tools/export_model.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o infer_quant:False -inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="SVTR" +inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="3,64,256" --rec_algorithm="SVTR" --use_gpu:True|False --enable_mkldnn:False --cpu_threads:6 @@ -50,4 +50,4 @@ inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbo --benchmark:True null:null ===========================infer_benchmark_params========================== -random_infer_input:[{float32,[1,224,224]}] +random_infer_input:[{float32,[1=3,64,256]}] From c269e1e42f79d33aa87d5e0ba49b488c56f97ae4 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Mon, 30 May 2022 08:36:11 +0000 Subject: [PATCH 04/43] modify vitstr name --- configs/rec/rec_mtb_nrtr.yml | 2 +- ...{rec_vitstr.yml => rec_vitstr_none_ce.yml} | 4 +-- doc/doc_ch/algorithm_overview.md | 2 +- doc/doc_ch/algorithm_rec_vitstr.md | 26 +++++++++---------- doc/doc_en/algorithm_overview_en.md | 2 +- doc/doc_en/algorithm_rec_vitstr_en.md | 14 +++++----- ppocr/losses/__init__.py | 7 +++-- .../{rec_ce_smooth_loss.py => rec_ce_loss.py} | 4 +-- .../configs/rec_mtb_nrtr/rec_mtb_nrtr.yml | 2 +- .../rec_mtb_nrtr/train_infer_python.txt | 4 +-- .../rec_vitstr_none_ce.yml} | 4 +-- .../train_infer_python.txt | 10 +++---- 12 files changed, 40 insertions(+), 41 deletions(-) rename configs/rec/{rec_vitstr.yml => rec_vitstr_none_ce.yml} (97%) rename ppocr/losses/{rec_ce_smooth_loss.py => rec_ce_loss.py} (93%) rename test_tipc/configs/{rec_vitstr/rec_vitstr.yml => rec_vitstr_none_ce/rec_vitstr_none_ce.yml} (97%) rename test_tipc/configs/{rec_vitstr => rec_vitstr_none_ce}/train_infer_python.txt (82%) diff --git a/configs/rec/rec_mtb_nrtr.yml b/configs/rec/rec_mtb_nrtr.yml index dfe2cc9811..405674882c 100644 --- a/configs/rec/rec_mtb_nrtr.yml +++ b/configs/rec/rec_mtb_nrtr.yml @@ -49,7 +49,7 @@ Architecture: Loss: - name: CESmoothingLoss + name: CELoss smoothing: True PostProcess: diff --git a/configs/rec/rec_vitstr.yml b/configs/rec/rec_vitstr_none_ce.yml similarity index 97% rename from configs/rec/rec_vitstr.yml rename to configs/rec/rec_vitstr_none_ce.yml index 005db0184a..065cc785f8 100644 --- a/configs/rec/rec_vitstr.yml +++ b/configs/rec/rec_vitstr_none_ce.yml @@ -3,7 +3,7 @@ Global: epoch_num: 20 log_smooth_window: 20 print_batch_step: 10 - save_model_dir: ./output/rec/vitstr/ + save_model_dir: ./output/rec/vitstr_none_ce/ save_epoch_step: 1 # evaluation is run every 2000 iterations after the 0th iteration# eval_batch_step: [0, 50] @@ -44,7 +44,7 @@ Architecture: name: CTCHead Loss: - name: CESmoothingLoss + name: CELoss smoothing: False with_all: True diff --git a/doc/doc_ch/algorithm_overview.md b/doc/doc_ch/algorithm_overview.md index 934ac08537..4196ec05af 100755 --- a/doc/doc_ch/algorithm_overview.md +++ b/doc/doc_ch/algorithm_overview.md @@ -85,7 +85,7 @@ |SAR|Resnet31| 87.20% | rec_r31_sar | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | -|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar) | diff --git a/doc/doc_ch/algorithm_rec_vitstr.md b/doc/doc_ch/algorithm_rec_vitstr.md index bd83b8d9c2..ab12be7202 100644 --- a/doc/doc_ch/algorithm_rec_vitstr.md +++ b/doc/doc_ch/algorithm_rec_vitstr.md @@ -27,7 +27,7 @@ |模型|骨干网络|配置文件|Acc|下载链接| | --- | --- | --- | --- | --- | -|ViTSTR|ViTSTR|[rec_vitstr.yml](../../configs/rec/rec_vitstr.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| +|ViTSTR|ViTSTR|[rec_vitstr_none_ce.yml](../../configs/rec/rec_vitstr_none_ce.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar)| ## 2. 环境配置 @@ -40,7 +40,7 @@ ### 3.1 模型训练 -请参考[文本识别训练教程](./recognition.md)。PaddleOCR对代码进行了模块化，训练`ViTSTR`识别模型时需要**更换配置文件**为`ViTSTR`的[配置文件](../../configs/rec/rec_ViTSTR.yml)。 +请参考[文本识别训练教程](./recognition.md)。PaddleOCR对代码进行了模块化，训练`ViTSTR`识别模型时需要**更换配置文件**为`ViTSTR`的[配置文件](../../configs/rec/rec_vitstr_none_ce.yml)。 #### 启动训练 @@ -48,10 +48,10 @@ 具体地，在完成数据准备后，便可以启动训练，训练命令如下： ```shell #单卡训练（训练周期长，不建议） -python3 tools/train.py -c configs/rec/rec_vitstr.yml +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.yml +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_vitstr_none_ce.yml ``` @@ -61,7 +61,7 @@ python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs ```shell # 注意将pretrained_model的路径设置为本地路径。 -python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy +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 ``` @@ -70,7 +70,7 @@ python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec 使用如下命令进行单张图片预测： ```shell # 注意将pretrained_model的路径设置为本地路径。 -python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_vitstr_train/best_accuracy +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/'。 ``` @@ -80,15 +80,15 @@ python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./ ### 4.1 Python推理 -首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) )，可以使用如下命令进行转换： +首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar) )，可以使用如下命令进行转换： ```shell # 注意将pretrained_model的路径设置为本地路径。 -python3 tools/export_model.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy Global.save_inference_dir=./inference/rec_vitstr/ +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`文件中的对应NRTR的`infer_shape`。 +- 如果您修改了训练时的输入大小，请修改`tools/export_model.py`文件中的对应ViTSTR的`infer_shape`。 转换成功后，在目录下有三个文件： ``` @@ -110,20 +110,20 @@ python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' 执行命令后，上面图像的预测结果（识别的文本和得分）会打印到屏幕上，示例如下： 结果如下： ```shell -Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +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`中NRTR的预处理为您的预处理方法。 +- 如果您修改了预处理方法，需修改`tools/infer/predict_rec.py`中ViTSTR的预处理为您的预处理方法。 ### 4.2 C++推理部署 -由于C++预处理后处理还未支持NRTR，所以暂未支持 +由于C++预处理后处理还未支持ViTSTR，所以暂未支持 ### 4.3 Serving服务化部署 @@ -139,7 +139,7 @@ Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) ## 5. FAQ 1. 在`ViTSTR`论文中，使用在ImageNet1k上的预训练权重进行初始化训练，我们在训练未采用预训练权重，最终精度没有变化甚至有所提高。 -2. 我们仅仅复现了`ViTSTR`中的tiny版本，如果有需要使用small、base版本，可直接使用源开源repo中的预训练权重转为Paddle权重即可使用。 +2. 我们仅仅复现了`ViTSTR`中的tiny版本，如果需要使用small、base版本，可将[ViTSTR源repo](https://github.com/roatienza/deep-text-recognition-benchmark) 中的预训练权重转为Paddle权重使用。 ## 引用 diff --git a/doc/doc_en/algorithm_overview_en.md b/doc/doc_en/algorithm_overview_en.md index 213d95807d..f2284542ed 100755 --- a/doc/doc_en/algorithm_overview_en.md +++ b/doc/doc_en/algorithm_overview_en.md @@ -84,7 +84,7 @@ Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation r |SAR|Resnet31| 87.20% | rec_r31_sar | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | -|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | diff --git a/doc/doc_en/algorithm_rec_vitstr_en.md b/doc/doc_en/algorithm_rec_vitstr_en.md index 666798c3e0..b6f26d39e2 100644 --- a/doc/doc_en/algorithm_rec_vitstr_en.md +++ b/doc/doc_en/algorithm_rec_vitstr_en.md @@ -25,7 +25,7 @@ Using MJSynth and SynthText two text recognition datasets for training, and eval |Model|Backbone|config|Acc|Download link| | --- | --- | --- | --- | --- | -|ViTSTR|ViTSTR|[rec_vitstr.yml](../../configs/rec/rec_vitstr.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| +|ViTSTR|ViTSTR|[rec_vitstr_none_ce.yml](../../configs/rec/rec_vitstr_none_ce.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| ## 2. Environment @@ -43,24 +43,24 @@ Specifically, after the data preparation is completed, the training can be start ``` #Single GPU training (long training period, not recommended) -python3 tools/train.py -c configs/rec/rec_vitstr.yml +python3 tools/train.py -c configs/rec/rec_vitstr_none_ce.yml #Multi GPU training, specify the gpu number through the --gpus parameter -python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_vitstr.yml +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_vitstr_none_ce.yml ``` Evaluation: ``` # GPU evaluation -python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model={path/to/weights}/best_accuracy +python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_vitstr_none_ce.yml -o Global.pretrained_model={path/to/weights}/best_accuracy ``` Prediction: ``` # The configuration file used for prediction must match the training -python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_vitstr_train/best_accuracy +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 ``` @@ -71,7 +71,7 @@ python3 tools/infer_rec.py -c configs/rec/rec_vitstr.yml -o Global.infer_img='./ First, the model saved during the ViTSTR text recognition training process is converted into an inference model. ( [Model download link](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)) ), you can use the following command to convert: ``` -python3 tools/export_model.py -c configs/rec/rec_vitstr.yml -o Global.pretrained_model=./rec_vitstr_train/best_accuracy Global.save_inference_dir=./inference/rec_vitstr +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 ``` **Note:** @@ -98,7 +98,7 @@ python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' After executing the command, the prediction result (recognized text and score) of the image above is printed to the screen, an example is as follows: The result is as follows: ```shell -Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9998350143432617) ``` diff --git a/ppocr/losses/__init__.py b/ppocr/losses/__init__.py index 6c4545eb21..7bea87f62f 100755 --- a/ppocr/losses/__init__.py +++ b/ppocr/losses/__init__.py @@ -30,7 +30,7 @@ from .det_fce_loss import FCELoss from .rec_ctc_loss import CTCLoss from .rec_att_loss import AttentionLoss from .rec_srn_loss import SRNLoss -from .rec_ce_smooth_loss import CESmoothingLoss +from .rec_ce_loss import CELoss from .rec_sar_loss import SARLoss from .rec_aster_loss import AsterLoss from .rec_pren_loss import PRENLoss @@ -60,9 +60,8 @@ def build_loss(config): support_dict = [ 'DBLoss', 'PSELoss', 'EASTLoss', 'SASTLoss', 'FCELoss', 'CTCLoss', 'ClsLoss', 'AttentionLoss', 'SRNLoss', 'PGLoss', 'CombinedLoss', - 'CESmoothingLoss', 'TableAttentionLoss', 'SARLoss', 'AsterLoss', - 'SDMGRLoss', 'VQASerTokenLayoutLMLoss', 'LossFromOutput', 'PRENLoss', - 'MultiLoss' + 'CELoss', 'TableAttentionLoss', 'SARLoss', 'AsterLoss', 'SDMGRLoss', + 'VQASerTokenLayoutLMLoss', 'LossFromOutput', 'PRENLoss', 'MultiLoss' ] config = copy.deepcopy(config) module_name = config.pop('name') diff --git a/ppocr/losses/rec_ce_smooth_loss.py b/ppocr/losses/rec_ce_loss.py similarity index 93% rename from ppocr/losses/rec_ce_smooth_loss.py rename to ppocr/losses/rec_ce_loss.py index 22243ed41f..b837ac27bb 100644 --- a/ppocr/losses/rec_ce_smooth_loss.py +++ b/ppocr/losses/rec_ce_loss.py @@ -3,9 +3,9 @@ from paddle import nn import paddle.nn.functional as F -class CESmoothingLoss(nn.Layer): +class CELoss(nn.Layer): def __init__(self, smoothing=True, with_all=False, **kwargs): - super(CESmoothingLoss, self).__init__() + super(CELoss, self).__init__() self.loss_func = nn.CrossEntropyLoss(reduction='mean', ignore_index=0) self.smoothing = smoothing self.with_all = with_all diff --git a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml index 3936ab58ad..ba6728d836 100644 --- a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml +++ b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml @@ -49,7 +49,7 @@ Architecture: Loss: - name: CESmoothingLoss + name: CELoss smoothing: True PostProcess: diff --git a/test_tipc/configs/rec_mtb_nrtr/train_infer_python.txt b/test_tipc/configs/rec_mtb_nrtr/train_infer_python.txt index de6de5a0ca..20d984e38b 100644 --- a/test_tipc/configs/rec_mtb_nrtr/train_infer_python.txt +++ b/test_tipc/configs/rec_mtb_nrtr/train_infer_python.txt @@ -40,10 +40,10 @@ infer_quant:False inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,32,100" --rec_algorithm="NRTR" --use_gpu:True|False --enable_mkldnn:True|False ---cpu_threads:1|6 +--cpu_threads:6 --rec_batch_num:1|6 --use_tensorrt:True|False ---precision:fp32|int8 +--precision:fp32 --rec_model_dir: --image_dir:./inference/rec_inference --save_log_path:./test/output/ diff --git a/test_tipc/configs/rec_vitstr/rec_vitstr.yml b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml similarity index 97% rename from test_tipc/configs/rec_vitstr/rec_vitstr.yml rename to test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml index 427bce4b5a..ccd93c9eaf 100644 --- a/test_tipc/configs/rec_vitstr/rec_vitstr.yml +++ b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml @@ -3,7 +3,7 @@ Global: epoch_num: 20 log_smooth_window: 20 print_batch_step: 10 - save_model_dir: ./output/rec/vitstr/ + save_model_dir: ./output/rec/vitstr_none_ce/ save_epoch_step: 1 # evaluation is run every 2000 iterations after the 0th iteration# eval_batch_step: [0, 2000] @@ -43,7 +43,7 @@ Architecture: name: CTCHead Loss: - name: CESmoothingLoss + name: CELoss smoothing: False with_all: True diff --git a/test_tipc/configs/rec_vitstr/train_infer_python.txt b/test_tipc/configs/rec_vitstr_none_ce/train_infer_python.txt similarity index 82% rename from test_tipc/configs/rec_vitstr/train_infer_python.txt rename to test_tipc/configs/rec_vitstr_none_ce/train_infer_python.txt index 8b58e220f7..04c5742ea2 100644 --- a/test_tipc/configs/rec_vitstr/train_infer_python.txt +++ b/test_tipc/configs/rec_vitstr_none_ce/train_infer_python.txt @@ -13,7 +13,7 @@ train_infer_img_dir:./inference/rec_inference null:null ## trainer:norm_train -norm_train:tools/train.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +norm_train:tools/train.py -c test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml -o pact_train:null fpgm_train:null distill_train:null @@ -21,21 +21,21 @@ null:null null:null ## ===========================eval_params=========================== -eval:tools/eval.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +eval:tools/eval.py -c test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml -o null:null ## ===========================infer_params=========================== Global.save_inference_dir:./output/ Global.checkpoints: -norm_export:tools/export_model.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +norm_export:tools/export_model.py -c test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml -o quant_export:null fpgm_export:null distill_export:null export1:null export2:null ## -train_model:./inference/rec_vitstr_train/best_accuracy -infer_export:tools/export_model.py -c test_tipc/configs/rec_vitstr/rec_vitstr.yml -o +train_model:./inference/rec_vitstr_none_ce_train/best_accuracy +infer_export:tools/export_model.py -c test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml -o infer_quant:False inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="1,224,224" --rec_algorithm="ViTSTR" --use_gpu:True|False From 275c5a2ba93d0791e7b8009e3a595f4ac7554989 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Mon, 30 May 2022 08:38:07 +0000 Subject: [PATCH 05/43] modify vitstr config batchsize --- configs/rec/rec_vitstr_none_ce.yml | 2 +- test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/configs/rec/rec_vitstr_none_ce.yml b/configs/rec/rec_vitstr_none_ce.yml index 065cc785f8..c71eba0cc2 100644 --- a/configs/rec/rec_vitstr_none_ce.yml +++ b/configs/rec/rec_vitstr_none_ce.yml @@ -75,7 +75,7 @@ Train: shuffle: True batch_size_per_card: 48 drop_last: True - num_workers: 2 + num_workers: 8 Eval: dataset: diff --git a/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml index ccd93c9eaf..4e8151f599 100644 --- a/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml +++ b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml @@ -75,7 +75,7 @@ Train: shuffle: True batch_size_per_card: 48 drop_last: True - num_workers: 2 + num_workers: 8 Eval: dataset: From c503dc2f9352272615dc3cc11737b833036c6ccc Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Wed, 22 Jun 2022 14:19:55 +0000 Subject: [PATCH 06/43] [New Rec] add vitstr and ABINet --- configs/rec/rec_mtb_nrtr.yml | 6 +- configs/rec/rec_r45_abinet.yml | 103 ++++ configs/rec/rec_svtrnet.yml | 2 +- configs/rec/rec_vitstr_none_ce.yml | 12 +- doc/doc_ch/algorithm_overview.md | 2 + doc/doc_ch/algorithm_rec_abinet.md | 155 ++++++ doc/doc_ch/algorithm_rec_nrtr.md | 140 ++++- doc/doc_en/algorithm_overview_en.md | 2 + doc/doc_en/algorithm_rec_abinet_en.md | 136 +++++ doc/doc_en/algorithm_rec_nrtr_en.md | 141 ++++- doc/doc_en/algorithm_rec_vitstr_en.md | 2 +- ppocr/data/imaug/__init__.py | 2 +- ppocr/data/imaug/label_ops.py | 158 +++--- ppocr/data/imaug/operators.py | 33 -- ppocr/data/imaug/rec_img_aug.py | 38 ++ ppocr/losses/rec_ce_loss.py | 76 ++- ppocr/modeling/backbones/__init__.py | 18 +- ppocr/modeling/backbones/rec_resnet.py | 280 ++++++++++ ppocr/modeling/heads/__init__.py | 3 +- ppocr/modeling/heads/multiheadAttention.py | 163 ------ ppocr/modeling/heads/rec_abinet_head.py | 296 +++++++++++ ppocr/modeling/heads/rec_nrtr_head.py | 489 ++++++------------ ppocr/postprocess/__init__.py | 4 +- ppocr/postprocess/rec_postprocess.py | 209 ++++---- .../configs/rec_mtb_nrtr/rec_mtb_nrtr.yml | 2 +- .../configs/rec_r45_abinet/rec_r45_abinet.yml | 105 ++++ .../rec_r45_abinet/train_infer_python.txt | 53 ++ test_tipc/configs/rec_svtrnet/rec_svtrnet.yml | 2 +- .../rec_svtrnet/train_infer_python.txt | 4 +- .../rec_vitstr_none_ce/rec_vitstr_none_ce.yml | 5 +- tools/export_model.py | 15 +- tools/infer/predict_rec.py | 51 +- tools/program.py | 2 +- 33 files changed, 1973 insertions(+), 736 deletions(-) create mode 100644 configs/rec/rec_r45_abinet.yml create mode 100644 doc/doc_ch/algorithm_rec_abinet.md create mode 100644 doc/doc_en/algorithm_rec_abinet_en.md create mode 100644 ppocr/modeling/backbones/rec_resnet.py delete mode 100755 ppocr/modeling/heads/multiheadAttention.py create mode 100644 ppocr/modeling/heads/rec_abinet_head.py create mode 100644 test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml create mode 100644 test_tipc/configs/rec_r45_abinet/train_infer_python.txt diff --git a/configs/rec/rec_mtb_nrtr.yml b/configs/rec/rec_mtb_nrtr.yml index 405674882c..4e5826adc9 100644 --- a/configs/rec/rec_mtb_nrtr.yml +++ b/configs/rec/rec_mtb_nrtr.yml @@ -9,7 +9,7 @@ Global: eval_batch_step: [0, 2000] cal_metric_during_train: True pretrained_model: - checkpoints: + checkpoints: save_inference_dir: use_visualdl: False infer_img: doc/imgs_words_en/word_10.png @@ -82,7 +82,7 @@ Train: Eval: dataset: name: LMDBDataSet - data_dir: ./train_data/data_lmdb_release/evaluation/ + data_dir: ./train_data/data_lmdb_release/validation/ transforms: - DecodeImage: # load image img_mode: BGR @@ -97,5 +97,5 @@ Eval: shuffle: False drop_last: False batch_size_per_card: 256 - num_workers: 1 + num_workers: 4 use_shared_memory: False diff --git a/configs/rec/rec_r45_abinet.yml b/configs/rec/rec_r45_abinet.yml new file mode 100644 index 0000000000..a756fead9f --- /dev/null +++ b/configs/rec/rec_r45_abinet.yml @@ -0,0 +1,103 @@ +Global: + use_gpu: True + epoch_num: 10 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/r45_abinet/ + save_epoch_step: 1 + # evaluation is run every 2000 iterations + eval_batch_step: [0, 2000] + cal_metric_during_train: True + pretrained_model: + checkpoints: + save_inference_dir: + use_visualdl: False + infer_img: doc/imgs_words_en/word_10.png + # for data or label process + character_dict_path: + character_type: en + max_text_length: 25 + infer_mode: False + use_space_char: False + save_res_path: ./output/rec/predicts_abinet.txt + +Optimizer: + name: Adam + beta1: 0.9 + beta2: 0.99 + clip_norm: 20.0 + lr: + name: Piecewise + decay_epochs: [6] + values: [0.0001, 0.00001] + regularizer: + name: 'L2' + factor: 0. + +Architecture: + model_type: rec + algorithm: ABINet + in_channels: 3 + Transform: + Backbone: + name: ResNet45 + + Head: + name: ABINetHead + use_lang: True + iter_size: 3 + + +Loss: + name: CELoss + ignore_index: &ignore_index 100 # Must be greater than the number of character classes + +PostProcess: + name: ABINetLabelDecode + +Metric: + name: RecMetric + main_indicator: acc + +Train: + dataset: + name: LMDBDataSet + data_dir: ./train_data/data_lmdb_release/training/ + transforms: + - DecodeImage: # load image + img_mode: RGB + channel_first: False + - ABINetLabelEncode: # Class handling label + ignore_index: *ignore_index + - ABINetRecResizeImg: + image_shape: [3, 32, 128] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: True + batch_size_per_card: 96 + drop_last: True + num_workers: 4 + +Eval: + dataset: + name: LMDBDataSet + data_dir: ./train_data/data_lmdb_release/validation/ + transforms: + - DecodeImage: # load image + img_mode: RGB + channel_first: False + - ABINetLabelEncode: # Class handling label + ignore_index: *ignore_index + - ABINetRecResizeImg: + image_shape: [3, 32, 128] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: False + drop_last: False + batch_size_per_card: 256 + num_workers: 4 + use_shared_memory: False diff --git a/configs/rec/rec_svtrnet.yml b/configs/rec/rec_svtrnet.yml index a700e4bd92..a3d292b6ba 100644 --- a/configs/rec/rec_svtrnet.yml +++ b/configs/rec/rec_svtrnet.yml @@ -26,7 +26,7 @@ Optimizer: name: AdamW beta1: 0.9 beta2: 0.99 - epsilon: 0.00000008 + epsilon: 8.e-8 weight_decay: 0.05 no_weight_decay_name: norm pos_embed one_dim_param_no_weight_decay: true diff --git a/configs/rec/rec_vitstr_none_ce.yml b/configs/rec/rec_vitstr_none_ce.yml index c71eba0cc2..b969c83a5d 100644 --- a/configs/rec/rec_vitstr_none_ce.yml +++ b/configs/rec/rec_vitstr_none_ce.yml @@ -6,7 +6,7 @@ Global: save_model_dir: ./output/rec/vitstr_none_ce/ save_epoch_step: 1 # evaluation is run every 2000 iterations after the 0th iteration# - eval_batch_step: [0, 50] + eval_batch_step: [0, 2000] cal_metric_during_train: True pretrained_model: checkpoints: @@ -23,7 +23,7 @@ Global: Optimizer: name: Adadelta - epsilon: 0.00000001 + epsilon: 1.e-8 rho: 0.95 clip_norm: 5.0 lr: @@ -45,8 +45,8 @@ Architecture: Loss: name: CELoss - smoothing: False with_all: True + ignore_index: &ignore_index 0 # Must be zero or greater than the number of character classes PostProcess: name: ViTSTRLabelDecode @@ -58,12 +58,13 @@ Metric: Train: dataset: name: LMDBDataSet - data_dir: ./train_data/data_lmdb_release/training + data_dir: ./train_data/data_lmdb_release/training/ transforms: - DecodeImage: # load image img_mode: BGR channel_first: False - ViTSTRLabelEncode: # Class handling label + ignore_index: *ignore_index - GrayRecResizeImg: image_shape: [224, 224] # W H resize_type: PIL # PIL or OpenCV @@ -80,12 +81,13 @@ Train: Eval: dataset: name: LMDBDataSet - data_dir: ./train_data/data_lmdb_release/validation + data_dir: ./train_data/data_lmdb_release/validation/ transforms: - DecodeImage: # load image img_mode: BGR channel_first: False - ViTSTRLabelEncode: # Class handling label + ignore_index: *ignore_index - GrayRecResizeImg: image_shape: [224, 224] # W H resize_type: PIL # PIL or OpenCV diff --git a/doc/doc_ch/algorithm_overview.md b/doc/doc_ch/algorithm_overview.md index 4196ec05af..eb81e4cd6d 100755 --- a/doc/doc_ch/algorithm_overview.md +++ b/doc/doc_ch/algorithm_overview.md @@ -67,6 +67,7 @@ - [x] [SEED](./algorithm_rec_seed.md) - [x] [SVTR](./algorithm_rec_svtr.md) - [x] [ViTSTR](./algorithm_rec_vitstr.md) +- [x] [ABINet](./algorithm_rec_abinet.md) 参考[DTRB](https://arxiv.org/abs/1904.01906)[3]文字识别训练和评估流程，使用MJSynth和SynthText两个文字识别数据集训练，在IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE数据集上进行评估，算法效果如下： @@ -86,6 +87,7 @@ |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | |ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar) | +|ABINet|Resnet45| 90.75% | rec_r45_abinet_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | diff --git a/doc/doc_ch/algorithm_rec_abinet.md b/doc/doc_ch/algorithm_rec_abinet.md new file mode 100644 index 0000000000..d20c703014 --- /dev/null +++ b/doc/doc_ch/algorithm_rec_abinet.md @@ -0,0 +1,155 @@ +# 场景文本识别算法-ABINet + +- [1. 算法简介](#1) +- [2. 环境配置](#2) +- [3. 模型训练、评估、预测](#3) + - [3.1 训练](#3-1) + - [3.2 评估](#3-2) + - [3.3 预测](#3-3) +- [4. 推理部署](#4) + - [4.1 Python推理](#4-1) + - [4.2 C++推理](#4-2) + - [4.3 Serving服务化部署](#4-3) + - [4.4 更多推理部署](#4-4) +- [5. FAQ](#5) + + +## 1. 算法简介 + +论文信息： +> [ABINet: Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition](https://openaccess.thecvf.com/content/CVPR2021/papers/Fang_Read_Like_Humans_Autonomous_Bidirectional_and_Iterative_Language_Modeling_for_CVPR_2021_paper.pdf) +> Shancheng Fang and Hongtao Xie and Yuxin Wang and Zhendong Mao and Yongdong Zhang +> CVPR, 2021 + + + +`ABINet`使用MJSynth和SynthText两个文字识别数据集训练，在IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE数据集上进行评估，算法复现效果如下： + +|模型|骨干网络|配置文件|Acc|下载链接| +| --- | --- | --- | --- | --- | +|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[训练模型]()/[预训练模型]| + + +## 2. 环境配置 +请先参考[《运行环境准备》](./environment.md)配置PaddleOCR运行环境，参考[《项目克隆》](./clone.md)克隆项目代码。 + + + +## 3. 模型训练、评估、预测 + + +### 3.1 模型训练 + +请参考[文本识别训练教程](./recognition.md)。PaddleOCR对代码进行了模块化，训练`ABINet`识别模型时需要**更换配置文件**为`ABINet`的[配置文件](../../configs/rec/rec_r45_abinet.yml)。 + +#### 启动训练 + + +具体地，在完成数据准备后，便可以启动训练，训练命令如下： +```shell +#单卡训练（训练周期长，不建议） +python3 tools/train.py -c configs/rec/rec_r45_abinet.yml + +#多卡训练，通过--gpus参数指定卡号 +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_r45_abinet.yml +``` + + +### 3.2 评估 + +可下载已训练完成的[模型文件](#model)，使用如下命令进行评估： + +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_r45_abinet.yml -o Global.pretrained_model=./rec_r45_abinet_train/best_accuracy +``` + + +### 3.3 预测 + +使用如下命令进行单张图片预测： +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 tools/infer_rec.py -c configs/rec/rec_r45_abinet.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_r45_abinet_train/best_accuracy +# 预测文件夹下所有图像时，可修改infer_img为文件夹，如 Global.infer_img='./doc/imgs_words_en/'。 +``` + + + +## 4. 推理部署 + + +### 4.1 Python推理 +首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址]() )，可以使用如下命令进行转换： + +```shell +# 注意将pretrained_model的路径设置为本地路径。 +python3 tools/export_model.py -c configs/rec/rec_r45_abinet.yml -o Global.pretrained_model=./rec_r45_abinet_train/best_accuracy Global.save_inference_dir=./inference/rec_r45_abinet/ +``` +**注意：** +- 如果您是在自己的数据集上训练的模型，并且调整了字典文件，请注意修改配置文件中的`character_dict_path`是否是所需要的字典文件。 +- 如果您修改了训练时的输入大小，请修改`tools/export_model.py`文件中的对应ABINet的`infer_shape`。 + +转换成功后，在目录下有三个文件： +``` +/inference/rec_r45_abinet/ + ├── inference.pdiparams # 识别inference模型的参数文件 + ├── inference.pdiparams.info # 识别inference模型的参数信息，可忽略 + └── inference.pdmodel # 识别inference模型的program文件 +``` + +执行如下命令进行模型推理： + +```shell +python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' --rec_model_dir='./inference/rec_r45_abinet/' --rec_algorithm='ABINet' --rec_image_shape='3,32,128' --rec_char_dict_path='./ppocr/utils/ic15_dict.txt' +# 预测文件夹下所有图像时，可修改image_dir为文件夹，如 --image_dir='./doc/imgs_words_en/'。 +``` + +![](../imgs_words_en/word_10.png) + +执行命令后，上面图像的预测结果（识别的文本和得分）会打印到屏幕上，示例如下： +结果如下： +```shell +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9999995231628418) +``` + +**注意**： + +- 训练上述模型采用的图像分辨率是[3，32，128]，需要通过参数`rec_image_shape`设置为您训练时的识别图像形状。 +- 在推理时需要设置参数`rec_char_dict_path`指定字典，如果您修改了字典，请修改该参数为您的字典文件。 +- 如果您修改了预处理方法，需修改`tools/infer/predict_rec.py`中ABINet的预处理为您的预处理方法。 + + + +### 4.2 C++推理部署 + +由于C++预处理后处理还未支持ABINet，所以暂未支持 + + +### 4.3 Serving服务化部署 + +暂不支持 + + +### 4.4 更多推理部署 + +暂不支持 + + +## 5. FAQ + +1. MJSynth和SynthText两种数据集来自于[ABINet源repo](https://github.com/FangShancheng/ABINet) 。 +2. 我们使用ABINet作者提供的预训练模型进行finetune训练。 + +## 引用 + +```bibtex +@article{Fang2021ABINet, + title = {ABINet: Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition}, + author = {Shancheng Fang and Hongtao Xie and Yuxin Wang and Zhendong Mao and Yongdong Zhang}, + booktitle = {CVPR}, + year = {2021}, + url = {https://arxiv.org/abs/2103.06495}, + pages = {7098-7107} +} +``` diff --git a/doc/doc_ch/algorithm_rec_nrtr.md b/doc/doc_ch/algorithm_rec_nrtr.md index d3b626d024..c619ac1dbc 100644 --- a/doc/doc_ch/algorithm_rec_nrtr.md +++ b/doc/doc_ch/algorithm_rec_nrtr.md @@ -12,6 +12,7 @@ - [4.3 Serving服务化部署](#4-3) - [4.4 更多推理部署](#4-4) - [5. FAQ](#5) +- [6. 发行公告](#6) ## 1. 算法简介 @@ -110,7 +111,7 @@ python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' 执行命令后，上面图像的预测结果（识别的文本和得分）会打印到屏幕上，示例如下： 结果如下： ```shell -Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9465042352676392) ``` **注意**： @@ -140,12 +141,147 @@ Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) 1. `NRTR`论文中使用Beam搜索进行解码字符，但是速度较慢，这里默认未使用Beam搜索，以贪婪搜索进行解码字符。 + +## 6. 发行公告 + +1. release/2.6更新NRTR代码结构，新版NRTR可加载旧版(release/2.5及之前)模型参数，使用下面示例代码将旧版模型参数转换为新版模型参数： + +```python + + params = paddle.load('path/' + '.pdparams') # 旧版本参数 + state_dict = model.state_dict() # 新版模型参数 + new_state_dict = {} + + for k1, v1 in state_dict.items(): + + k = k1 + if 'encoder' in k and 'self_attn' in k and 'qkv' in k and 'weight' in k: + + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')].transpose((1, 0, 2, 3)) + k = params[k_para.replace('qkv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('qkv', 'conv3')].transpose((1, 0, 2, 3)) + + new_state_dict[k1] = np.concatenate([q[:, :, 0, 0], k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'encoder' in k and 'self_attn' in k and 'qkv' in k and 'bias' in k: + + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')] + k = params[k_para.replace('qkv', 'conv2')] + v = params[k_para.replace('qkv', 'conv3')] + + new_state_dict[k1] = np.concatenate([q, k, v], -1) + + elif 'encoder' in k and 'self_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + elif 'encoder' in k and 'norm3' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para.replace('norm3', 'norm2')] + + elif 'encoder' in k and 'norm1' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + + elif 'decoder' in k and 'self_attn' in k and 'qkv' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')].transpose((1, 0, 2, 3)) + k = params[k_para.replace('qkv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('qkv', 'conv3')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = np.concatenate([q[:, :, 0, 0], k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'decoder' in k and 'self_attn' in k and 'qkv' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')] + k = params[k_para.replace('qkv', 'conv2')] + v = params[k_para.replace('qkv', 'conv3')] + new_state_dict[k1] = np.concatenate([q, k, v], -1) + + elif 'decoder' in k and 'self_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + elif 'decoder' in k and 'cross_attn' in k and 'q' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + q = params[k_para.replace('q', 'conv1')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = q[:, :, 0, 0] + + elif 'decoder' in k and 'cross_attn' in k and 'q' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + q = params[k_para.replace('q', 'conv1')] + new_state_dict[k1] = q + + elif 'decoder' in k and 'cross_attn' in k and 'kv' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + k = params[k_para.replace('kv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('kv', 'conv3')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = np.concatenate([k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'decoder' in k and 'cross_attn' in k and 'kv' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + k = params[k_para.replace('kv', 'conv2')] + v = params[k_para.replace('kv', 'conv3')] + new_state_dict[k1] = np.concatenate([k, v], -1) + + elif 'decoder' in k and 'cross_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + new_state_dict[k1] = params[k_para] + elif 'decoder' in k and 'norm' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + elif 'mlp' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('fc', 'conv') + k_para = k_para.replace('mlp.', '') + w = params[k_para].transpose((1, 0, 2, 3)) + new_state_dict[k1] = w[:, :, 0, 0] + elif 'mlp' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('fc', 'conv') + k_para = k_para.replace('mlp.', '') + w = params[k_para] + new_state_dict[k1] = w + + else: + new_state_dict[k1] = params[k1] + + if list(new_state_dict[k1].shape) != list(v1.shape): + print(k1) + + + for k, v1 in state_dict.items(): + if k not in new_state_dict.keys(): + print(1, k) + elif list(new_state_dict[k].shape) != list(v1.shape): + print(2, k) + + + + model.set_state_dict(new_state_dict) + paddle.save(model.state_dict(), 'nrtrnew_from_old_params.pdparams') + +``` + +2. 新版相比与旧版，代码结构简洁，推理速度有所提高。 + + ## 引用 ```bibtex @article{Sheng2019NRTR, title = {NRTR: A No-Recurrence Sequence-to-Sequence Model For Scene Text Recognition}, - author = {Fenfen Sheng and Zhineng Chen andBo Xu}, + author = {Fenfen Sheng and Zhineng Chen and Bo Xu}, booktitle = {ICDAR}, year = {2019}, url = {http://arxiv.org/abs/1806.00926}, diff --git a/doc/doc_en/algorithm_overview_en.md b/doc/doc_en/algorithm_overview_en.md index f2284542ed..28aca7c0d1 100755 --- a/doc/doc_en/algorithm_overview_en.md +++ b/doc/doc_en/algorithm_overview_en.md @@ -66,6 +66,7 @@ Supported text recognition algorithms (Click the link to get the tutorial): - [x] [SEED](./algorithm_rec_seed_en.md) - [x] [SVTR](./algorithm_rec_svtr_en.md) - [x] [ViTSTR](./algorithm_rec_vitstr_en.md) +- [x] [ABINet](./algorithm_rec_abinet_en.md) Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation result of these above text recognition (using MJSynth and SynthText for training, evaluate on IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE) is as follow: @@ -85,6 +86,7 @@ Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation r |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | |ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | +|ABINet|Resnet45| 90.75% | rec_r45_abinet_en | [trained model](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | diff --git a/doc/doc_en/algorithm_rec_abinet_en.md b/doc/doc_en/algorithm_rec_abinet_en.md new file mode 100644 index 0000000000..3b0f6c09a9 --- /dev/null +++ b/doc/doc_en/algorithm_rec_abinet_en.md @@ -0,0 +1,136 @@ +# ABINet + +- [1. Introduction](#1) +- [2. Environment](#2) +- [3. Model Training / Evaluation / Prediction](#3) + - [3.1 Training](#3-1) + - [3.2 Evaluation](#3-2) + - [3.3 Prediction](#3-3) +- [4. Inference and Deployment](#4) + - [4.1 Python Inference](#4-1) + - [4.2 C++ Inference](#4-2) + - [4.3 Serving](#4-3) + - [4.4 More](#4-4) +- [5. FAQ](#5) + + +## 1. Introduction + +Paper: +> [ABINet: Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition](https://openaccess.thecvf.com/content/CVPR2021/papers/Fang_Read_Like_Humans_Autonomous_Bidirectional_and_Iterative_Language_Modeling_for_CVPR_2021_paper.pdf) +> Shancheng Fang and Hongtao Xie and Yuxin Wang and Zhendong Mao and Yongdong Zhang +> CVPR, 2021 + +Using MJSynth and SynthText two text recognition datasets for training, and evaluating on IIIT, SVT, IC03, IC13, IC15, SVTP, CUTE datasets, the algorithm reproduction effect is as follows: + +|Model|Backbone|config|Acc|Download link| +| --- | --- | --- | --- | --- | +|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[trained model]()/[pretrained model]()| + + +## 2. Environment +Please refer to ["Environment Preparation"](./environment_en.md) to configure the PaddleOCR environment, and refer to ["Project Clone"](./clone_en.md) to clone the project code. + + + +## 3. Model Training / Evaluation / Prediction + +Please refer to [Text Recognition Tutorial](./recognition_en.md). PaddleOCR modularizes the code, and training different recognition models only requires **changing the configuration file**. + +Training: + +Specifically, after the data preparation is completed, the training can be started. The training command is as follows: + +``` +#Single GPU training (long training period, not recommended) +python3 tools/train.py -c configs/rec/rec_r45_abinet.yml + +#Multi GPU training, specify the gpu number through the --gpus parameter +python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py -c configs/rec/rec_r45_abinet.yml +``` + +Evaluation: + +``` +# GPU evaluation +python3 -m paddle.distributed.launch --gpus '0' tools/eval.py -c configs/rec/rec_r45_abinet.yml -o Global.pretrained_model={path/to/weights}/best_accuracy +``` + +Prediction: + +``` +# The configuration file used for prediction must match the training +python3 tools/infer_rec.py -c configs/rec/rec_r45_abinet.yml -o Global.infer_img='./doc/imgs_words_en/word_10.png' Global.pretrained_model=./rec_r45_abinet_train/best_accuracy +``` + + +## 4. Inference and Deployment + + +### 4.1 Python Inference +First, the model saved during the ABINet text recognition training process is converted into an inference model. ( [Model download link]()) ), you can use the following command to convert: + +``` +python3 tools/export_model.py -c configs/rec/rec_r45_abinet.yml -o Global.pretrained_model=./rec_r45_abinet_train/best_accuracy Global.save_inference_dir=./inference/rec_r45_abinet +``` + +**Note:** +- If you are training the model on your own dataset and have modified the dictionary file, please pay attention to modify the `character_dict_path` in the configuration file to the modified dictionary file. +- If you modified the input size during training, please modify the `infer_shape` corresponding to ABINet in the `tools/export_model.py` file. + +After the conversion is successful, there are three files in the directory: +``` +/inference/rec_r45_abinet/ + ├── inference.pdiparams + ├── inference.pdiparams.info + └── inference.pdmodel +``` + + +For ABINet text recognition model inference, the following commands can be executed: + +``` +python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' --rec_model_dir='./inference/rec_r45_abinet/' --rec_algorithm='ABINet' --rec_image_shape='3,32,128' --rec_char_dict_path='./ppocr/utils/ic15_dict.txt' +``` + +![](../imgs_words_en/word_10.png) + +After executing the command, the prediction result (recognized text and score) of the image above is printed to the screen, an example is as follows: +The result is as follows: +```shell +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9999995231628418) +``` + + +### 4.2 C++ Inference + +Not supported + + +### 4.3 Serving + +Not supported + + +### 4.4 More + +Not supported + + +## 5. FAQ + +1. Note that the MJSynth and SynthText datasets come from [ABINet repo](https://github.com/FangShancheng/ABINet). +2. We use the pre-trained model provided by the ABINet authors for finetune training. + +## Citation + +```bibtex +@article{Fang2021ABINet, + title = {ABINet: Read Like Humans: Autonomous, Bidirectional and Iterative Language Modeling for Scene Text Recognition}, + author = {Shancheng Fang and Hongtao Xie and Yuxin Wang and Zhendong Mao and Yongdong Zhang}, + booktitle = {CVPR}, + year = {2021}, + url = {https://arxiv.org/abs/2103.06495}, + pages = {7098-7107} +} +``` diff --git a/doc/doc_en/algorithm_rec_nrtr_en.md b/doc/doc_en/algorithm_rec_nrtr_en.md index 40c9b91629..309d7ab123 100644 --- a/doc/doc_en/algorithm_rec_nrtr_en.md +++ b/doc/doc_en/algorithm_rec_nrtr_en.md @@ -12,6 +12,7 @@ - [4.3 Serving](#4-3) - [4.4 More](#4-4) - [5. FAQ](#5) +- [6. Release Note](#6) ## 1. Introduction @@ -25,7 +26,7 @@ Using MJSynth and SynthText two text recognition datasets for training, and eval |Model|Backbone|config|Acc|Download link| | --- | --- | --- | --- | --- | -|NRTR|MTB|[rec_mtb_nrtr.yml](../../configs/rec/rec_mtb_nrtr.yml)|84.21%|[train model](https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/rec_mtb_nrtr_train.tar)| +|NRTR|MTB|[rec_mtb_nrtr.yml](../../configs/rec/rec_mtb_nrtr.yml)|84.21%|[trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/rec_mtb_nrtr_train.tar)| ## 2. Environment @@ -98,7 +99,7 @@ python3 tools/infer/predict_rec.py --image_dir='./doc/imgs_words_en/word_10.png' After executing the command, the prediction result (recognized text and score) of the image above is printed to the screen, an example is as follows: The result is as follows: ```shell -Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9265879392623901) +Predicts of ./doc/imgs_words_en/word_10.png:('pain', 0.9465042352676392) ``` @@ -121,12 +122,146 @@ Not supported 1. In the `NRTR` paper, Beam search is used to decode characters, but the speed is slow. Beam search is not used by default here, and greedy search is used to decode characters. + +## 6. Release Note + +1. The release/2.6 version updates the NRTR code structure. The new version of NRTR can load the model parameters of the old version (release/2.5 and before), and you may use the following code to convert the old version model parameters to the new version model parameters: + +```python + + params = paddle.load('path/' + '.pdparams') # the old version parameters + state_dict = model.state_dict() # the new version model parameters + new_state_dict = {} + + for k1, v1 in state_dict.items(): + + k = k1 + if 'encoder' in k and 'self_attn' in k and 'qkv' in k and 'weight' in k: + + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')].transpose((1, 0, 2, 3)) + k = params[k_para.replace('qkv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('qkv', 'conv3')].transpose((1, 0, 2, 3)) + + new_state_dict[k1] = np.concatenate([q[:, :, 0, 0], k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'encoder' in k and 'self_attn' in k and 'qkv' in k and 'bias' in k: + + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')] + k = params[k_para.replace('qkv', 'conv2')] + v = params[k_para.replace('qkv', 'conv3')] + + new_state_dict[k1] = np.concatenate([q, k, v], -1) + + elif 'encoder' in k and 'self_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + elif 'encoder' in k and 'norm3' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para.replace('norm3', 'norm2')] + + elif 'encoder' in k and 'norm1' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + + elif 'decoder' in k and 'self_attn' in k and 'qkv' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')].transpose((1, 0, 2, 3)) + k = params[k_para.replace('qkv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('qkv', 'conv3')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = np.concatenate([q[:, :, 0, 0], k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'decoder' in k and 'self_attn' in k and 'qkv' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + q = params[k_para.replace('qkv', 'conv1')] + k = params[k_para.replace('qkv', 'conv2')] + v = params[k_para.replace('qkv', 'conv3')] + new_state_dict[k1] = np.concatenate([q, k, v], -1) + + elif 'decoder' in k and 'self_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + + elif 'decoder' in k and 'cross_attn' in k and 'q' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + q = params[k_para.replace('q', 'conv1')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = q[:, :, 0, 0] + + elif 'decoder' in k and 'cross_attn' in k and 'q' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + q = params[k_para.replace('q', 'conv1')] + new_state_dict[k1] = q + + elif 'decoder' in k and 'cross_attn' in k and 'kv' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + k = params[k_para.replace('kv', 'conv2')].transpose((1, 0, 2, 3)) + v = params[k_para.replace('kv', 'conv3')].transpose((1, 0, 2, 3)) + new_state_dict[k1] = np.concatenate([k[:, :, 0, 0], v[:, :, 0, 0]], -1) + + elif 'decoder' in k and 'cross_attn' in k and 'kv' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + k = params[k_para.replace('kv', 'conv2')] + v = params[k_para.replace('kv', 'conv3')] + new_state_dict[k1] = np.concatenate([k, v], -1) + + elif 'decoder' in k and 'cross_attn' in k and 'out_proj' in k: + + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('cross_attn', 'multihead_attn') + new_state_dict[k1] = params[k_para] + elif 'decoder' in k and 'norm' in k: + k_para = k[:13] + 'layers.' + k[13:] + new_state_dict[k1] = params[k_para] + elif 'mlp' in k and 'weight' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('fc', 'conv') + k_para = k_para.replace('mlp.', '') + w = params[k_para].transpose((1, 0, 2, 3)) + new_state_dict[k1] = w[:, :, 0, 0] + elif 'mlp' in k and 'bias' in k: + k_para = k[:13] + 'layers.' + k[13:] + k_para = k_para.replace('fc', 'conv') + k_para = k_para.replace('mlp.', '') + w = params[k_para] + new_state_dict[k1] = w + + else: + new_state_dict[k1] = params[k1] + + if list(new_state_dict[k1].shape) != list(v1.shape): + print(k1) + + + for k, v1 in state_dict.items(): + if k not in new_state_dict.keys(): + print(1, k) + elif list(new_state_dict[k].shape) != list(v1.shape): + print(2, k) + + + + model.set_state_dict(new_state_dict) + paddle.save(model.state_dict(), 'nrtrnew_from_old_params.pdparams') + +``` + +2. The new version has a clean code structure and improved inference speed compared with the old version. + ## Citation ```bibtex @article{Sheng2019NRTR, title = {NRTR: A No-Recurrence Sequence-to-Sequence Model For Scene Text Recognition}, - author = {Fenfen Sheng and Zhineng Chen andBo Xu}, + author = {Fenfen Sheng and Zhineng Chen and Bo Xu}, booktitle = {ICDAR}, year = {2019}, url = {http://arxiv.org/abs/1806.00926}, diff --git a/doc/doc_en/algorithm_rec_vitstr_en.md b/doc/doc_en/algorithm_rec_vitstr_en.md index b6f26d39e2..a6f9e2f15d 100644 --- a/doc/doc_en/algorithm_rec_vitstr_en.md +++ b/doc/doc_en/algorithm_rec_vitstr_en.md @@ -25,7 +25,7 @@ Using MJSynth and SynthText two text recognition datasets for training, and eval |Model|Backbone|config|Acc|Download link| | --- | --- | --- | --- | --- | -|ViTSTR|ViTSTR|[rec_vitstr_none_ce.yml](../../configs/rec/rec_vitstr_none_ce.yml)|79.82%|[训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| +|ViTSTR|ViTSTR|[rec_vitstr_none_ce.yml](../../configs/rec/rec_vitstr_none_ce.yml)|79.82%|[trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar)| ## 2. Environment diff --git a/ppocr/data/imaug/__init__.py b/ppocr/data/imaug/__init__.py index 88c9603fb0..58e8a5c780 100644 --- a/ppocr/data/imaug/__init__.py +++ b/ppocr/data/imaug/__init__.py @@ -24,7 +24,7 @@ from .make_pse_gt import MakePseGt from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ - SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg + SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, ABINetRecResizeImg from .ssl_img_aug import SSLRotateResize from .randaugment import RandAugment from .copy_paste import CopyPaste diff --git a/ppocr/data/imaug/label_ops.py b/ppocr/data/imaug/label_ops.py index 0366e3f785..d96ecf639a 100644 --- a/ppocr/data/imaug/label_ops.py +++ b/ppocr/data/imaug/label_ops.py @@ -157,37 +157,6 @@ class BaseRecLabelEncode(object): return text_list -class NRTRLabelEncode(BaseRecLabelEncode): - """ Convert between text-label and text-index """ - - def __init__(self, - max_text_length, - character_dict_path=None, - use_space_char=False, - **kwargs): - - super(NRTRLabelEncode, self).__init__( - max_text_length, character_dict_path, use_space_char) - - def __call__(self, data): - text = data['label'] - text = self.encode(text) - if text is None: - return None - if len(text) >= self.max_text_len - 1: - return None - data['length'] = np.array(len(text)) - text.insert(0, 2) - text.append(3) - text = text + [0] * (self.max_text_len - len(text)) - data['label'] = np.array(text) - return data - - def add_special_char(self, dict_character): - dict_character = ['blank', '', '', ''] + dict_character - return dict_character - - class CTCLabelEncode(BaseRecLabelEncode): """ Convert between text-label and text-index """ @@ -840,37 +809,6 @@ class PRENLabelEncode(BaseRecLabelEncode): return data -class ViTSTRLabelEncode(BaseRecLabelEncode): - """ Convert between text-label and text-index """ - - def __init__(self, - max_text_length, - character_dict_path=None, - use_space_char=False, - **kwargs): - - super(ViTSTRLabelEncode, self).__init__( - max_text_length, character_dict_path, use_space_char) - - def __call__(self, data): - text = data['label'] - text = self.encode(text) - if text is None: - return None - if len(text) >= self.max_text_len: - return None - data['length'] = np.array(len(text)) - text.insert(0, 0) - text.append(1) - text = text + [0] * (self.max_text_len + 2 - len(text)) - data['label'] = np.array(text) - return data - - def add_special_char(self, dict_character): - dict_character = ['', ''] + dict_character - return dict_character - - class VQATokenLabelEncode(object): """ Label encode for NLP VQA methods @@ -1077,3 +1015,99 @@ class MultiLabelEncode(BaseRecLabelEncode): data_out['label_sar'] = sar['label'] data_out['length'] = ctc['length'] return data_out + + +class NRTRLabelEncode(BaseRecLabelEncode): + """ Convert between text-label and text-index """ + + def __init__(self, + max_text_length, + character_dict_path=None, + use_space_char=False, + **kwargs): + + super(NRTRLabelEncode, self).__init__( + max_text_length, character_dict_path, use_space_char) + + def __call__(self, data): + text = data['label'] + text = self.encode(text) + if text is None: + return None + if len(text) >= self.max_text_len - 1: + return None + data['length'] = np.array(len(text)) + text.insert(0, 2) + text.append(3) + text = text + [0] * (self.max_text_len - len(text)) + data['label'] = np.array(text) + return data + + def add_special_char(self, dict_character): + dict_character = ['blank', '', '', ''] + dict_character + return dict_character + + +class ViTSTRLabelEncode(BaseRecLabelEncode): + """ Convert between text-label and text-index """ + + def __init__(self, + max_text_length, + character_dict_path=None, + use_space_char=False, + ignore_index=0, + **kwargs): + + super(ViTSTRLabelEncode, self).__init__( + max_text_length, character_dict_path, use_space_char) + self.ignore_index = ignore_index + + def __call__(self, data): + text = data['label'] + text = self.encode(text) + if text is None: + return None + if len(text) >= self.max_text_len: + return None + data['length'] = np.array(len(text)) + text.insert(0, self.ignore_index) + text.append(1) + text = text + [self.ignore_index] * (self.max_text_len + 2 - len(text)) + data['label'] = np.array(text) + return data + + def add_special_char(self, dict_character): + dict_character = ['', ''] + dict_character + return dict_character + + +class ABINetLabelEncode(BaseRecLabelEncode): + """ Convert between text-label and text-index """ + + def __init__(self, + max_text_length, + character_dict_path=None, + use_space_char=False, + ignore_index=100, + **kwargs): + + super(ABINetLabelEncode, self).__init__( + max_text_length, character_dict_path, use_space_char) + self.ignore_index = ignore_index + + def __call__(self, data): + text = data['label'] + text = self.encode(text) + if text is None: + return None + if len(text) >= self.max_text_len: + return None + data['length'] = np.array(len(text)) + text.append(0) + text = text + [self.ignore_index] * (self.max_text_len + 1 - len(text)) + data['label'] = np.array(text) + return data + + def add_special_char(self, dict_character): + dict_character = [''] + dict_character + return dict_character diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index 09736515e7..5397d71ccb 100644 --- a/ppocr/data/imaug/operators.py +++ b/ppocr/data/imaug/operators.py @@ -67,39 +67,6 @@ class DecodeImage(object): return data -class NRTRDecodeImage(object): - """ decode image """ - - def __init__(self, img_mode='RGB', channel_first=False, **kwargs): - self.img_mode = img_mode - self.channel_first = channel_first - - def __call__(self, data): - img = data['image'] - if six.PY2: - assert type(img) is str and len( - img) > 0, "invalid input 'img' in DecodeImage" - else: - assert type(img) is bytes and len( - img) > 0, "invalid input 'img' in DecodeImage" - img = np.frombuffer(img, dtype='uint8') - - img = cv2.imdecode(img, 1) - - if img is None: - return None - if self.img_mode == 'GRAY': - img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) - elif self.img_mode == 'RGB': - assert img.shape[2] == 3, 'invalid shape of image[%s]' % (img.shape) - img = img[:, :, ::-1] - img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) - if self.channel_first: - img = img.transpose((2, 0, 1)) - data['image'] = img - return data - - class NormalizeImage(object): """ normalize image such as substract mean, divide std """ diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index d3120f5950..a8b3b813b3 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -279,6 +279,24 @@ class PRENResizeImg(object): return data +class ABINetRecResizeImg(object): + def __init__(self, + image_shape, + infer_mode=False, + character_dict_path=None, + **kwargs): + self.image_shape = image_shape + self.infer_mode = infer_mode + self.character_dict_path = character_dict_path + + def __call__(self, data): + img = data['image'] + norm_img, valid_ratio = resize_norm_img_abinet(img, self.image_shape) + data['image'] = norm_img + data['valid_ratio'] = valid_ratio + return data + + def resize_norm_img_sar(img, image_shape, width_downsample_ratio=0.25): imgC, imgH, imgW_min, imgW_max = image_shape h = img.shape[0] @@ -397,6 +415,26 @@ def resize_norm_img_srn(img, image_shape): return np.reshape(img_black, (c, row, col)).astype(np.float32) +def resize_norm_img_abinet(img, image_shape): + imgC, imgH, imgW = image_shape + + resized_image = cv2.resize( + img, (imgW, imgH), interpolation=cv2.INTER_LINEAR) + resized_w = imgW + resized_image = resized_image.astype('float32') + resized_image = resized_image / 255. + + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + resized_image = ( + resized_image - mean[None, None, ...]) / std[None, None, ...] + resized_image = resized_image.transpose((2, 0, 1)) + resized_image = resized_image.astype('float32') + + valid_ratio = min(1.0, float(resized_w / imgW)) + return resized_image, valid_ratio + + def srn_other_inputs(image_shape, num_heads, max_text_length): imgC, imgH, imgW = image_shape diff --git a/ppocr/losses/rec_ce_loss.py b/ppocr/losses/rec_ce_loss.py index b837ac27bb..614384de86 100644 --- a/ppocr/losses/rec_ce_loss.py +++ b/ppocr/losses/rec_ce_loss.py @@ -4,31 +4,63 @@ import paddle.nn.functional as F class CELoss(nn.Layer): - def __init__(self, smoothing=True, with_all=False, **kwargs): + def __init__(self, + smoothing=False, + with_all=False, + ignore_index=-1, + **kwargs): super(CELoss, self).__init__() - self.loss_func = nn.CrossEntropyLoss(reduction='mean', ignore_index=0) + if ignore_index >= 0: + self.loss_func = nn.CrossEntropyLoss( + reduction='mean', ignore_index=ignore_index) + else: + self.loss_func = nn.CrossEntropyLoss(reduction='mean') self.smoothing = smoothing self.with_all = with_all def forward(self, pred, batch): - pred = pred.reshape([-1, pred.shape[2]]) - if self.with_all: - tgt = batch[1] + + if isinstance(pred, dict): # for ABINet + loss = {} + loss_sum = [] + for name, logits in pred.items(): + if isinstance(logits, list): + logit_num = len(logits) + all_tgt = paddle.concat([batch[1]] * logit_num, 0) + all_logits = paddle.concat(logits, 0) + flt_logtis = all_logits.reshape([-1, all_logits.shape[2]]) + flt_tgt = all_tgt.reshape([-1]) + else: + flt_logtis = logits.reshape([-1, logits.shape[2]]) + flt_tgt = batch[1].reshape([-1]) + loss[name + '_loss'] = self.loss_func(flt_logtis, flt_tgt) + loss_sum.append(loss[name + '_loss']) + loss['loss'] = sum(loss_sum) + return loss else: - max_len = batch[2].max() - tgt = batch[1][:, 1:2 + max_len] - tgt = tgt.reshape([-1]) - if self.smoothing: - eps = 0.1 - n_class = pred.shape[1] - one_hot = F.one_hot(tgt, pred.shape[1]) - one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1) - log_prb = F.log_softmax(pred, axis=1) - non_pad_mask = paddle.not_equal( - tgt, paddle.zeros( - tgt.shape, dtype=tgt.dtype)) - loss = -(one_hot * log_prb).sum(axis=1) - loss = loss.masked_select(non_pad_mask).mean() - else: - loss = self.loss_func(pred, tgt) - return {'loss': loss} + if self.with_all: # for ViTSTR + tgt = batch[1] + pred = pred.reshape([-1, pred.shape[2]]) + tgt = tgt.reshape([-1]) + loss = self.loss_func(pred, tgt) + return {'loss': loss} + else: # for NRTR + max_len = batch[2].max() + tgt = batch[1][:, 1:2 + max_len] + pred = pred.reshape([-1, pred.shape[2]]) + tgt = tgt.reshape([-1]) + if self.smoothing: + eps = 0.1 + n_class = pred.shape[1] + one_hot = F.one_hot(tgt, pred.shape[1]) + one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / ( + n_class - 1) + log_prb = F.log_softmax(pred, axis=1) + non_pad_mask = paddle.not_equal( + tgt, paddle.zeros( + tgt.shape, dtype=tgt.dtype)) + loss = -(one_hot * log_prb).sum(axis=1) + loss = loss.masked_select(non_pad_mask).mean() + else: + loss = self.loss_func(pred, tgt) + return {'loss': loss} diff --git a/ppocr/modeling/backbones/__init__.py b/ppocr/modeling/backbones/__init__.py index a368e74816..56d2c2b9db 100755 --- a/ppocr/modeling/backbones/__init__.py +++ b/ppocr/modeling/backbones/__init__.py @@ -27,7 +27,7 @@ def build_backbone(config, model_type): from .rec_resnet_fpn import ResNetFPN from .rec_mv1_enhance import MobileNetV1Enhance from .rec_nrtr_mtb import MTB - from .rec_resnet_31 import ResNet31 + from .rec_resnet import ResNet31, ResNet45 from .rec_resnet_aster import ResNet_ASTER from .rec_micronet import MicroNet from .rec_efficientb3_pren import EfficientNetb3_PREN @@ -35,29 +35,29 @@ def build_backbone(config, model_type): from .rec_vitstr import ViTSTR support_dict = [ 'MobileNetV1Enhance', 'MobileNetV3', 'ResNet', 'ResNetFPN', 'MTB', - "ResNet31", "ResNet_ASTER", 'MicroNet', 'EfficientNetb3_PREN', - 'SVTRNet', 'ViTSTR' + 'ResNet31', 'ResNet45', 'ResNet_ASTER', 'MicroNet', + 'EfficientNetb3_PREN', 'SVTRNet', 'ViTSTR' ] - elif model_type == "e2e": + elif model_type == 'e2e': from .e2e_resnet_vd_pg import ResNet support_dict = ['ResNet'] elif model_type == 'kie': from .kie_unet_sdmgr import Kie_backbone support_dict = ['Kie_backbone'] - elif model_type == "table": + elif model_type == 'table': from .table_resnet_vd import ResNet from .table_mobilenet_v3 import MobileNetV3 - support_dict = ["ResNet", "MobileNetV3"] + support_dict = ['ResNet', 'MobileNetV3'] elif model_type == 'vqa': from .vqa_layoutlm import LayoutLMForSer, LayoutLMv2ForSer, LayoutLMv2ForRe, LayoutXLMForSer, LayoutXLMForRe support_dict = [ - "LayoutLMForSer", "LayoutLMv2ForSer", 'LayoutLMv2ForRe', - "LayoutXLMForSer", 'LayoutXLMForRe' + 'LayoutLMForSer', 'LayoutLMv2ForSer', 'LayoutLMv2ForRe', + 'LayoutXLMForSer', 'LayoutXLMForRe' ] else: raise NotImplementedError - module_name = config.pop("name") + module_name = config.pop('name') assert module_name in support_dict, Exception( "when model typs is {}, backbone only support {}".format(model_type, support_dict)) diff --git a/ppocr/modeling/backbones/rec_resnet.py b/ppocr/modeling/backbones/rec_resnet.py new file mode 100644 index 0000000000..ce76b642cc --- /dev/null +++ b/ppocr/modeling/backbones/rec_resnet.py @@ -0,0 +1,280 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This code is refer from: +https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/layers/conv_layer.py +https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/backbones/resnet31_ocr.py +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import paddle +from paddle import ParamAttr +from paddle.nn.initializer import KaimingNormal +import paddle.nn as nn +import paddle.nn.functional as F +import numpy as np +import math + +__all__ = ["ResNet31", "ResNet45"] + + +def conv1x1(in_planes, out_planes, stride=1): + return nn.Conv2D( + in_planes, + out_planes, + kernel_size=1, + stride=stride, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + + +def conv3x3(in_channel, out_channel, stride=1): + return nn.Conv2D( + in_channel, + out_channel, + kernel_size=3, + stride=stride, + padding=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + + +class BasicBlock(nn.Layer): + expansion = 1 + + def __init__(self, in_channels, channels, stride=1, downsample=None): + super().__init__() + self.conv1 = conv1x1(in_channels, channels) + self.bn1 = nn.BatchNorm2D(channels) + self.relu = nn.ReLU() + self.conv2 = conv3x3(channels, channels, stride) + self.bn2 = nn.BatchNorm2D(channels) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + out += residual + out = self.relu(out) + + return out + + +class ResNet31(nn.Layer): + ''' + Args: + in_channels (int): Number of channels of input image tensor. + layers (list[int]): List of BasicBlock number for each stage. + channels (list[int]): List of out_channels of Conv2d layer. + out_indices (None | Sequence[int]): Indices of output stages. + last_stage_pool (bool): If True, add `MaxPool2d` layer to last stage. + ''' + + def __init__(self, + in_channels=3, + layers=[1, 2, 5, 3], + channels=[64, 128, 256, 256, 512, 512, 512], + out_indices=None, + last_stage_pool=False): + super(ResNet31, self).__init__() + assert isinstance(in_channels, int) + assert isinstance(last_stage_pool, bool) + + self.out_indices = out_indices + self.last_stage_pool = last_stage_pool + + # conv 1 (Conv Conv) + self.conv1_1 = nn.Conv2D( + in_channels, channels[0], kernel_size=3, stride=1, padding=1) + self.bn1_1 = nn.BatchNorm2D(channels[0]) + self.relu1_1 = nn.ReLU() + + self.conv1_2 = nn.Conv2D( + channels[0], channels[1], kernel_size=3, stride=1, padding=1) + self.bn1_2 = nn.BatchNorm2D(channels[1]) + self.relu1_2 = nn.ReLU() + + # conv 2 (Max-pooling, Residual block, Conv) + self.pool2 = nn.MaxPool2D( + kernel_size=2, stride=2, padding=0, ceil_mode=True) + self.block2 = self._make_layer(channels[1], channels[2], layers[0]) + self.conv2 = nn.Conv2D( + channels[2], channels[2], kernel_size=3, stride=1, padding=1) + self.bn2 = nn.BatchNorm2D(channels[2]) + self.relu2 = nn.ReLU() + + # conv 3 (Max-pooling, Residual block, Conv) + self.pool3 = nn.MaxPool2D( + kernel_size=2, stride=2, padding=0, ceil_mode=True) + self.block3 = self._make_layer(channels[2], channels[3], layers[1]) + self.conv3 = nn.Conv2D( + channels[3], channels[3], kernel_size=3, stride=1, padding=1) + self.bn3 = nn.BatchNorm2D(channels[3]) + self.relu3 = nn.ReLU() + + # conv 4 (Max-pooling, Residual block, Conv) + self.pool4 = nn.MaxPool2D( + kernel_size=(2, 1), stride=(2, 1), padding=0, ceil_mode=True) + self.block4 = self._make_layer(channels[3], channels[4], layers[2]) + self.conv4 = nn.Conv2D( + channels[4], channels[4], kernel_size=3, stride=1, padding=1) + self.bn4 = nn.BatchNorm2D(channels[4]) + self.relu4 = nn.ReLU() + + # conv 5 ((Max-pooling), Residual block, Conv) + self.pool5 = None + if self.last_stage_pool: + self.pool5 = nn.MaxPool2D( + kernel_size=2, stride=2, padding=0, ceil_mode=True) + self.block5 = self._make_layer(channels[4], channels[5], layers[3]) + self.conv5 = nn.Conv2D( + channels[5], channels[5], kernel_size=3, stride=1, padding=1) + self.bn5 = nn.BatchNorm2D(channels[5]) + self.relu5 = nn.ReLU() + + self.out_channels = channels[-1] + + def _make_layer(self, input_channels, output_channels, blocks): + layers = [] + for _ in range(blocks): + downsample = None + if input_channels != output_channels: + downsample = nn.Sequential( + nn.Conv2D( + input_channels, + output_channels, + kernel_size=1, + stride=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False), + nn.BatchNorm2D(output_channels), ) + + layers.append( + BasicBlock( + input_channels, output_channels, downsample=downsample)) + input_channels = output_channels + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1_1(x) + x = self.bn1_1(x) + x = self.relu1_1(x) + + x = self.conv1_2(x) + x = self.bn1_2(x) + x = self.relu1_2(x) + + outs = [] + for i in range(4): + layer_index = i + 2 + pool_layer = getattr(self, f'pool{layer_index}') + block_layer = getattr(self, f'block{layer_index}') + conv_layer = getattr(self, f'conv{layer_index}') + bn_layer = getattr(self, f'bn{layer_index}') + relu_layer = getattr(self, f'relu{layer_index}') + + if pool_layer is not None: + x = pool_layer(x) + x = block_layer(x) + x = conv_layer(x) + x = bn_layer(x) + x = relu_layer(x) + + outs.append(x) + + if self.out_indices is not None: + return tuple([outs[i] for i in self.out_indices]) + + return x + + +class ResNet(nn.Layer): + def __init__(self, block, layers, in_channels=3): + self.inplanes = 32 + super(ResNet, self).__init__() + self.conv1 = nn.Conv2D( + 3, + 32, + kernel_size=3, + stride=1, + padding=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + self.bn1 = nn.BatchNorm2D(32) + self.relu = nn.ReLU() + + self.layer1 = self._make_layer(block, 32, layers[0], stride=2) + self.layer2 = self._make_layer(block, 64, layers[1], stride=1) + self.layer3 = self._make_layer(block, 128, layers[2], stride=2) + self.layer4 = self._make_layer(block, 256, layers[3], stride=1) + self.layer5 = self._make_layer(block, 512, layers[4], stride=1) + self.out_channels = 512 + + # for m in self.modules(): + # if isinstance(m, nn.Conv2D): + # n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels + # m.weight.data.normal_(0, math.sqrt(2. / n)) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + # downsample = True + downsample = nn.Sequential( + nn.Conv2D( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False), + nn.BatchNorm2D(planes * block.expansion), ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + # print(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + # print(x) + x = self.layer4(x) + x = self.layer5(x) + return x + + +def ResNet45(in_channels=3): + return ResNet(BasicBlock, [3, 4, 6, 6, 3], in_channels=in_channels) diff --git a/ppocr/modeling/heads/__init__.py b/ppocr/modeling/heads/__init__.py index 1670ea38e6..14e6aab854 100755 --- a/ppocr/modeling/heads/__init__.py +++ b/ppocr/modeling/heads/__init__.py @@ -33,6 +33,7 @@ def build_head(config): from .rec_aster_head import AsterHead from .rec_pren_head import PRENHead from .rec_multi_head import MultiHead + from .rec_abinet_head import ABINetHead # cls head from .cls_head import ClsHead @@ -46,7 +47,7 @@ def build_head(config): 'DBHead', 'PSEHead', 'FCEHead', 'EASTHead', 'SASTHead', 'CTCHead', 'ClsHead', 'AttentionHead', 'SRNHead', 'PGHead', 'Transformer', 'TableAttentionHead', 'SARHead', 'AsterHead', 'SDMGRHead', 'PRENHead', - 'MultiHead' + 'MultiHead', 'ABINetHead' ] #table head diff --git a/ppocr/modeling/heads/multiheadAttention.py b/ppocr/modeling/heads/multiheadAttention.py deleted file mode 100755 index 900865ba1a..0000000000 --- a/ppocr/modeling/heads/multiheadAttention.py +++ /dev/null @@ -1,163 +0,0 @@ -# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import paddle -from paddle import nn -import paddle.nn.functional as F -from paddle.nn import Linear -from paddle.nn.initializer import XavierUniform as xavier_uniform_ -from paddle.nn.initializer import Constant as constant_ -from paddle.nn.initializer import XavierNormal as xavier_normal_ - -zeros_ = constant_(value=0.) -ones_ = constant_(value=1.) - - -class MultiheadAttention(nn.Layer): - """Allows the model to jointly attend to information - from different representation subspaces. - See reference: Attention Is All You Need - - .. math:: - \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O - \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) - - Args: - embed_dim: total dimension of the model - num_heads: parallel attention layers, or heads - - """ - - def __init__(self, - embed_dim, - num_heads, - dropout=0., - bias=True, - add_bias_kv=False, - add_zero_attn=False): - super(MultiheadAttention, self).__init__() - self.embed_dim = embed_dim - self.num_heads = num_heads - self.dropout = dropout - self.head_dim = embed_dim // num_heads - assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" - self.scaling = self.head_dim**-0.5 - self.out_proj = Linear(embed_dim, embed_dim, bias_attr=bias) - self._reset_parameters() - self.conv1 = paddle.nn.Conv2D( - in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)) - self.conv2 = paddle.nn.Conv2D( - in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)) - self.conv3 = paddle.nn.Conv2D( - in_channels=embed_dim, out_channels=embed_dim, kernel_size=(1, 1)) - - def _reset_parameters(self): - xavier_uniform_(self.out_proj.weight) - - def forward(self, - query, - key, - value, - key_padding_mask=None, - incremental_state=None, - attn_mask=None): - """ - Inputs of forward function - query: [target length, batch size, embed dim] - key: [sequence length, batch size, embed dim] - value: [sequence length, batch size, embed dim] - key_padding_mask: if True, mask padding based on batch size - incremental_state: if provided, previous time steps are cashed - need_weights: output attn_output_weights - static_kv: key and value are static - - Outputs of forward function - attn_output: [target length, batch size, embed dim] - attn_output_weights: [batch size, target length, sequence length] - """ - q_shape = paddle.shape(query) - src_shape = paddle.shape(key) - q = self._in_proj_q(query) - k = self._in_proj_k(key) - v = self._in_proj_v(value) - q *= self.scaling - q = paddle.transpose( - paddle.reshape( - q, [q_shape[0], q_shape[1], self.num_heads, self.head_dim]), - [1, 2, 0, 3]) - k = paddle.transpose( - paddle.reshape( - k, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]), - [1, 2, 0, 3]) - v = paddle.transpose( - paddle.reshape( - v, [src_shape[0], q_shape[1], self.num_heads, self.head_dim]), - [1, 2, 0, 3]) - if key_padding_mask is not None: - assert key_padding_mask.shape[0] == q_shape[1] - assert key_padding_mask.shape[1] == src_shape[0] - attn_output_weights = paddle.matmul(q, - paddle.transpose(k, [0, 1, 3, 2])) - if attn_mask is not None: - attn_mask = paddle.unsqueeze(paddle.unsqueeze(attn_mask, 0), 0) - attn_output_weights += attn_mask - if key_padding_mask is not None: - attn_output_weights = paddle.reshape( - attn_output_weights, - [q_shape[1], self.num_heads, q_shape[0], src_shape[0]]) - key = paddle.unsqueeze(paddle.unsqueeze(key_padding_mask, 1), 2) - key = paddle.cast(key, 'float32') - y = paddle.full( - shape=paddle.shape(key), dtype='float32', fill_value='-inf') - y = paddle.where(key == 0., key, y) - attn_output_weights += y - attn_output_weights = F.softmax( - attn_output_weights.astype('float32'), - axis=-1, - dtype=paddle.float32 if attn_output_weights.dtype == paddle.float16 - else attn_output_weights.dtype) - attn_output_weights = F.dropout( - attn_output_weights, p=self.dropout, training=self.training) - - attn_output = paddle.matmul(attn_output_weights, v) - attn_output = paddle.reshape( - paddle.transpose(attn_output, [2, 0, 1, 3]), - [q_shape[0], q_shape[1], self.embed_dim]) - attn_output = self.out_proj(attn_output) - - return attn_output - - def _in_proj_q(self, query): - query = paddle.transpose(query, [1, 2, 0]) - query = paddle.unsqueeze(query, axis=2) - res = self.conv1(query) - res = paddle.squeeze(res, axis=2) - res = paddle.transpose(res, [2, 0, 1]) - return res - - def _in_proj_k(self, key): - key = paddle.transpose(key, [1, 2, 0]) - key = paddle.unsqueeze(key, axis=2) - res = self.conv2(key) - res = paddle.squeeze(res, axis=2) - res = paddle.transpose(res, [2, 0, 1]) - return res - - def _in_proj_v(self, value): - value = paddle.transpose(value, [1, 2, 0]) #(1, 2, 0) - value = paddle.unsqueeze(value, axis=2) - res = self.conv3(value) - res = paddle.squeeze(res, axis=2) - res = paddle.transpose(res, [2, 0, 1]) - return res diff --git a/ppocr/modeling/heads/rec_abinet_head.py b/ppocr/modeling/heads/rec_abinet_head.py new file mode 100644 index 0000000000..a0f60f1be1 --- /dev/null +++ b/ppocr/modeling/heads/rec_abinet_head.py @@ -0,0 +1,296 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This code is refer from: +https://github.com/FangShancheng/ABINet/tree/main/modules +""" + +import math +import paddle +from paddle import nn +import paddle.nn.functional as F +from paddle.nn import LayerList +from ppocr.modeling.heads.rec_nrtr_head import TransformerBlock, PositionalEncoding + + +class BCNLanguage(nn.Layer): + def __init__(self, + d_model=512, + nhead=8, + num_layers=4, + dim_feedforward=2048, + dropout=0., + max_length=25, + detach=True, + num_classes=37): + super().__init__() + + self.d_model = d_model + self.detach = detach + self.max_length = max_length + 1 # additional stop token + self.proj = nn.Linear(num_classes, d_model, bias_attr=False) + self.token_encoder = PositionalEncoding( + dropout=0.1, dim=d_model, max_len=self.max_length) + self.pos_encoder = PositionalEncoding( + dropout=0, dim=d_model, max_len=self.max_length) + + self.decoder = nn.LayerList([ + TransformerBlock( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + attention_dropout_rate=dropout, + residual_dropout_rate=dropout, + with_self_attn=False, + with_cross_attn=True) for i in range(num_layers) + ]) + + self.cls = nn.Linear(d_model, num_classes) + + def forward(self, tokens, lengths): + """ + Args: + tokens: (B, N, C) where N is length, B is batch size and C is classes number + lengths: (B,) + """ + if self.detach: tokens = tokens.detach() + embed = self.proj(tokens) # (B, N, C) + embed = self.token_encoder(embed) # (B, N, C) + padding_mask = _get_mask(lengths, self.max_length) + zeros = paddle.zeros_like(embed) # (B, N, C) + qeury = self.pos_encoder(zeros) + for decoder_layer in self.decoder: + qeury = decoder_layer(qeury, embed, cross_mask=padding_mask) + output = qeury # (B, N, C) + + logits = self.cls(output) # (B, N, C) + + return output, logits + + +def encoder_layer(in_c, out_c, k=3, s=2, p=1): + return nn.Sequential( + nn.Conv2D(in_c, out_c, k, s, p), nn.BatchNorm2D(out_c), nn.ReLU()) + + +def decoder_layer(in_c, + out_c, + k=3, + s=1, + p=1, + mode='nearest', + scale_factor=None, + size=None): + align_corners = False if mode == 'nearest' else True + return nn.Sequential( + nn.Upsample( + size=size, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners), + nn.Conv2D(in_c, out_c, k, s, p), + nn.BatchNorm2D(out_c), + nn.ReLU()) + + +class PositionAttention(nn.Layer): + def __init__(self, + max_length, + in_channels=512, + num_channels=64, + h=8, + w=32, + mode='nearest', + **kwargs): + super().__init__() + self.max_length = max_length + self.k_encoder = nn.Sequential( + encoder_layer( + in_channels, num_channels, s=(1, 2)), + encoder_layer( + num_channels, num_channels, s=(2, 2)), + encoder_layer( + num_channels, num_channels, s=(2, 2)), + encoder_layer( + num_channels, num_channels, s=(2, 2))) + self.k_decoder = nn.Sequential( + decoder_layer( + num_channels, num_channels, scale_factor=2, mode=mode), + decoder_layer( + num_channels, num_channels, scale_factor=2, mode=mode), + decoder_layer( + num_channels, num_channels, scale_factor=2, mode=mode), + decoder_layer( + num_channels, in_channels, size=(h, w), mode=mode)) + + self.pos_encoder = PositionalEncoding( + dropout=0, dim=in_channels, max_len=max_length) + self.project = nn.Linear(in_channels, in_channels) + + def forward(self, x): + B, C, H, W = x.shape + k, v = x, x + + # calculate key vector + features = [] + for i in range(0, len(self.k_encoder)): + k = self.k_encoder[i](k) + features.append(k) + for i in range(0, len(self.k_decoder) - 1): + k = self.k_decoder[i](k) + # print(k.shape, features[len(self.k_decoder) - 2 - i].shape) + k = k + features[len(self.k_decoder) - 2 - i] + k = self.k_decoder[-1](k) + + # calculate query vector + # TODO q=f(q,k) + zeros = paddle.zeros( + (B, self.max_length, C), dtype=x.dtype) # (T, N, C) + q = self.pos_encoder(zeros) # (B, N, C) + q = self.project(q) # (B, N, C) + + # calculate attention + attn_scores = q @k.flatten(2) # (B, N, (H*W)) + attn_scores = attn_scores / (C**0.5) + attn_scores = F.softmax(attn_scores, axis=-1) + + v = v.flatten(2).transpose([0, 2, 1]) # (B, (H*W), C) + attn_vecs = attn_scores @v # (B, N, C) + + return attn_vecs, attn_scores.reshape([0, self.max_length, H, W]) + + +class ABINetHead(nn.Layer): + def __init__(self, + in_channels, + out_channels, + d_model=512, + nhead=8, + num_layers=3, + dim_feedforward=2048, + dropout=0.1, + max_length=25, + use_lang=False, + iter_size=1): + super().__init__() + self.max_length = max_length + 1 + self.pos_encoder = PositionalEncoding( + dropout=0.1, dim=d_model, max_len=8 * 32) + self.encoder = nn.LayerList([ + TransformerBlock( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + attention_dropout_rate=dropout, + residual_dropout_rate=dropout, + with_self_attn=True, + with_cross_attn=False) for i in range(num_layers) + ]) + self.decoder = PositionAttention( + max_length=max_length + 1, # additional stop token + mode='nearest', ) + self.out_channels = out_channels + self.cls = nn.Linear(d_model, self.out_channels) + self.use_lang = use_lang + if use_lang: + self.iter_size = iter_size + self.language = BCNLanguage( + d_model=d_model, + nhead=nhead, + num_layers=4, + dim_feedforward=dim_feedforward, + dropout=dropout, + max_length=max_length, + num_classes=self.out_channels) + # alignment + self.w_att_align = nn.Linear(2 * d_model, d_model) + self.cls_align = nn.Linear(d_model, self.out_channels) + + def forward(self, x, targets=None): + x = x.transpose([0, 2, 3, 1]) + _, H, W, C = x.shape + feature = x.flatten(1, 2) + feature = self.pos_encoder(feature) + for encoder_layer in self.encoder: + feature = encoder_layer(feature) + feature = feature.reshape([0, H, W, C]).transpose([0, 3, 1, 2]) + v_feature, attn_scores = self.decoder( + feature) # (B, N, C), (B, C, H, W) + vis_logits = self.cls(v_feature) # (B, N, C) + logits = vis_logits + vis_lengths = _get_length(vis_logits) + if self.use_lang: + align_logits = vis_logits + align_lengths = vis_lengths + all_l_res, all_a_res = [], [] + for i in range(self.iter_size): + tokens = F.softmax(align_logits, axis=-1) + lengths = align_lengths + lengths = paddle.clip( + lengths, 2, self.max_length) # TODO:move to langauge model + l_feature, l_logits = self.language(tokens, lengths) + + # alignment + all_l_res.append(l_logits) + fuse = paddle.concat((l_feature, v_feature), -1) + f_att = F.sigmoid(self.w_att_align(fuse)) + output = f_att * v_feature + (1 - f_att) * l_feature + align_logits = self.cls_align(output) # (B, N, C) + + align_lengths = _get_length(align_logits) + all_a_res.append(align_logits) + if self.training: + return { + 'align': all_a_res, + 'lang': all_l_res, + 'vision': vis_logits + } + else: + logits = align_logits + if self.training: + return logits + else: + return F.softmax(logits, -1) + + +def _get_length(logit): + """ Greed decoder to obtain length from logit""" + out = (logit.argmax(-1) == 0) + abn = out.any(-1) + out_int = out.cast('int32') + out = (out_int.cumsum(-1) == 1) & out + out = out.cast('int32') + out = out.argmax(-1) + out = out + 1 + out = paddle.where(abn, out, paddle.to_tensor(logit.shape[1])) + return out + + +def _get_mask(length, max_length): + """Generate a square mask for the sequence. The masked positions are filled with float('-inf'). + Unmasked positions are filled with float(0.0). + """ + length = length.unsqueeze(-1) + B = paddle.shape(length)[0] + grid = paddle.arange(0, max_length).unsqueeze(0).tile([B, 1]) + zero_mask = paddle.zeros([B, max_length], dtype='float32') + inf_mask = paddle.full([B, max_length], '-inf', dtype='float32') + diag_mask = paddle.diag( + paddle.full( + [max_length], '-inf', dtype=paddle.float32), + offset=0, + name=None) + mask = paddle.where(grid >= length, inf_mask, zero_mask) + mask = mask.unsqueeze(1) + diag_mask + return mask.unsqueeze(1) diff --git a/ppocr/modeling/heads/rec_nrtr_head.py b/ppocr/modeling/heads/rec_nrtr_head.py index 38ba0c9178..bf9ef56145 100644 --- a/ppocr/modeling/heads/rec_nrtr_head.py +++ b/ppocr/modeling/heads/rec_nrtr_head.py @@ -14,20 +14,15 @@ import math import paddle -import copy from paddle import nn import paddle.nn.functional as F from paddle.nn import LayerList -from paddle.nn.initializer import XavierNormal as xavier_uniform_ -from paddle.nn import Dropout, Linear, LayerNorm, Conv2D +# from paddle.nn.initializer import XavierNormal as xavier_uniform_ +from paddle.nn import Dropout, Linear, LayerNorm import numpy as np -from ppocr.modeling.heads.multiheadAttention import MultiheadAttention -from paddle.nn.initializer import Constant as constant_ +from ppocr.modeling.backbones.rec_svtrnet import Mlp, zeros_, ones_ from paddle.nn.initializer import XavierNormal as xavier_normal_ -zeros_ = constant_(value=0.) -ones_ = constant_(value=1.) - class Transformer(nn.Layer): """A transformer model. User is able to modify the attributes as needed. The architechture @@ -45,7 +40,6 @@ class Transformer(nn.Layer): dropout: the dropout value (default=0.1). custom_encoder: custom encoder (default=None). custom_decoder: custom decoder (default=None). - """ def __init__(self, @@ -54,45 +48,49 @@ class Transformer(nn.Layer): num_encoder_layers=6, beam_size=0, num_decoder_layers=6, + max_len=25, dim_feedforward=1024, attention_dropout_rate=0.0, residual_dropout_rate=0.1, - custom_encoder=None, - custom_decoder=None, in_channels=0, out_channels=0, scale_embedding=True): super(Transformer, self).__init__() self.out_channels = out_channels + 1 + self.max_len = max_len self.embedding = Embeddings( d_model=d_model, vocab=self.out_channels, padding_idx=0, scale_embedding=scale_embedding) self.positional_encoding = PositionalEncoding( - dropout=residual_dropout_rate, - dim=d_model, ) - if custom_encoder is not None: - self.encoder = custom_encoder - else: - if num_encoder_layers > 0: - encoder_layer = TransformerEncoderLayer( - d_model, nhead, dim_feedforward, attention_dropout_rate, - residual_dropout_rate) - self.encoder = TransformerEncoder(encoder_layer, - num_encoder_layers) - else: - self.encoder = None + dropout=residual_dropout_rate, dim=d_model) - if custom_decoder is not None: - self.decoder = custom_decoder + if num_encoder_layers > 0: + self.encoder = nn.LayerList([ + TransformerBlock( + d_model, + nhead, + dim_feedforward, + attention_dropout_rate, + residual_dropout_rate, + with_self_attn=True, + with_cross_attn=False) for i in range(num_encoder_layers) + ]) else: - decoder_layer = TransformerDecoderLayer( - d_model, nhead, dim_feedforward, attention_dropout_rate, - residual_dropout_rate) - self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers) + self.encoder = None + + self.decoder = nn.LayerList([ + TransformerBlock( + d_model, + nhead, + dim_feedforward, + attention_dropout_rate, + residual_dropout_rate, + with_self_attn=True, + with_cross_attn=True) for i in range(num_decoder_layers) + ]) - self._reset_parameters() self.beam_size = beam_size self.d_model = d_model self.nhead = nhead @@ -105,7 +103,7 @@ class Transformer(nn.Layer): def _init_weights(self, m): - if isinstance(m, nn.Conv2D): + if isinstance(m, nn.Linear): xavier_normal_(m.weight) if m.bias is not None: zeros_(m.bias) @@ -113,24 +111,20 @@ class Transformer(nn.Layer): def forward_train(self, src, tgt): tgt = tgt[:, :-1] - tgt_key_padding_mask = self.generate_padding_mask(tgt) - tgt = self.embedding(tgt).transpose([1, 0, 2]) + tgt = self.embedding(tgt) tgt = self.positional_encoding(tgt) - tgt_mask = self.generate_square_subsequent_mask(tgt.shape[0]) + tgt_mask = self.generate_square_subsequent_mask(tgt.shape[1]) if self.encoder is not None: - src = self.positional_encoding(src.transpose([1, 0, 2])) - memory = self.encoder(src) + src = self.positional_encoding(src) + for encoder_layer in self.encoder: + src = encoder_layer(src) + memory = src # B N C else: - memory = src.squeeze(2).transpose([2, 0, 1]) - output = self.decoder( - tgt, - memory, - tgt_mask=tgt_mask, - memory_mask=None, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=None) - output = output.transpose([1, 0, 2]) + memory = src # B N C + for decoder_layer in self.decoder: + tgt = decoder_layer(tgt, memory, self_mask=tgt_mask) + output = tgt logit = self.tgt_word_prj(output) return logit @@ -140,8 +134,8 @@ class Transformer(nn.Layer): src: the sequence to the encoder (required). tgt: the sequence to the decoder (required). Shape: - - src: :math:`(S, N, E)`. - - tgt: :math:`(T, N, E)`. + - src: :math:`(B, sN, C)`. + - tgt: :math:`(B, tN, C)`. Examples: >>> output = transformer_model(src, tgt) """ @@ -157,36 +151,35 @@ class Transformer(nn.Layer): return self.forward_test(src) def forward_test(self, src): + bs = paddle.shape(src)[0] if self.encoder is not None: - src = self.positional_encoding(paddle.transpose(src, [1, 0, 2])) - memory = self.encoder(src) + src = self.positional_encoding(src) + for encoder_layer in self.encoder: + src = encoder_layer(src) + memory = src # B N C else: - memory = paddle.transpose(paddle.squeeze(src, 2), [2, 0, 1]) + memory = src dec_seq = paddle.full((bs, 1), 2, dtype=paddle.int64) dec_prob = paddle.full((bs, 1), 1., dtype=paddle.float32) - for len_dec_seq in range(1, 25): - dec_seq_embed = paddle.transpose(self.embedding(dec_seq), [1, 0, 2]) + for len_dec_seq in range(1, self.max_len): + dec_seq_embed = self.embedding(dec_seq) dec_seq_embed = self.positional_encoding(dec_seq_embed) tgt_mask = self.generate_square_subsequent_mask( - paddle.shape(dec_seq_embed)[0]) - output = self.decoder( - dec_seq_embed, - memory, - tgt_mask=tgt_mask, - memory_mask=None, - tgt_key_padding_mask=None, - memory_key_padding_mask=None) - dec_output = paddle.transpose(output, [1, 0, 2]) + paddle.shape(dec_seq_embed)[1]) + tgt = dec_seq_embed + for decoder_layer in self.decoder: + tgt = decoder_layer(tgt, memory, self_mask=tgt_mask) + dec_output = tgt dec_output = dec_output[:, -1, :] - word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=1) - preds_idx = paddle.argmax(word_prob, axis=1) + word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=-1) + preds_idx = paddle.argmax(word_prob, axis=-1) if paddle.equal_all( preds_idx, paddle.full( paddle.shape(preds_idx), 3, dtype='int64')): break - preds_prob = paddle.max(word_prob, axis=1) + preds_prob = paddle.max(word_prob, axis=-1) dec_seq = paddle.concat( [dec_seq, paddle.reshape(preds_idx, [-1, 1])], axis=1) dec_prob = paddle.concat( @@ -194,10 +187,10 @@ class Transformer(nn.Layer): return [dec_seq, dec_prob] def forward_beam(self, images): - ''' Translation work in one batch ''' + """ Translation work in one batch """ def get_inst_idx_to_tensor_position_map(inst_idx_list): - ''' Indicate the position of an instance in a tensor. ''' + """ Indicate the position of an instance in a tensor. """ return { inst_idx: tensor_position for tensor_position, inst_idx in enumerate(inst_idx_list) @@ -205,7 +198,7 @@ class Transformer(nn.Layer): def collect_active_part(beamed_tensor, curr_active_inst_idx, n_prev_active_inst, n_bm): - ''' Collect tensor parts associated to active instances. ''' + """ Collect tensor parts associated to active instances. """ beamed_tensor_shape = paddle.shape(beamed_tensor) n_curr_active_inst = len(curr_active_inst_idx) @@ -237,9 +230,8 @@ class Transformer(nn.Layer): return active_src_enc, active_inst_idx_to_position_map def beam_decode_step(inst_dec_beams, len_dec_seq, enc_output, - inst_idx_to_position_map, n_bm, - memory_key_padding_mask): - ''' Decode and update beam status, and then return active beam idx ''' + inst_idx_to_position_map, n_bm): + """ Decode and update beam status, and then return active beam idx """ def prepare_beam_dec_seq(inst_dec_beams, len_dec_seq): dec_partial_seq = [ @@ -249,19 +241,15 @@ class Transformer(nn.Layer): dec_partial_seq = dec_partial_seq.reshape([-1, len_dec_seq]) return dec_partial_seq - def predict_word(dec_seq, enc_output, n_active_inst, n_bm, - memory_key_padding_mask): - dec_seq = paddle.transpose(self.embedding(dec_seq), [1, 0, 2]) + def predict_word(dec_seq, enc_output, n_active_inst, n_bm): + dec_seq = self.embedding(dec_seq) dec_seq = self.positional_encoding(dec_seq) tgt_mask = self.generate_square_subsequent_mask( - paddle.shape(dec_seq)[0]) - dec_output = self.decoder( - dec_seq, - enc_output, - tgt_mask=tgt_mask, - tgt_key_padding_mask=None, - memory_key_padding_mask=memory_key_padding_mask, ) - dec_output = paddle.transpose(dec_output, [1, 0, 2]) + paddle.shape(dec_seq)[1]) + tgt = dec_seq + for decoder_layer in self.decoder: + tgt = decoder_layer(tgt, enc_output, self_mask=tgt_mask) + dec_output = tgt dec_output = dec_output[:, -1, :] # Pick the last step: (bh * bm) * d_h word_prob = F.softmax(self.tgt_word_prj(dec_output), axis=1) @@ -281,8 +269,7 @@ class Transformer(nn.Layer): n_active_inst = len(inst_idx_to_position_map) dec_seq = prepare_beam_dec_seq(inst_dec_beams, len_dec_seq) - word_prob = predict_word(dec_seq, enc_output, n_active_inst, n_bm, - None) + word_prob = predict_word(dec_seq, enc_output, n_active_inst, n_bm) # Update the beam with predicted word prob information and collect incomplete instances active_inst_idx_list = collect_active_inst_idx_list( inst_dec_beams, word_prob, inst_idx_to_position_map) @@ -303,10 +290,10 @@ class Transformer(nn.Layer): with paddle.no_grad(): #-- Encode if self.encoder is not None: - src = self.positional_encoding(images.transpose([1, 0, 2])) + src = self.positional_encoding(images) src_enc = self.encoder(src) else: - src_enc = images.squeeze(2).transpose([0, 2, 1]) + src_enc = images n_bm = self.beam_size src_shape = paddle.shape(src_enc) @@ -317,11 +304,11 @@ class Transformer(nn.Layer): inst_idx_to_position_map = get_inst_idx_to_tensor_position_map( active_inst_idx_list) # Decode - for len_dec_seq in range(1, 25): + for len_dec_seq in range(1, self.max_len): src_enc_copy = src_enc.clone() active_inst_idx_list = beam_decode_step( inst_dec_beams, len_dec_seq, src_enc_copy, - inst_idx_to_position_map, n_bm, None) + inst_idx_to_position_map, n_bm) if not active_inst_idx_list: break # all instances have finished their path to src_enc, inst_idx_to_position_map = collate_active_info( @@ -354,261 +341,124 @@ class Transformer(nn.Layer): shape=[sz, sz], dtype='float32', fill_value='-inf'), diagonal=1) mask = mask + mask_inf - return mask - - def generate_padding_mask(self, x): - padding_mask = paddle.equal(x, paddle.to_tensor(0, dtype=x.dtype)) - return padding_mask - - def _reset_parameters(self): - """Initiate parameters in the transformer model.""" - - for p in self.parameters(): - if p.dim() > 1: - xavier_uniform_(p) + return mask.unsqueeze([0, 1]) -class TransformerEncoder(nn.Layer): - """TransformerEncoder is a stack of N encoder layers - Args: - encoder_layer: an instance of the TransformerEncoderLayer() class (required). - num_layers: the number of sub-encoder-layers in the encoder (required). - norm: the layer normalization component (optional). - """ +class MultiheadAttention(nn.Layer): + """Allows the model to jointly attend to information + from different representation subspaces. + See reference: Attention Is All You Need - def __init__(self, encoder_layer, num_layers): - super(TransformerEncoder, self).__init__() - self.layers = _get_clones(encoder_layer, num_layers) - self.num_layers = num_layers - - def forward(self, src): - """Pass the input through the endocder layers in turn. - Args: - src: the sequnce to the encoder (required). - mask: the mask for the src sequence (optional). - src_key_padding_mask: the mask for the src keys per batch (optional). - """ - output = src - - for i in range(self.num_layers): - output = self.layers[i](output, - src_mask=None, - src_key_padding_mask=None) - - return output - - -class TransformerDecoder(nn.Layer): - """TransformerDecoder is a stack of N decoder layers + .. math:: + \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O + \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) Args: - decoder_layer: an instance of the TransformerDecoderLayer() class (required). - num_layers: the number of sub-decoder-layers in the decoder (required). - norm: the layer normalization component (optional). + embed_dim: total dimension of the model + num_heads: parallel attention layers, or heads """ - def __init__(self, decoder_layer, num_layers): - super(TransformerDecoder, self).__init__() - self.layers = _get_clones(decoder_layer, num_layers) - self.num_layers = num_layers + def __init__(self, embed_dim, num_heads, dropout=0., self_attn=False): + super(MultiheadAttention, self).__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + # self.dropout = dropout + self.head_dim = embed_dim // num_heads + assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" + self.scale = self.head_dim**-0.5 + self.self_attn = self_attn + if self_attn: + self.qkv = nn.Linear(embed_dim, embed_dim * 3) + else: + self.q = nn.Linear(embed_dim, embed_dim) + self.kv = nn.Linear(embed_dim, embed_dim * 2) + self.attn_drop = nn.Dropout(dropout) + self.out_proj = nn.Linear(embed_dim, embed_dim) - def forward(self, - tgt, - memory, - tgt_mask=None, - memory_mask=None, - tgt_key_padding_mask=None, - memory_key_padding_mask=None): - """Pass the inputs (and mask) through the decoder layer in turn. + def forward(self, query, key=None, attn_mask=None): - Args: - tgt: the sequence to the decoder (required). - memory: the sequnce from the last layer of the encoder (required). - tgt_mask: the mask for the tgt sequence (optional). - memory_mask: the mask for the memory sequence (optional). - tgt_key_padding_mask: the mask for the tgt keys per batch (optional). - memory_key_padding_mask: the mask for the memory keys per batch (optional). - """ - output = tgt - for i in range(self.num_layers): - output = self.layers[i]( - output, - memory, - tgt_mask=tgt_mask, - memory_mask=memory_mask, - tgt_key_padding_mask=tgt_key_padding_mask, - memory_key_padding_mask=memory_key_padding_mask) + qN = query.shape[1] - return output + if self.self_attn: + qkv = self.qkv(query).reshape( + (0, qN, 3, self.num_heads, self.head_dim)).transpose( + (2, 0, 3, 1, 4)) + q, k, v = qkv[0], qkv[1], qkv[2] + else: + kN = key.shape[1] + q = self.q(query).reshape( + [0, qN, self.num_heads, self.head_dim]).transpose([0, 2, 1, 3]) + kv = self.kv(key).reshape( + (0, kN, 2, self.num_heads, self.head_dim)).transpose( + (2, 0, 3, 1, 4)) + k, v = kv[0], kv[1] + + attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale + + if attn_mask is not None: + attn += attn_mask + + attn = F.softmax(attn, axis=-1) + attn = self.attn_drop(attn) + + x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape( + (0, qN, self.embed_dim)) + x = self.out_proj(x) + + return x -class TransformerEncoderLayer(nn.Layer): - """TransformerEncoderLayer is made up of self-attn and feedforward network. - This standard encoder layer is based on the paper "Attention Is All You Need". - Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, - Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in - Neural Information Processing Systems, pages 6000-6010. Users may modify or implement - in a different way during application. - - Args: - d_model: the number of expected features in the input (required). - nhead: the number of heads in the multiheadattention models (required). - dim_feedforward: the dimension of the feedforward network model (default=2048). - dropout: the dropout value (default=0.1). - - """ - +class TransformerBlock(nn.Layer): def __init__(self, d_model, nhead, dim_feedforward=2048, attention_dropout_rate=0.0, - residual_dropout_rate=0.1): - super(TransformerEncoderLayer, self).__init__() - self.self_attn = MultiheadAttention( - d_model, nhead, dropout=attention_dropout_rate) + residual_dropout_rate=0.1, + with_self_attn=True, + with_cross_attn=False, + epsilon=1e-5): + super(TransformerBlock, self).__init__() + self.with_self_attn = with_self_attn + if with_self_attn: + self.self_attn = MultiheadAttention( + d_model, + nhead, + dropout=attention_dropout_rate, + self_attn=with_self_attn) + self.norm1 = LayerNorm(d_model, epsilon=epsilon) + self.dropout1 = Dropout(residual_dropout_rate) + self.with_cross_attn = with_cross_attn + if with_cross_attn: + self.cross_attn = MultiheadAttention( #for self_attn of encoder or cross_attn of decoder + d_model, + nhead, + dropout=attention_dropout_rate) + self.norm2 = LayerNorm(d_model, epsilon=epsilon) + self.dropout2 = Dropout(residual_dropout_rate) - self.conv1 = Conv2D( - in_channels=d_model, - out_channels=dim_feedforward, - kernel_size=(1, 1)) - self.conv2 = Conv2D( - in_channels=dim_feedforward, - out_channels=d_model, - kernel_size=(1, 1)) + self.mlp = Mlp(in_features=d_model, + hidden_features=dim_feedforward, + act_layer=nn.ReLU, + drop=residual_dropout_rate) - self.norm1 = LayerNorm(d_model) - self.norm2 = LayerNorm(d_model) - self.dropout1 = Dropout(residual_dropout_rate) - self.dropout2 = Dropout(residual_dropout_rate) + self.norm3 = LayerNorm(d_model, epsilon=epsilon) - def forward(self, src, src_mask=None, src_key_padding_mask=None): - """Pass the input through the endocder layer. - Args: - src: the sequnce to the encoder layer (required). - src_mask: the mask for the src sequence (optional). - src_key_padding_mask: the mask for the src keys per batch (optional). - """ - src2 = self.self_attn( - src, - src, - src, - attn_mask=src_mask, - key_padding_mask=src_key_padding_mask) - src = src + self.dropout1(src2) - src = self.norm1(src) - - src = paddle.transpose(src, [1, 2, 0]) - src = paddle.unsqueeze(src, 2) - src2 = self.conv2(F.relu(self.conv1(src))) - src2 = paddle.squeeze(src2, 2) - src2 = paddle.transpose(src2, [2, 0, 1]) - src = paddle.squeeze(src, 2) - src = paddle.transpose(src, [2, 0, 1]) - - src = src + self.dropout2(src2) - src = self.norm2(src) - return src - - -class TransformerDecoderLayer(nn.Layer): - """TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network. - This standard decoder layer is based on the paper "Attention Is All You Need". - Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, - Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in - Neural Information Processing Systems, pages 6000-6010. Users may modify or implement - in a different way during application. - - Args: - d_model: the number of expected features in the input (required). - nhead: the number of heads in the multiheadattention models (required). - dim_feedforward: the dimension of the feedforward network model (default=2048). - dropout: the dropout value (default=0.1). - - """ - - def __init__(self, - d_model, - nhead, - dim_feedforward=2048, - attention_dropout_rate=0.0, - residual_dropout_rate=0.1): - super(TransformerDecoderLayer, self).__init__() - self.self_attn = MultiheadAttention( - d_model, nhead, dropout=attention_dropout_rate) - self.multihead_attn = MultiheadAttention( - d_model, nhead, dropout=attention_dropout_rate) - - self.conv1 = Conv2D( - in_channels=d_model, - out_channels=dim_feedforward, - kernel_size=(1, 1)) - self.conv2 = Conv2D( - in_channels=dim_feedforward, - out_channels=d_model, - kernel_size=(1, 1)) - - self.norm1 = LayerNorm(d_model) - self.norm2 = LayerNorm(d_model) - self.norm3 = LayerNorm(d_model) - self.dropout1 = Dropout(residual_dropout_rate) - self.dropout2 = Dropout(residual_dropout_rate) self.dropout3 = Dropout(residual_dropout_rate) - def forward(self, - tgt, - memory, - tgt_mask=None, - memory_mask=None, - tgt_key_padding_mask=None, - memory_key_padding_mask=None): - """Pass the inputs (and mask) through the decoder layer. + def forward(self, tgt, memory=None, self_mask=None, cross_mask=None): + if self.with_self_attn: + tgt1 = self.self_attn(tgt, attn_mask=self_mask) + tgt = self.norm1(tgt + self.dropout1(tgt1)) - Args: - tgt: the sequence to the decoder layer (required). - memory: the sequnce from the last layer of the encoder (required). - tgt_mask: the mask for the tgt sequence (optional). - memory_mask: the mask for the memory sequence (optional). - tgt_key_padding_mask: the mask for the tgt keys per batch (optional). - memory_key_padding_mask: the mask for the memory keys per batch (optional). - - """ - tgt2 = self.self_attn( - tgt, - tgt, - tgt, - attn_mask=tgt_mask, - key_padding_mask=tgt_key_padding_mask) - tgt = tgt + self.dropout1(tgt2) - tgt = self.norm1(tgt) - tgt2 = self.multihead_attn( - tgt, - memory, - memory, - attn_mask=memory_mask, - key_padding_mask=memory_key_padding_mask) - tgt = tgt + self.dropout2(tgt2) - tgt = self.norm2(tgt) - - # default - tgt = paddle.transpose(tgt, [1, 2, 0]) - tgt = paddle.unsqueeze(tgt, 2) - tgt2 = self.conv2(F.relu(self.conv1(tgt))) - tgt2 = paddle.squeeze(tgt2, 2) - tgt2 = paddle.transpose(tgt2, [2, 0, 1]) - tgt = paddle.squeeze(tgt, 2) - tgt = paddle.transpose(tgt, [2, 0, 1]) - - tgt = tgt + self.dropout3(tgt2) - tgt = self.norm3(tgt) + if self.with_cross_attn: + tgt2 = self.cross_attn(tgt, key=memory, attn_mask=cross_mask) + tgt = self.norm2(tgt + self.dropout2(tgt2)) + tgt = self.norm3(tgt + self.dropout3(self.mlp(tgt))) return tgt -def _get_clones(module, N): - return LayerList([copy.deepcopy(module) for i in range(N)]) - - class PositionalEncoding(nn.Layer): """Inject some information about the relative or absolute position of the tokens in the sequence. The positional encodings have the same dimension as @@ -651,8 +501,9 @@ class PositionalEncoding(nn.Layer): Examples: >>> output = pos_encoder(x) """ + x = x.transpose([1, 0, 2]) x = x + self.pe[:paddle.shape(x)[0], :] - return self.dropout(x) + return self.dropout(x).transpose([1, 0, 2]) class PositionalEncoding_2d(nn.Layer): @@ -725,7 +576,7 @@ class PositionalEncoding_2d(nn.Layer): class Embeddings(nn.Layer): - def __init__(self, d_model, vocab, padding_idx, scale_embedding): + def __init__(self, d_model, vocab, padding_idx=None, scale_embedding=True): super(Embeddings, self).__init__() self.embedding = nn.Embedding(vocab, d_model, padding_idx=padding_idx) w0 = np.random.normal(0.0, d_model**-0.5, @@ -742,7 +593,7 @@ class Embeddings(nn.Layer): class Beam(): - ''' Beam search ''' + """ Beam search """ def __init__(self, size, device=False): diff --git a/ppocr/postprocess/__init__.py b/ppocr/postprocess/__init__.py index 4f900ee1fc..2635117c84 100644 --- a/ppocr/postprocess/__init__.py +++ b/ppocr/postprocess/__init__.py @@ -27,7 +27,7 @@ from .sast_postprocess import SASTPostProcess from .fce_postprocess import FCEPostProcess from .rec_postprocess import CTCLabelDecode, AttnLabelDecode, SRNLabelDecode, \ DistillationCTCLabelDecode, TableLabelDecode, NRTRLabelDecode, SARLabelDecode, \ - SEEDLabelDecode, PRENLabelDecode, ViTSTRLabelDecode + SEEDLabelDecode, PRENLabelDecode, ViTSTRLabelDecode, ABINetLabelDecode from .cls_postprocess import ClsPostProcess from .pg_postprocess import PGPostProcess from .vqa_token_ser_layoutlm_postprocess import VQASerTokenLayoutLMPostProcess @@ -42,7 +42,7 @@ def build_post_process(config, global_config=None): 'DistillationDBPostProcess', 'NRTRLabelDecode', 'SARLabelDecode', 'SEEDLabelDecode', 'VQASerTokenLayoutLMPostProcess', 'VQAReTokenLayoutLMPostProcess', 'PRENLabelDecode', - 'DistillationSARLabelDecode', 'ViTSTRLabelDecode' + 'DistillationSARLabelDecode', 'ViTSTRLabelDecode', 'ABINetLabelDecode' ] if config['name'] == 'PSEPostProcess': diff --git a/ppocr/postprocess/rec_postprocess.py b/ppocr/postprocess/rec_postprocess.py index df6203fada..c77420ad19 100644 --- a/ppocr/postprocess/rec_postprocess.py +++ b/ppocr/postprocess/rec_postprocess.py @@ -140,96 +140,6 @@ class DistillationCTCLabelDecode(CTCLabelDecode): return output -class NRTRLabelDecode(BaseRecLabelDecode): - """ Convert between text-label and text-index """ - - def __init__(self, character_dict_path=None, use_space_char=True, **kwargs): - super(NRTRLabelDecode, self).__init__(character_dict_path, - use_space_char) - - def __call__(self, preds, label=None, *args, **kwargs): - - if len(preds) == 2: - preds_id = preds[0] - preds_prob = preds[1] - if isinstance(preds_id, paddle.Tensor): - preds_id = preds_id.numpy() - if isinstance(preds_prob, paddle.Tensor): - preds_prob = preds_prob.numpy() - if preds_id[0][0] == 2: - preds_idx = preds_id[:, 1:] - preds_prob = preds_prob[:, 1:] - else: - preds_idx = preds_id - text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) - if label is None: - return text - label = self.decode(label[:, 1:]) - else: - if isinstance(preds, paddle.Tensor): - preds = preds.numpy() - preds_idx = preds.argmax(axis=2) - preds_prob = preds.max(axis=2) - text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) - if label is None: - return text - label = self.decode(label[:, 1:]) - return text, label - - def add_special_char(self, dict_character): - dict_character = ['blank', '', '', ''] + dict_character - return dict_character - - def decode(self, text_index, text_prob=None, is_remove_duplicate=False): - """ convert text-index into text-label. """ - result_list = [] - batch_size = len(text_index) - for batch_idx in range(batch_size): - char_list = [] - conf_list = [] - for idx in range(len(text_index[batch_idx])): - try: - char_idx = self.character[int(text_index[batch_idx][idx])] - except: - continue - if char_idx == '': # end - break - char_list.append(char_idx) - if text_prob is not None: - conf_list.append(text_prob[batch_idx][idx]) - else: - conf_list.append(1) - text = ''.join(char_list) - result_list.append((text.lower(), np.mean(conf_list).tolist())) - return result_list - - -class ViTSTRLabelDecode(NRTRLabelDecode): - """ Convert between text-label and text-index """ - - def __init__(self, character_dict_path=None, use_space_char=False, - **kwargs): - super(ViTSTRLabelDecode, self).__init__(character_dict_path, - use_space_char) - - def __call__(self, preds, label=None, *args, **kwargs): - if isinstance(preds, paddle.Tensor): - preds = preds[:, 1:].numpy() - else: - preds = preds[:, 1:] - preds_idx = preds.argmax(axis=2) - preds_prob = preds.max(axis=2) - text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) - if label is None: - return text - label = self.decode(label[:, 1:]) - return text, label - - def add_special_char(self, dict_character): - dict_character = ['', ''] + dict_character - return dict_character - - class AttnLabelDecode(BaseRecLabelDecode): """ Convert between text-label and text-index """ @@ -778,3 +688,122 @@ class PRENLabelDecode(BaseRecLabelDecode): return text label = self.decode(label) return text, label + + +class NRTRLabelDecode(BaseRecLabelDecode): + """ Convert between text-label and text-index """ + + def __init__(self, character_dict_path=None, use_space_char=True, **kwargs): + super(NRTRLabelDecode, self).__init__(character_dict_path, + use_space_char) + + def __call__(self, preds, label=None, *args, **kwargs): + + if len(preds) == 2: + preds_id = preds[0] + preds_prob = preds[1] + if isinstance(preds_id, paddle.Tensor): + preds_id = preds_id.numpy() + if isinstance(preds_prob, paddle.Tensor): + preds_prob = preds_prob.numpy() + if preds_id[0][0] == 2: + preds_idx = preds_id[:, 1:] + preds_prob = preds_prob[:, 1:] + else: + preds_idx = preds_id + text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) + if label is None: + return text + label = self.decode(label[:, 1:]) + else: + if isinstance(preds, paddle.Tensor): + preds = preds.numpy() + preds_idx = preds.argmax(axis=2) + preds_prob = preds.max(axis=2) + text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) + if label is None: + return text + label = self.decode(label[:, 1:]) + return text, label + + def add_special_char(self, dict_character): + dict_character = ['blank', '', '', ''] + dict_character + return dict_character + + def decode(self, text_index, text_prob=None, is_remove_duplicate=False): + """ convert text-index into text-label. """ + result_list = [] + batch_size = len(text_index) + for batch_idx in range(batch_size): + char_list = [] + conf_list = [] + for idx in range(len(text_index[batch_idx])): + try: + char_idx = self.character[int(text_index[batch_idx][idx])] + except: + continue + if char_idx == '': # end + break + char_list.append(char_idx) + if text_prob is not None: + conf_list.append(text_prob[batch_idx][idx]) + else: + conf_list.append(1) + text = ''.join(char_list) + result_list.append((text.lower(), np.mean(conf_list).tolist())) + return result_list + + +class ViTSTRLabelDecode(NRTRLabelDecode): + """ Convert between text-label and text-index """ + + def __init__(self, character_dict_path=None, use_space_char=False, + **kwargs): + super(ViTSTRLabelDecode, self).__init__(character_dict_path, + use_space_char) + + def __call__(self, preds, label=None, *args, **kwargs): + if isinstance(preds, paddle.Tensor): + preds = preds[:, 1:].numpy() + else: + preds = preds[:, 1:] + preds_idx = preds.argmax(axis=2) + preds_prob = preds.max(axis=2) + text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) + if label is None: + return text + label = self.decode(label[:, 1:]) + return text, label + + def add_special_char(self, dict_character): + dict_character = ['', ''] + dict_character + return dict_character + + +class ABINetLabelDecode(NRTRLabelDecode): + """ Convert between text-label and text-index """ + + def __init__(self, character_dict_path=None, use_space_char=False, + **kwargs): + super(ABINetLabelDecode, self).__init__(character_dict_path, + use_space_char) + + def __call__(self, preds, label=None, *args, **kwargs): + if isinstance(preds, dict): + preds = preds['align'][-1].numpy() + elif isinstance(preds, paddle.Tensor): + preds = preds.numpy() + else: + preds = preds + + preds_idx = preds.argmax(axis=2) + preds_prob = preds.max(axis=2) + text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False) + if label is None: + return text + label = self.decode(label) + return text, label + + def add_special_char(self, dict_character): + dict_character = [''] + dict_character + return dict_character diff --git a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml index ba6728d836..8118d58724 100644 --- a/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml +++ b/test_tipc/configs/rec_mtb_nrtr/rec_mtb_nrtr.yml @@ -99,5 +99,5 @@ Eval: shuffle: False drop_last: False batch_size_per_card: 256 - num_workers: 1 + num_workers: 4 use_shared_memory: False diff --git a/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml b/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml new file mode 100644 index 0000000000..306c49d683 --- /dev/null +++ b/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml @@ -0,0 +1,105 @@ +Global: + use_gpu: True + epoch_num: 10 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/r45_abinet/ + save_epoch_step: 1 + # evaluation is run every 2000 iterations + eval_batch_step: [0, 2000] + cal_metric_during_train: True + pretrained_model: + checkpoints: + save_inference_dir: + use_visualdl: False + infer_img: doc/imgs_words_en/word_10.png + # for data or label process + character_dict_path: + character_type: en + max_text_length: 25 + infer_mode: False + use_space_char: False + save_res_path: ./output/rec/predicts_abinet.txt + +Optimizer: + name: Adam + beta1: 0.9 + beta2: 0.99 + clip_norm: 20.0 + lr: + name: Piecewise + decay_epochs: [6] + values: [0.0001, 0.00001] + regularizer: + name: 'L2' + factor: 0. + +Architecture: + model_type: rec + algorithm: ABINet + in_channels: 3 + Transform: + Backbone: + name: ResNet45 + + Head: + name: ABINetHead + use_lang: True + iter_size: 3 + + +Loss: + name: CELoss + ignore_index: &ignore_index 100 # Must be greater than the number of character classes + +PostProcess: + name: ABINetLabelDecode + +Metric: + name: RecMetric + main_indicator: acc + +Train: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data/ + label_file_list: ["./train_data/ic15_data/rec_gt_train.txt"] + transforms: + - DecodeImage: # load image + img_mode: RGB + channel_first: False + - ABINetLabelEncode: # Class handling label + ignore_index: *ignore_index + - ABINetRecResizeImg: + image_shape: [3, 32, 128] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: True + batch_size_per_card: 96 + drop_last: True + num_workers: 4 + +Eval: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ic15_data + label_file_list: ["./train_data/ic15_data/rec_gt_test.txt"] + transforms: + - DecodeImage: # load image + img_mode: RGB + channel_first: False + - ABINetLabelEncode: # Class handling label + ignore_index: *ignore_index + - ABINetRecResizeImg: + image_shape: [3, 32, 128] + padding: False + - KeepKeys: + keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order + loader: + shuffle: False + drop_last: False + batch_size_per_card: 256 + num_workers: 4 + use_shared_memory: False diff --git a/test_tipc/configs/rec_r45_abinet/train_infer_python.txt b/test_tipc/configs/rec_r45_abinet/train_infer_python.txt new file mode 100644 index 0000000000..ecab1bcbbd --- /dev/null +++ b/test_tipc/configs/rec_r45_abinet/train_infer_python.txt @@ -0,0 +1,53 @@ +===========================train_params=========================== +model_name:rec_abinet +python:python3.7 +gpu_list:0|0,1 +Global.use_gpu:True|True +Global.auto_cast:null +Global.epoch_num:lite_train_lite_infer=2|whole_train_whole_infer=300 +Global.save_model_dir:./output/ +Train.loader.batch_size_per_card:lite_train_lite_infer=16|whole_train_whole_infer=64 +Global.pretrained_model:null +train_model_name:latest +train_infer_img_dir:./inference/rec_inference +null:null +## +trainer:norm_train +norm_train:tools/train.py -c test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml -o +pact_train:null +fpgm_train:null +distill_train:null +null:null +null:null +## +===========================eval_params=========================== +eval:tools/eval.py -c test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml -o +null:null +## +===========================infer_params=========================== +Global.save_inference_dir:./output/ +Global.checkpoints: +norm_export:tools/export_model.py -c test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml -o +quant_export:null +fpgm_export:null +distill_export:null +export1:null +export2:null +## +train_model:./inference/rec_r45_abinet_train/best_accuracy +infer_export:tools/export_model.py -c test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml -o +infer_quant:False +inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/ic15_dict.txt --rec_image_shape="3,32,128" --rec_algorithm="ABINet" +--use_gpu:True|False +--enable_mkldnn:False +--cpu_threads:6 +--rec_batch_num:1|6 +--use_tensorrt:False +--precision:fp32 +--rec_model_dir: +--image_dir:./inference/rec_inference +--save_log_path:./test/output/ +--benchmark:True +null:null +===========================infer_benchmark_params========================== +random_infer_input:[{float32,[3,32,128]}] diff --git a/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml index 26facca34d..bffd8ac09a 100644 --- a/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml +++ b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml @@ -26,7 +26,7 @@ Optimizer: name: AdamW beta1: 0.9 beta2: 0.99 - epsilon: 0.00000008 + epsilon: 8.e-8 weight_decay: 0.05 no_weight_decay_name: norm pos_embed one_dim_param_no_weight_decay: true diff --git a/test_tipc/configs/rec_svtrnet/train_infer_python.txt b/test_tipc/configs/rec_svtrnet/train_infer_python.txt index 60fd25ad4f..a7e4a24063 100644 --- a/test_tipc/configs/rec_svtrnet/train_infer_python.txt +++ b/test_tipc/configs/rec_svtrnet/train_infer_python.txt @@ -37,7 +37,7 @@ export2:null train_model:./inference/rec_svtrnet_train/best_accuracy infer_export:tools/export_model.py -c test_tipc/configs/rec_svtrnet/rec_svtrnet.yml -o infer_quant:False -inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbol_dict.txt --rec_image_shape="3,64,256" --rec_algorithm="SVTR" +inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/ic15_dict.txt --rec_image_shape="3,64,256" --rec_algorithm="SVTR" --use_gpu:True|False --enable_mkldnn:False --cpu_threads:6 @@ -50,4 +50,4 @@ inference:tools/infer/predict_rec.py --rec_char_dict_path=./ppocr/utils/EN_symbo --benchmark:True null:null ===========================infer_benchmark_params========================== -random_infer_input:[{float32,[1=3,64,256]}] +random_infer_input:[{float32,[3,64,256]}] diff --git a/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml index 4e8151f599..a0aed48875 100644 --- a/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml +++ b/test_tipc/configs/rec_vitstr_none_ce/rec_vitstr_none_ce.yml @@ -23,7 +23,7 @@ Global: Optimizer: name: Adadelta - epsilon: 0.00000001 + epsilon: 1.e-8 rho: 0.95 clip_norm: 5.0 lr: @@ -46,6 +46,7 @@ Loss: name: CELoss smoothing: False with_all: True + ignore_index: &ignore_index 0 # Must be zero or greater than the number of character classes PostProcess: name: ViTSTRLabelDecode @@ -64,6 +65,7 @@ Train: img_mode: BGR channel_first: False - ViTSTRLabelEncode: # Class handling label + ignore_index: *ignore_index - GrayRecResizeImg: image_shape: [224, 224] # W H resize_type: PIL # PIL or OpenCV @@ -87,6 +89,7 @@ Eval: img_mode: BGR channel_first: False - ViTSTRLabelEncode: # Class handling label + ignore_index: *ignore_index - GrayRecResizeImg: image_shape: [224, 224] # W H resize_type: PIL # PIL or OpenCV diff --git a/tools/export_model.py b/tools/export_model.py index 7ef2c36b12..e2673239b1 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -79,6 +79,19 @@ def export_single_model(model, arch_config, save_path, logger, quanter=None): shape=[None, 1, 224, 224], dtype="float32"), ] model = to_static(model, input_spec=other_shape) + elif arch_config["algorithm"] == "ABINet": + other_shape = [ + paddle.static.InputSpec( + shape=[None, 3, 32, 128], dtype="float32"), + ] + # print([None, 3, 32, 128]) + model = to_static(model, input_spec=other_shape) + elif arch_config["algorithm"] == "NRTR": + other_shape = [ + paddle.static.InputSpec( + shape=[None, 1, 32, 100], dtype="float32"), + ] + model = to_static(model, input_spec=other_shape) else: infer_shape = [3, -1, -1] if arch_config["model_type"] == "rec": @@ -90,8 +103,6 @@ def export_single_model(model, arch_config, save_path, logger, quanter=None): "When there is tps in the network, variable length input is not supported, and the input size needs to be the same as during training" ) infer_shape[-1] = 100 - if arch_config["algorithm"] == "NRTR": - infer_shape = [1, 32, 100] elif arch_config["model_type"] == "table": infer_shape = [3, 488, 488] model = to_static( diff --git a/tools/infer/predict_rec.py b/tools/infer/predict_rec.py index 1945667972..a95f555966 100755 --- a/tools/infer/predict_rec.py +++ b/tools/infer/predict_rec.py @@ -75,6 +75,12 @@ class TextRecognizer(object): "character_dict_path": args.rec_char_dict_path, "use_space_char": args.use_space_char } + elif self.rec_algorithm == 'ABINet': + postprocess_params = { + 'name': 'ABINetLabelDecode', + "character_dict_path": args.rec_char_dict_path, + "use_space_char": args.use_space_char + } self.postprocess_op = build_post_process(postprocess_params) self.predictor, self.input_tensor, self.output_tensors, self.config = \ utility.create_predictor(args, 'rec', logger) @@ -145,17 +151,6 @@ class TextRecognizer(object): padding_im[:, :, 0:resized_w] = resized_image return padding_im - def resize_norm_img_svtr(self, img, image_shape): - - imgC, imgH, imgW = image_shape - resized_image = cv2.resize( - img, (imgW, imgH), interpolation=cv2.INTER_LINEAR) - resized_image = resized_image.astype('float32') - resized_image = resized_image.transpose((2, 0, 1)) / 255 - resized_image -= 0.5 - resized_image /= 0.5 - return resized_image - def resize_norm_img_srn(self, img, image_shape): imgC, imgH, imgW = image_shape @@ -263,6 +258,35 @@ class TextRecognizer(object): return padding_im, resize_shape, pad_shape, valid_ratio + def resize_norm_img_svtr(self, img, image_shape): + + imgC, imgH, imgW = image_shape + resized_image = cv2.resize( + img, (imgW, imgH), interpolation=cv2.INTER_LINEAR) + resized_image = resized_image.astype('float32') + resized_image = resized_image.transpose((2, 0, 1)) / 255 + resized_image -= 0.5 + resized_image /= 0.5 + return resized_image + + def resize_norm_img_abinet(self, img, image_shape): + + imgC, imgH, imgW = image_shape + + resized_image = cv2.resize( + img, (imgW, imgH), interpolation=cv2.INTER_LINEAR) + resized_image = resized_image.astype('float32') + resized_image = resized_image / 255. + + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + resized_image = ( + resized_image - mean[None, None, ...]) / std[None, None, ...] + resized_image = resized_image.transpose((2, 0, 1)) + resized_image = resized_image.astype('float32') + + return resized_image + def __call__(self, img_list): img_num = len(img_list) # Calculate the aspect ratio of all text bars @@ -313,6 +337,11 @@ class TextRecognizer(object): self.rec_image_shape) norm_img = norm_img[np.newaxis, :] norm_img_batch.append(norm_img) + elif self.rec_algorithm == "ABINet": + norm_img = self.resize_norm_img_abinet( + img_list[indices[ino]], self.rec_image_shape) + norm_img = norm_img[np.newaxis, :] + norm_img_batch.append(norm_img) else: norm_img = self.resize_norm_img(img_list[indices[ino]], max_wh_ratio) diff --git a/tools/program.py b/tools/program.py index 745c28b872..f61d3f9af7 100755 --- a/tools/program.py +++ b/tools/program.py @@ -575,7 +575,7 @@ def preprocess(is_train=False): 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', - 'ViTSTR' + 'ViTSTR', 'ABINet' ] if use_xpu: From 37f22e3e33400599fa0856777b94ab6af0e2f31b Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Thu, 23 Jun 2022 05:03:15 +0000 Subject: [PATCH 07/43] add rec_resnet45 --- ppocr/modeling/backbones/__init__.py | 3 +- ppocr/modeling/backbones/rec_resnet.py | 280 ---------------------- ppocr/modeling/backbones/rec_resnet_45.py | 147 ++++++++++++ 3 files changed, 149 insertions(+), 281 deletions(-) delete mode 100644 ppocr/modeling/backbones/rec_resnet.py create mode 100644 ppocr/modeling/backbones/rec_resnet_45.py diff --git a/ppocr/modeling/backbones/__init__.py b/ppocr/modeling/backbones/__init__.py index 56d2c2b9db..f8959e263e 100755 --- a/ppocr/modeling/backbones/__init__.py +++ b/ppocr/modeling/backbones/__init__.py @@ -27,7 +27,8 @@ def build_backbone(config, model_type): from .rec_resnet_fpn import ResNetFPN from .rec_mv1_enhance import MobileNetV1Enhance from .rec_nrtr_mtb import MTB - from .rec_resnet import ResNet31, ResNet45 + from .rec_resnet_31 import ResNet31 + from .rec_resnet_45 import ResNet45 from .rec_resnet_aster import ResNet_ASTER from .rec_micronet import MicroNet from .rec_efficientb3_pren import EfficientNetb3_PREN diff --git a/ppocr/modeling/backbones/rec_resnet.py b/ppocr/modeling/backbones/rec_resnet.py deleted file mode 100644 index ce76b642cc..0000000000 --- a/ppocr/modeling/backbones/rec_resnet.py +++ /dev/null @@ -1,280 +0,0 @@ -# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -This code is refer from: -https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/layers/conv_layer.py -https://github.com/open-mmlab/mmocr/blob/main/mmocr/models/textrecog/backbones/resnet31_ocr.py -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import paddle -from paddle import ParamAttr -from paddle.nn.initializer import KaimingNormal -import paddle.nn as nn -import paddle.nn.functional as F -import numpy as np -import math - -__all__ = ["ResNet31", "ResNet45"] - - -def conv1x1(in_planes, out_planes, stride=1): - return nn.Conv2D( - in_planes, - out_planes, - kernel_size=1, - stride=stride, - weight_attr=ParamAttr(initializer=KaimingNormal()), - bias_attr=False) - - -def conv3x3(in_channel, out_channel, stride=1): - return nn.Conv2D( - in_channel, - out_channel, - kernel_size=3, - stride=stride, - padding=1, - weight_attr=ParamAttr(initializer=KaimingNormal()), - bias_attr=False) - - -class BasicBlock(nn.Layer): - expansion = 1 - - def __init__(self, in_channels, channels, stride=1, downsample=None): - super().__init__() - self.conv1 = conv1x1(in_channels, channels) - self.bn1 = nn.BatchNorm2D(channels) - self.relu = nn.ReLU() - self.conv2 = conv3x3(channels, channels, stride) - self.bn2 = nn.BatchNorm2D(channels) - self.downsample = downsample - self.stride = stride - - def forward(self, x): - residual = x - - out = self.conv1(x) - out = self.bn1(out) - out = self.relu(out) - - out = self.conv2(out) - out = self.bn2(out) - - if self.downsample is not None: - residual = self.downsample(x) - out += residual - out = self.relu(out) - - return out - - -class ResNet31(nn.Layer): - ''' - Args: - in_channels (int): Number of channels of input image tensor. - layers (list[int]): List of BasicBlock number for each stage. - channels (list[int]): List of out_channels of Conv2d layer. - out_indices (None | Sequence[int]): Indices of output stages. - last_stage_pool (bool): If True, add `MaxPool2d` layer to last stage. - ''' - - def __init__(self, - in_channels=3, - layers=[1, 2, 5, 3], - channels=[64, 128, 256, 256, 512, 512, 512], - out_indices=None, - last_stage_pool=False): - super(ResNet31, self).__init__() - assert isinstance(in_channels, int) - assert isinstance(last_stage_pool, bool) - - self.out_indices = out_indices - self.last_stage_pool = last_stage_pool - - # conv 1 (Conv Conv) - self.conv1_1 = nn.Conv2D( - in_channels, channels[0], kernel_size=3, stride=1, padding=1) - self.bn1_1 = nn.BatchNorm2D(channels[0]) - self.relu1_1 = nn.ReLU() - - self.conv1_2 = nn.Conv2D( - channels[0], channels[1], kernel_size=3, stride=1, padding=1) - self.bn1_2 = nn.BatchNorm2D(channels[1]) - self.relu1_2 = nn.ReLU() - - # conv 2 (Max-pooling, Residual block, Conv) - self.pool2 = nn.MaxPool2D( - kernel_size=2, stride=2, padding=0, ceil_mode=True) - self.block2 = self._make_layer(channels[1], channels[2], layers[0]) - self.conv2 = nn.Conv2D( - channels[2], channels[2], kernel_size=3, stride=1, padding=1) - self.bn2 = nn.BatchNorm2D(channels[2]) - self.relu2 = nn.ReLU() - - # conv 3 (Max-pooling, Residual block, Conv) - self.pool3 = nn.MaxPool2D( - kernel_size=2, stride=2, padding=0, ceil_mode=True) - self.block3 = self._make_layer(channels[2], channels[3], layers[1]) - self.conv3 = nn.Conv2D( - channels[3], channels[3], kernel_size=3, stride=1, padding=1) - self.bn3 = nn.BatchNorm2D(channels[3]) - self.relu3 = nn.ReLU() - - # conv 4 (Max-pooling, Residual block, Conv) - self.pool4 = nn.MaxPool2D( - kernel_size=(2, 1), stride=(2, 1), padding=0, ceil_mode=True) - self.block4 = self._make_layer(channels[3], channels[4], layers[2]) - self.conv4 = nn.Conv2D( - channels[4], channels[4], kernel_size=3, stride=1, padding=1) - self.bn4 = nn.BatchNorm2D(channels[4]) - self.relu4 = nn.ReLU() - - # conv 5 ((Max-pooling), Residual block, Conv) - self.pool5 = None - if self.last_stage_pool: - self.pool5 = nn.MaxPool2D( - kernel_size=2, stride=2, padding=0, ceil_mode=True) - self.block5 = self._make_layer(channels[4], channels[5], layers[3]) - self.conv5 = nn.Conv2D( - channels[5], channels[5], kernel_size=3, stride=1, padding=1) - self.bn5 = nn.BatchNorm2D(channels[5]) - self.relu5 = nn.ReLU() - - self.out_channels = channels[-1] - - def _make_layer(self, input_channels, output_channels, blocks): - layers = [] - for _ in range(blocks): - downsample = None - if input_channels != output_channels: - downsample = nn.Sequential( - nn.Conv2D( - input_channels, - output_channels, - kernel_size=1, - stride=1, - weight_attr=ParamAttr(initializer=KaimingNormal()), - bias_attr=False), - nn.BatchNorm2D(output_channels), ) - - layers.append( - BasicBlock( - input_channels, output_channels, downsample=downsample)) - input_channels = output_channels - return nn.Sequential(*layers) - - def forward(self, x): - x = self.conv1_1(x) - x = self.bn1_1(x) - x = self.relu1_1(x) - - x = self.conv1_2(x) - x = self.bn1_2(x) - x = self.relu1_2(x) - - outs = [] - for i in range(4): - layer_index = i + 2 - pool_layer = getattr(self, f'pool{layer_index}') - block_layer = getattr(self, f'block{layer_index}') - conv_layer = getattr(self, f'conv{layer_index}') - bn_layer = getattr(self, f'bn{layer_index}') - relu_layer = getattr(self, f'relu{layer_index}') - - if pool_layer is not None: - x = pool_layer(x) - x = block_layer(x) - x = conv_layer(x) - x = bn_layer(x) - x = relu_layer(x) - - outs.append(x) - - if self.out_indices is not None: - return tuple([outs[i] for i in self.out_indices]) - - return x - - -class ResNet(nn.Layer): - def __init__(self, block, layers, in_channels=3): - self.inplanes = 32 - super(ResNet, self).__init__() - self.conv1 = nn.Conv2D( - 3, - 32, - kernel_size=3, - stride=1, - padding=1, - weight_attr=ParamAttr(initializer=KaimingNormal()), - bias_attr=False) - self.bn1 = nn.BatchNorm2D(32) - self.relu = nn.ReLU() - - self.layer1 = self._make_layer(block, 32, layers[0], stride=2) - self.layer2 = self._make_layer(block, 64, layers[1], stride=1) - self.layer3 = self._make_layer(block, 128, layers[2], stride=2) - self.layer4 = self._make_layer(block, 256, layers[3], stride=1) - self.layer5 = self._make_layer(block, 512, layers[4], stride=1) - self.out_channels = 512 - - # for m in self.modules(): - # if isinstance(m, nn.Conv2D): - # n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels - # m.weight.data.normal_(0, math.sqrt(2. / n)) - - def _make_layer(self, block, planes, blocks, stride=1): - downsample = None - if stride != 1 or self.inplanes != planes * block.expansion: - # downsample = True - downsample = nn.Sequential( - nn.Conv2D( - self.inplanes, - planes * block.expansion, - kernel_size=1, - stride=stride, - weight_attr=ParamAttr(initializer=KaimingNormal()), - bias_attr=False), - nn.BatchNorm2D(planes * block.expansion), ) - - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) - self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - - return nn.Sequential(*layers) - - def forward(self, x): - - x = self.conv1(x) - x = self.bn1(x) - x = self.relu(x) - # print(x) - x = self.layer1(x) - x = self.layer2(x) - x = self.layer3(x) - # print(x) - x = self.layer4(x) - x = self.layer5(x) - return x - - -def ResNet45(in_channels=3): - return ResNet(BasicBlock, [3, 4, 6, 6, 3], in_channels=in_channels) diff --git a/ppocr/modeling/backbones/rec_resnet_45.py b/ppocr/modeling/backbones/rec_resnet_45.py new file mode 100644 index 0000000000..9093d0bc99 --- /dev/null +++ b/ppocr/modeling/backbones/rec_resnet_45.py @@ -0,0 +1,147 @@ +# copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This code is refer from: +https://github.com/FangShancheng/ABINet/tree/main/modules +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import paddle +from paddle import ParamAttr +from paddle.nn.initializer import KaimingNormal +import paddle.nn as nn +import paddle.nn.functional as F +import numpy as np +import math + +__all__ = ["ResNet45"] + + +def conv1x1(in_planes, out_planes, stride=1): + return nn.Conv2D( + in_planes, + out_planes, + kernel_size=1, + stride=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + + +def conv3x3(in_channel, out_channel, stride=1): + return nn.Conv2D( + in_channel, + out_channel, + kernel_size=3, + stride=stride, + padding=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + + +class BasicBlock(nn.Layer): + expansion = 1 + + def __init__(self, in_channels, channels, stride=1, downsample=None): + super().__init__() + self.conv1 = conv1x1(in_channels, channels) + self.bn1 = nn.BatchNorm2D(channels) + self.relu = nn.ReLU() + self.conv2 = conv3x3(channels, channels, stride) + self.bn2 = nn.BatchNorm2D(channels) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + residual = self.downsample(x) + out += residual + out = self.relu(out) + + return out + + +class ResNet45(nn.Layer): + def __init__(self, block=BasicBlock, layers=[3, 4, 6, 6, 3], in_channels=3): + self.inplanes = 32 + super(ResNet45, self).__init__() + self.conv1 = nn.Conv2D( + 3, + 32, + kernel_size=3, + stride=1, + padding=1, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False) + self.bn1 = nn.BatchNorm2D(32) + self.relu = nn.ReLU() + + self.layer1 = self._make_layer(block, 32, layers[0], stride=2) + self.layer2 = self._make_layer(block, 64, layers[1], stride=1) + self.layer3 = self._make_layer(block, 128, layers[2], stride=2) + self.layer4 = self._make_layer(block, 256, layers[3], stride=1) + self.layer5 = self._make_layer(block, 512, layers[4], stride=1) + self.out_channels = 512 + + # for m in self.modules(): + # if isinstance(m, nn.Conv2D): + # n = m._kernel_size[0] * m._kernel_size[1] * m._out_channels + # m.weight.data.normal_(0, math.sqrt(2. / n)) + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + # downsample = True + downsample = nn.Sequential( + nn.Conv2D( + self.inplanes, + planes * block.expansion, + kernel_size=1, + stride=stride, + weight_attr=ParamAttr(initializer=KaimingNormal()), + bias_attr=False), + nn.BatchNorm2D(planes * block.expansion), ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + # print(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + # print(x) + x = self.layer4(x) + x = self.layer5(x) + return x From 0a343fd3aae14b866281d897474d0669b1b26282 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Thu, 23 Jun 2022 08:17:00 +0000 Subject: [PATCH 08/43] svtr ch large model --- configs/rec/rec_svtrnet.yml | 4 +- configs/rec/rec_svtrnet_ch.yml | 159 ++++++++++++++++++++++++++++ doc/doc_ch/algorithm_rec_svtr.md | 1 - doc/doc_en/algorithm_rec_svtr_en.md | 1 - ppocr/data/imaug/__init__.py | 3 +- ppocr/data/imaug/rec_img_aug.py | 124 +++++++++++++--------- tools/export_model.py | 19 +++- 7 files changed, 252 insertions(+), 59 deletions(-) create mode 100644 configs/rec/rec_svtrnet_ch.yml diff --git a/configs/rec/rec_svtrnet.yml b/configs/rec/rec_svtrnet.yml index a3d292b6ba..5896b10515 100644 --- a/configs/rec/rec_svtrnet.yml +++ b/configs/rec/rec_svtrnet.yml @@ -83,7 +83,7 @@ Train: img_mode: BGR channel_first: False - CTCLabelEncode: # Class handling label - - RecResizeImg: + - SVTRRecResizeImg: character_dict_path: image_shape: [3, 64, 256] padding: False @@ -104,7 +104,7 @@ Eval: img_mode: BGR channel_first: False - CTCLabelEncode: # Class handling label - - RecResizeImg: + - SVTRRecResizeImg: character_dict_path: image_shape: [3, 64, 256] padding: False diff --git a/configs/rec/rec_svtrnet_ch.yml b/configs/rec/rec_svtrnet_ch.yml new file mode 100644 index 0000000000..b327edf499 --- /dev/null +++ b/configs/rec/rec_svtrnet_ch.yml @@ -0,0 +1,159 @@ +Global: + use_gpu: true + epoch_num: 100 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/rec/svtr_ch_all/ + save_epoch_step: 10 + eval_batch_step: + - 0 + - 2000 + cal_metric_during_train: true + pretrained_model: null + checkpoints: null + save_inference_dir: null + use_visualdl: false + infer_img: doc/imgs_words/ch/word_1.jpg + character_dict_path: ppocr/utils/ppocr_keys_v1.txt + max_text_length: 25 + infer_mode: false + use_space_char: true + save_res_path: ./output/rec/predicts_svtr_tiny_ch_all.txt +Optimizer: + name: AdamW + beta1: 0.9 + beta2: 0.99 + epsilon: 8.0e-08 + weight_decay: 0.05 + no_weight_decay_name: norm pos_embed + one_dim_param_no_weight_decay: true + lr: + name: Cosine + learning_rate: 0.0005 + warmup_epoch: 2 +Architecture: + model_type: rec + algorithm: SVTR + Transform: null + Backbone: + name: SVTRNet + img_size: + - 32 + - 320 + out_char_num: 40 + out_channels: 96 + patch_merging: Conv + embed_dim: + - 64 + - 128 + - 256 + depth: + - 3 + - 6 + - 3 + num_heads: + - 2 + - 4 + - 8 + mixer: + - Local + - Local + - Local + - Local + - Local + - Local + - Global + - Global + - Global + - Global + - Global + - Global + local_mixer: + - - 7 + - 11 + - - 7 + - 11 + - - 7 + - 11 + last_stage: true + prenorm: false + Neck: + name: SequenceEncoder + encoder_type: reshape + Head: + name: CTCHead +Loss: + name: CTCLoss +PostProcess: + name: CTCLabelDecode +Metric: + name: RecMetric + main_indicator: acc +Train: + dataset: + name: SimpleDataSet + label_file_list: + - /paddle/data/ocr_all/train_all_list.txt + data_dir: /paddle/data/ocr_all + ext_op_transform_idx: 1 + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - RecConAug: + prob: 0.5 + ext_data_num: 2 + image_shape: + - 32 + - 320 + - 3 + - RecAug: null + - CTCLabelEncode: null + - SVTRRecResizeImg: + character_dict_path: null + infer_mode: False + image_shape: + - 3 + - 32 + - 320 + padding: true + - KeepKeys: + keep_keys: + - image + - label + - length + loader: + shuffle: true + batch_size_per_card: 256 + drop_last: true + num_workers: 8 +Eval: + dataset: + name: SimpleDataSet + data_dir: /paddle/data/ocr_all + label_file_list: + - /paddle/data/ocr_all/val_list.txt + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - CTCLabelEncode: null + - SVTRRecResizeImg: + character_dict_path: null + infer_mode: False + image_shape: + - 3 + - 32 + - 320 + padding: true + - KeepKeys: + keep_keys: + - image + - label + - length + loader: + shuffle: false + drop_last: false + batch_size_per_card: 256 + num_workers: 2 +profiler_options: null diff --git a/doc/doc_ch/algorithm_rec_svtr.md b/doc/doc_ch/algorithm_rec_svtr.md index 41a22ca65c..c0e26433e9 100644 --- a/doc/doc_ch/algorithm_rec_svtr.md +++ b/doc/doc_ch/algorithm_rec_svtr.md @@ -111,7 +111,6 @@ python3 tools/export_model.py -c ./rec_svtr_tiny_none_ctc_en_train/rec_svtr_tiny **注意：** - 如果您是在自己的数据集上训练的模型，并且调整了字典文件，请注意修改配置文件中的`character_dict_path`是否为所正确的字典文件。 -- 如果您修改了训练时的输入大小，请修改`tools/export_model.py`文件中的对应SVTR的`infer_shape`。 转换成功后，在目录下有三个文件： ``` diff --git a/doc/doc_en/algorithm_rec_svtr_en.md b/doc/doc_en/algorithm_rec_svtr_en.md index d402a6b491..37cd35f35a 100644 --- a/doc/doc_en/algorithm_rec_svtr_en.md +++ b/doc/doc_en/algorithm_rec_svtr_en.md @@ -88,7 +88,6 @@ python3 tools/export_model.py -c configs/rec/rec_svtrnet.yml -o Global.pretraine **Note:** - If you are training the model on your own dataset and have modified the dictionary file, please pay attention to modify the `character_dict_path` in the configuration file to the modified dictionary file. -- If you modified the input size during training, please modify the `infer_shape` corresponding to SVTR in the `tools/export_model.py` file. After the conversion is successful, there are three files in the directory: ``` diff --git a/ppocr/data/imaug/__init__.py b/ppocr/data/imaug/__init__.py index 58e8a5c780..437e0152fa 100644 --- a/ppocr/data/imaug/__init__.py +++ b/ppocr/data/imaug/__init__.py @@ -24,7 +24,8 @@ from .make_pse_gt import MakePseGt from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ - SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, ABINetRecResizeImg + SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, \ + ABINetRecResizeImg, SVTRRecResizeImg from .ssl_img_aug import SSLRotateResize from .randaugment import RandAugment from .copy_paste import CopyPaste diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index a8b3b813b3..874d9aa084 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -148,57 +148,6 @@ class ClsResizeImg(object): return data -class GrayRecResizeImg(object): - def __init__(self, - image_shape, - resize_type, - inter_type='Image.ANTIALIAS', - scale=True, - padding=False, - **kwargs): - self.image_shape = image_shape - self.resize_type = resize_type - self.padding = padding - self.inter_type = eval(inter_type) - self.scale = scale - - def __call__(self, data): - img = data['image'] - img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) - image_shape = self.image_shape - if self.padding: - imgC, imgH, imgW = image_shape - # todo: change to 0 and modified image shape - h = img.shape[0] - w = img.shape[1] - ratio = w / float(h) - if math.ceil(imgH * ratio) > imgW: - resized_w = imgW - else: - resized_w = int(math.ceil(imgH * ratio)) - resized_image = cv2.resize(img, (resized_w, imgH)) - norm_img = np.expand_dims(resized_image, -1) - norm_img = norm_img.transpose((2, 0, 1)) - resized_image = norm_img.astype(np.float32) / 128. - 1. - padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32) - padding_im[:, :, 0:resized_w] = resized_image - data['image'] = padding_im - return data - if self.resize_type == 'PIL': - image_pil = Image.fromarray(np.uint8(img)) - img = image_pil.resize(self.image_shape, self.inter_type) - img = np.array(img) - if self.resize_type == 'OpenCV': - img = cv2.resize(img, self.image_shape) - norm_img = np.expand_dims(img, -1) - norm_img = norm_img.transpose((2, 0, 1)) - if self.scale: - data['image'] = norm_img.astype(np.float32) / 128. - 1. - else: - data['image'] = norm_img.astype(np.float32) / 255. - return data - - class RecResizeImg(object): def __init__(self, image_shape, @@ -279,6 +228,57 @@ class PRENResizeImg(object): return data +class GrayRecResizeImg(object): + def __init__(self, + image_shape, + resize_type, + inter_type='Image.ANTIALIAS', + scale=True, + padding=False, + **kwargs): + self.image_shape = image_shape + self.resize_type = resize_type + self.padding = padding + self.inter_type = eval(inter_type) + self.scale = scale + + def __call__(self, data): + img = data['image'] + img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + image_shape = self.image_shape + if self.padding: + imgC, imgH, imgW = image_shape + # todo: change to 0 and modified image shape + h = img.shape[0] + w = img.shape[1] + ratio = w / float(h) + if math.ceil(imgH * ratio) > imgW: + resized_w = imgW + else: + resized_w = int(math.ceil(imgH * ratio)) + resized_image = cv2.resize(img, (resized_w, imgH)) + norm_img = np.expand_dims(resized_image, -1) + norm_img = norm_img.transpose((2, 0, 1)) + resized_image = norm_img.astype(np.float32) / 128. - 1. + padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32) + padding_im[:, :, 0:resized_w] = resized_image + data['image'] = padding_im + return data + if self.resize_type == 'PIL': + image_pil = Image.fromarray(np.uint8(img)) + img = image_pil.resize(self.image_shape, self.inter_type) + img = np.array(img) + if self.resize_type == 'OpenCV': + img = cv2.resize(img, self.image_shape) + norm_img = np.expand_dims(img, -1) + norm_img = norm_img.transpose((2, 0, 1)) + if self.scale: + data['image'] = norm_img.astype(np.float32) / 128. - 1. + else: + data['image'] = norm_img.astype(np.float32) / 255. + return data + + class ABINetRecResizeImg(object): def __init__(self, image_shape, @@ -297,6 +297,28 @@ class ABINetRecResizeImg(object): return data +class SVTRRecResizeImg(object): + def __init__(self, + image_shape, + infer_mode=False, + character_dict_path='./ppocr/utils/ppocr_keys_v1.txt', + padding=True, + **kwargs): + self.image_shape = image_shape + self.infer_mode = infer_mode + self.character_dict_path = character_dict_path + self.padding = padding + + def __call__(self, data): + img = data['image'] + + norm_img, valid_ratio = resize_norm_img(img, self.image_shape, + self.padding) + data['image'] = norm_img + data['valid_ratio'] = valid_ratio + return data + + def resize_norm_img_sar(img, image_shape, width_downsample_ratio=0.25): imgC, imgH, imgW_min, imgW_max = image_shape h = img.shape[0] diff --git a/tools/export_model.py b/tools/export_model.py index e2673239b1..b10d41d5b2 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -31,7 +31,12 @@ from ppocr.utils.logging import get_logger from tools.program import load_config, merge_config, ArgsParser -def export_single_model(model, arch_config, save_path, logger, quanter=None): +def export_single_model(model, + arch_config, + save_path, + logger, + input_shape=None, + quanter=None): if arch_config["algorithm"] == "SRN": max_text_length = arch_config["Head"]["max_text_length"] other_shape = [ @@ -64,7 +69,7 @@ def export_single_model(model, arch_config, save_path, logger, quanter=None): else: other_shape = [ paddle.static.InputSpec( - shape=[None, 3, 64, 256], dtype="float32"), + shape=[None] + input_shape, dtype="float32"), ] model = to_static(model, input_spec=other_shape) elif arch_config["algorithm"] == "PREN": @@ -174,6 +179,13 @@ def main(): arch_config = config["Architecture"] + if arch_config["algorithm"] == "SVTR" and arch_config["Head"][ + "name"] != 'MultiHead': + input_shape = config["Eval"]["dataset"]["transforms"][-2][ + 'SVTRRecResizeImg']['image_shape'] + else: + input_shape = None + if arch_config["algorithm"] in ["Distillation", ]: # distillation model archs = list(arch_config["Models"].values()) for idx, name in enumerate(model.model_name_list): @@ -182,7 +194,8 @@ def main(): sub_model_save_path, logger) else: save_path = os.path.join(save_path, "inference") - export_single_model(model, arch_config, save_path, logger) + export_single_model( + model, arch_config, save_path, logger, input_shape=input_shape) if __name__ == "__main__": From 31ab300ca258243a311ae08e1496faa789567a0c Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Thu, 23 Jun 2022 13:09:48 +0000 Subject: [PATCH 09/43] [application] svtr ch model --- applications/高精度中文识别模型.md | 101 +++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 applications/高精度中文识别模型.md diff --git a/applications/高精度中文识别模型.md b/applications/高精度中文识别模型.md new file mode 100644 index 0000000000..27b63abf3c --- /dev/null +++ b/applications/高精度中文识别模型.md @@ -0,0 +1,101 @@ +# 高精度中文场景文本识别模型SVTR + +## 1. 简介 + +PP-OCRv3是百度开源的超轻量级场景文本检测识别模型库，其中超轻量的场景中文识别模型SVTR_LCNet使用了SVTR算法结构。为了保证速度，SVTR_LCNet将SVTR模型的Local Blocks替换为LCNet，使用两层Global Blocks。在中文场景中，经过多次优化，SVTR_LCNet的最终精度为79.4%。具体的： +- GTC：Attention指导CTC训练策略； +- TextConAug：挖掘文字上下文信息的数据增广策略； +- TextRotNet：自监督的预训练模型； +- UDML：联合互学习策略； +- UIM：无标注数据挖掘方案。 + +其中 *UIM：无标注数据挖掘方案* 使用了高精度的SVTR中文模型，该模型在PP-OCRv3识别的数据集上训练，精度为82.5%。本项目主要介绍SVTR中文模型如何使用。 + +aistudio项目链接: [高精度中文场景文本识别模型SVTR](https://aistudio.baidu.com/aistudio/projectdetail/4263032) + +## 2. SVTR中文模型使用 + +### 环境准备 + + +本任务基于Aistudio完成, 具体环境如下： + +- 操作系统: Linux +- PaddlePaddle: 2.3 +- PaddleOCR: dygraph + +下载 PaddleOCR代码 + +```bash +git clone -b dygraph https://github.com/PaddlePaddle/PaddleOCR +``` + +安装依赖库 + +```bash +pip install -r PaddleOCR/requirements.txt -i https://mirror.baidu.com/pypi/simple +``` + +### 快速使用 + +获取SVTR中文模型文件，请扫码填写问卷，加入PaddleOCR官方交流群获取全部OCR垂类模型下载链接、《动手学OCR》电子书等全套OCR学习资料🎁 +
+ +
+ +```bash +# 解压模型文件 +tar xf svtr_ch_high_accuracy.tar +``` + +预测中文文本，以下图为例： +![](https://ai-studio-static-online.cdn.bcebos.com/511f6d65f0e84ccc8c5668eafbe7feb73b26c40befec4bd58dde64b59bf340a1) + +预测命令： + +```bash +# CPU预测 +python tools/infer_rec.py -c configs/rec/rec_svtrnet_ch.yml -o Global.pretrained_model=./svtr_ch_high_accuracy/best_accuracy Global.infer_img=./doc/imgs_words/ch/word_1.jpg Global.use_gpu=False + +# GPU预测 +#python tools/infer_rec.py -c configs/rec/rec_svtrnet_ch.yml -o Global.pretrained_model=./svtr_ch_high_accuracy/best_accuracy Global.infer_img=./doc/imgs_words/ch/word_1.jpg Global.use_gpu=True +``` + +可以看到最后打印结果为 +- result: 韩国小馆 0.9853458404541016 + +0.9853458404541016为预测置信度。 + +### 推理模型导出与预测 + +inference 模型（paddle.jit.save保存的模型） 一般是模型训练，把模型结构和模型参数保存在文件中的固化模型，多用于预测部署场景。 训练过程中保存的模型是checkpoints模型，保存的只有模型的参数，多用于恢复训练等。 与checkpoints模型相比，inference 模型会额外保存模型的结构信息，在预测部署、加速推理上性能优越，灵活方便，适合于实际系统集成。 + +运行识别模型转inference模型命令，如下： + +```bash +python tools/export_model.py -c configs/rec/rec_svtrnet_ch.yml -o Global.pretrained_model=./svtr_ch_high_accuracy/best_accuracy Global.save_inference_dir=./inference/svtr_ch +``` + +转换成功后，在目录下有三个文件： +```shell +inference/svtr_ch/ + ├── inference.pdiparams # 识别inference模型的参数文件 + ├── inference.pdiparams.info # 识别inference模型的参数信息，可忽略 + └── inference.pdmodel # 识别inference模型的program文件 +``` + +inference模型预测，命令如下： + +```bash +# CPU预测 +python3 tools/infer/predict_rec.py --image_dir="./doc/imgs_words/ch/word_1.jpg" --rec_algorithm='SVTR' --rec_model_dir=./inference/svtr_ch/ --rec_image_shape='3, 32, 320' --rec_char_dict_path=ppocr/utils/ppocr_keys_v1.txt --use_gpu=False + +# GPU预测 +#python3 tools/infer/predict_rec.py --image_dir="./doc/imgs_words/ch/word_1.jpg" --rec_algorithm='SVTR' --rec_model_dir=./inference/svtr_ch/ --rec_image_shape='3, 32, 320' --rec_char_dict_path=ppocr/utils/ppocr_keys_v1.txt --use_gpu=True +``` + +**注意** + +- 使用PP-OCRv3识别SVTR_LCNet进行推理时，不需要使用--rec_algorithm指定算法名称，使用默认的推理方式即为PP-OCRv3识别的推理过程。但是使用SVTR算法时，需要指定--rec_algorithm='SVTR' +- 如果使用自定义字典训练的模型，需要将--rec_char_dict_path=ppocr/utils/ppocr_keys_v1.txt修改为自定义的字典 +- --rec_image_shape='3, 32, 320' 该参数不能去掉 From 6ef2c8762ed7ec3a90a6680292a6095236428103 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Thu, 23 Jun 2022 13:25:31 +0000 Subject: [PATCH 10/43] [application] svtr ch model --- applications/高精度中文识别模型.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/applications/高精度中文识别模型.md b/applications/高精度中文识别模型.md index 27b63abf3c..31d68abbff 100644 --- a/applications/高精度中文识别模型.md +++ b/applications/高精度中文识别模型.md @@ -9,7 +9,12 @@ PP-OCRv3是百度开源的超轻量级场景文本检测识别模型库，其中 - UDML：联合互学习策略； - UIM：无标注数据挖掘方案。 -其中 *UIM：无标注数据挖掘方案* 使用了高精度的SVTR中文模型，该模型在PP-OCRv3识别的数据集上训练，精度为82.5%。本项目主要介绍SVTR中文模型如何使用。 +其中 *UIM：无标注数据挖掘方案* 使用了高精度的SVTR中文模型进行无标注文件的刷库，该模型在PP-OCRv3识别的数据集上训练，精度对比如下表。 + +|中文识别算法|模型|精度| +| --- | --- | --- | +|PP-OCRv3|SVTR_LCNet|79.4%| +|SVTR|SVTR-Tiny|82.5%| aistudio项目链接: [高精度中文场景文本识别模型SVTR](https://aistudio.baidu.com/aistudio/projectdetail/4263032) @@ -96,6 +101,6 @@ python3 tools/infer/predict_rec.py --image_dir="./doc/imgs_words/ch/word_1.jpg" **注意** -- 使用PP-OCRv3识别SVTR_LCNet进行推理时，不需要使用--rec_algorithm指定算法名称，使用默认的推理方式即为PP-OCRv3识别的推理过程。但是使用SVTR算法时，需要指定--rec_algorithm='SVTR' +- 使用SVTR算法时，需要指定--rec_algorithm='SVTR' - 如果使用自定义字典训练的模型，需要将--rec_char_dict_path=ppocr/utils/ppocr_keys_v1.txt修改为自定义的字典 - --rec_image_shape='3, 32, 320' 该参数不能去掉 From 1fc0655d857f25704a54e6550a37c7c68f7d5fdd Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Thu, 23 Jun 2022 13:31:57 +0000 Subject: [PATCH 11/43] [application] svtr ch model --- applications/高精度中文识别模型.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/高精度中文识别模型.md b/applications/高精度中文识别模型.md index 31d68abbff..e3d0693ece 100644 --- a/applications/高精度中文识别模型.md +++ b/applications/高精度中文识别模型.md @@ -54,7 +54,7 @@ tar xf svtr_ch_high_accuracy.tar ``` 预测中文文本，以下图为例： -![](https://ai-studio-static-online.cdn.bcebos.com/511f6d65f0e84ccc8c5668eafbe7feb73b26c40befec4bd58dde64b59bf340a1) +![](../doc/imgs_words/ch/word_1.jpg) 预测命令： From 773b687f5a5a6d3115e27ac15070fd13ea6bb5ce Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Fri, 24 Jun 2022 07:25:50 +0000 Subject: [PATCH 12/43] add abinet_rec_aug and trained model --- configs/rec/rec_r45_abinet.yml | 3 +- doc/doc_ch/algorithm_rec_abinet.md | 4 +- doc/doc_en/algorithm_rec_abinet_en.md | 4 +- ppocr/data/imaug/__init__.py | 2 +- ppocr/data/imaug/abinet_aug.py | 407 ++++++++++++++++++ ppocr/data/imaug/rec_img_aug.py | 27 ++ .../configs/rec_r45_abinet/rec_r45_abinet.yml | 1 + 7 files changed, 442 insertions(+), 6 deletions(-) create mode 100644 ppocr/data/imaug/abinet_aug.py diff --git a/configs/rec/rec_r45_abinet.yml b/configs/rec/rec_r45_abinet.yml index a756fead9f..3cf4cddb04 100644 --- a/configs/rec/rec_r45_abinet.yml +++ b/configs/rec/rec_r45_abinet.yml @@ -8,7 +8,7 @@ Global: # evaluation is run every 2000 iterations eval_batch_step: [0, 2000] cal_metric_during_train: True - pretrained_model: + pretrained_model: ./rec_r45_abinet_train/abinet_vl_pretrained checkpoints: save_inference_dir: use_visualdl: False @@ -67,6 +67,7 @@ Train: - DecodeImage: # load image img_mode: RGB channel_first: False + - ABINetRecAug: - ABINetLabelEncode: # Class handling label ignore_index: *ignore_index - ABINetRecResizeImg: diff --git a/doc/doc_ch/algorithm_rec_abinet.md b/doc/doc_ch/algorithm_rec_abinet.md index d20c703014..47507c36c7 100644 --- a/doc/doc_ch/algorithm_rec_abinet.md +++ b/doc/doc_ch/algorithm_rec_abinet.md @@ -27,7 +27,7 @@ |模型|骨干网络|配置文件|Acc|下载链接| | --- | --- | --- | --- | --- | -|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[训练模型]()/[预训练模型]| +|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[预训练、训练模型](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar)| ## 2. 环境配置 @@ -80,7 +80,7 @@ python3 tools/infer_rec.py -c configs/rec/rec_r45_abinet.yml -o Global.infer_img ### 4.1 Python推理 -首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址]() )，可以使用如下命令进行转换： +首先将训练得到best模型，转换成inference model。这里以训练完成的模型为例（[模型下载地址](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) )，可以使用如下命令进行转换： ```shell # 注意将pretrained_model的路径设置为本地路径。 diff --git a/doc/doc_en/algorithm_rec_abinet_en.md b/doc/doc_en/algorithm_rec_abinet_en.md index 3b0f6c09a9..767ca65f64 100644 --- a/doc/doc_en/algorithm_rec_abinet_en.md +++ b/doc/doc_en/algorithm_rec_abinet_en.md @@ -25,7 +25,7 @@ Using MJSynth and SynthText two text recognition datasets for training, and eval |Model|Backbone|config|Acc|Download link| | --- | --- | --- | --- | --- | -|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[trained model]()/[pretrained model]()| +|ABINet|ResNet45|[rec_r45_abinet.yml](../../configs/rec/rec_r45_abinet.yml)|90.75%|[pretrained & trained model](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar)| ## 2. Environment @@ -68,7 +68,7 @@ python3 tools/infer_rec.py -c configs/rec/rec_r45_abinet.yml -o Global.infer_img ### 4.1 Python Inference -First, the model saved during the ABINet text recognition training process is converted into an inference model. ( [Model download link]()) ), you can use the following command to convert: +First, the model saved during the ABINet text recognition training process is converted into an inference model. ( [Model download link](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar)) ), you can use the following command to convert: ``` python3 tools/export_model.py -c configs/rec/rec_r45_abinet.yml -o Global.pretrained_model=./rec_r45_abinet_train/best_accuracy Global.save_inference_dir=./inference/rec_r45_abinet diff --git a/ppocr/data/imaug/__init__.py b/ppocr/data/imaug/__init__.py index 437e0152fa..63dfda91f8 100644 --- a/ppocr/data/imaug/__init__.py +++ b/ppocr/data/imaug/__init__.py @@ -25,7 +25,7 @@ from .make_pse_gt import MakePseGt from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, \ - ABINetRecResizeImg, SVTRRecResizeImg + ABINetRecResizeImg, SVTRRecResizeImg, ABINetRecAug from .ssl_img_aug import SSLRotateResize from .randaugment import RandAugment from .copy_paste import CopyPaste diff --git a/ppocr/data/imaug/abinet_aug.py b/ppocr/data/imaug/abinet_aug.py new file mode 100644 index 0000000000..eefdc75d5a --- /dev/null +++ b/ppocr/data/imaug/abinet_aug.py @@ -0,0 +1,407 @@ +# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +This code is refer from: +https://github.com/FangShancheng/ABINet/blob/main/transforms.py +""" +import math +import numbers +import random + +import cv2 +import numpy as np +from paddle.vision.transforms import Compose, ColorJitter + + +def sample_asym(magnitude, size=None): + return np.random.beta(1, 4, size) * magnitude + + +def sample_sym(magnitude, size=None): + return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude + + +def sample_uniform(low, high, size=None): + return np.random.uniform(low, high, size=size) + + +def get_interpolation(type='random'): + if type == 'random': + choice = [ + cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA + ] + interpolation = choice[random.randint(0, len(choice) - 1)] + elif type == 'nearest': + interpolation = cv2.INTER_NEAREST + elif type == 'linear': + interpolation = cv2.INTER_LINEAR + elif type == 'cubic': + interpolation = cv2.INTER_CUBIC + elif type == 'area': + interpolation = cv2.INTER_AREA + else: + raise TypeError( + 'Interpolation types only nearest, linear, cubic, area are supported!' + ) + return interpolation + + +class CVRandomRotation(object): + def __init__(self, degrees=15): + assert isinstance(degrees, + numbers.Number), "degree should be a single number." + assert degrees >= 0, "degree must be positive." + self.degrees = degrees + + @staticmethod + def get_params(degrees): + return sample_sym(degrees) + + def __call__(self, img): + angle = self.get_params(self.degrees) + src_h, src_w = img.shape[:2] + M = cv2.getRotationMatrix2D( + center=(src_w / 2, src_h / 2), angle=angle, scale=1.0) + abs_cos, abs_sin = abs(M[0, 0]), abs(M[0, 1]) + dst_w = int(src_h * abs_sin + src_w * abs_cos) + dst_h = int(src_h * abs_cos + src_w * abs_sin) + M[0, 2] += (dst_w - src_w) / 2 + M[1, 2] += (dst_h - src_h) / 2 + + flags = get_interpolation() + return cv2.warpAffine( + img, + M, (dst_w, dst_h), + flags=flags, + borderMode=cv2.BORDER_REPLICATE) + + +class CVRandomAffine(object): + def __init__(self, degrees, translate=None, scale=None, shear=None): + assert isinstance(degrees, + numbers.Number), "degree should be a single number." + assert degrees >= 0, "degree must be positive." + self.degrees = degrees + + if translate is not None: + assert isinstance(translate, (tuple, list)) and len(translate) == 2, \ + "translate should be a list or tuple and it must be of length 2." + for t in translate: + if not (0.0 <= t <= 1.0): + raise ValueError( + "translation values should be between 0 and 1") + self.translate = translate + + if scale is not None: + assert isinstance(scale, (tuple, list)) and len(scale) == 2, \ + "scale should be a list or tuple and it must be of length 2." + for s in scale: + if s <= 0: + raise ValueError("scale values should be positive") + self.scale = scale + + if shear is not None: + if isinstance(shear, numbers.Number): + if shear < 0: + raise ValueError( + "If shear is a single number, it must be positive.") + self.shear = [shear] + else: + assert isinstance(shear, (tuple, list)) and (len(shear) == 2), \ + "shear should be a list or tuple and it must be of length 2." + self.shear = shear + else: + self.shear = shear + + def _get_inverse_affine_matrix(self, center, angle, translate, scale, + shear): + # https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717 + from numpy import sin, cos, tan + + if isinstance(shear, numbers.Number): + shear = [shear, 0] + + if not isinstance(shear, (tuple, list)) and len(shear) == 2: + raise ValueError( + "Shear should be a single value or a tuple/list containing " + + "two values. Got {}".format(shear)) + + rot = math.radians(angle) + sx, sy = [math.radians(s) for s in shear] + + cx, cy = center + tx, ty = translate + + # RSS without scaling + a = cos(rot - sy) / cos(sy) + b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot) + c = sin(rot - sy) / cos(sy) + d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot) + + # Inverted rotation matrix with scale and shear + # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1 + M = [d, -b, 0, -c, a, 0] + M = [x / scale for x in M] + + # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1 + M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty) + M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty) + + # Apply center translation: C * RSS^-1 * C^-1 * T^-1 + M[2] += cx + M[5] += cy + return M + + @staticmethod + def get_params(degrees, translate, scale_ranges, shears, height): + angle = sample_sym(degrees) + if translate is not None: + max_dx = translate[0] * height + max_dy = translate[1] * height + translations = (np.round(sample_sym(max_dx)), + np.round(sample_sym(max_dy))) + else: + translations = (0, 0) + + if scale_ranges is not None: + scale = sample_uniform(scale_ranges[0], scale_ranges[1]) + else: + scale = 1.0 + + if shears is not None: + if len(shears) == 1: + shear = [sample_sym(shears[0]), 0.] + elif len(shears) == 2: + shear = [sample_sym(shears[0]), sample_sym(shears[1])] + else: + shear = 0.0 + + return angle, translations, scale, shear + + def __call__(self, img): + src_h, src_w = img.shape[:2] + angle, translate, scale, shear = self.get_params( + self.degrees, self.translate, self.scale, self.shear, src_h) + + M = self._get_inverse_affine_matrix((src_w / 2, src_h / 2), angle, + (0, 0), scale, shear) + M = np.array(M).reshape(2, 3) + + startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), + (0, src_h - 1)] + project = lambda x, y, a, b, c: int(a * x + b * y + c) + endpoints = [(project(x, y, *M[0]), project(x, y, *M[1])) + for x, y in startpoints] + + rect = cv2.minAreaRect(np.array(endpoints)) + bbox = cv2.boxPoints(rect).astype(dtype=np.int) + max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max() + min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min() + + dst_w = int(max_x - min_x) + dst_h = int(max_y - min_y) + M[0, 2] += (dst_w - src_w) / 2 + M[1, 2] += (dst_h - src_h) / 2 + + # add translate + dst_w += int(abs(translate[0])) + dst_h += int(abs(translate[1])) + if translate[0] < 0: M[0, 2] += abs(translate[0]) + if translate[1] < 0: M[1, 2] += abs(translate[1]) + + flags = get_interpolation() + return cv2.warpAffine( + img, + M, (dst_w, dst_h), + flags=flags, + borderMode=cv2.BORDER_REPLICATE) + + +class CVRandomPerspective(object): + def __init__(self, distortion=0.5): + self.distortion = distortion + + def get_params(self, width, height, distortion): + offset_h = sample_asym( + distortion * height / 2, size=4).astype(dtype=np.int) + offset_w = sample_asym( + distortion * width / 2, size=4).astype(dtype=np.int) + topleft = (offset_w[0], offset_h[0]) + topright = (width - 1 - offset_w[1], offset_h[1]) + botright = (width - 1 - offset_w[2], height - 1 - offset_h[2]) + botleft = (offset_w[3], height - 1 - offset_h[3]) + + startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), + (0, height - 1)] + endpoints = [topleft, topright, botright, botleft] + return np.array( + startpoints, dtype=np.float32), np.array( + endpoints, dtype=np.float32) + + def __call__(self, img): + height, width = img.shape[:2] + startpoints, endpoints = self.get_params(width, height, self.distortion) + M = cv2.getPerspectiveTransform(startpoints, endpoints) + + # TODO: more robust way to crop image + rect = cv2.minAreaRect(endpoints) + bbox = cv2.boxPoints(rect).astype(dtype=np.int) + max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max() + min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min() + min_x, min_y = max(min_x, 0), max(min_y, 0) + + flags = get_interpolation() + img = cv2.warpPerspective( + img, + M, (max_x, max_y), + flags=flags, + borderMode=cv2.BORDER_REPLICATE) + img = img[min_y:, min_x:] + return img + + +class CVRescale(object): + def __init__(self, factor=4, base_size=(128, 512)): + """ Define image scales using gaussian pyramid and rescale image to target scale. + + Args: + factor: the decayed factor from base size, factor=4 keeps target scale by default. + base_size: base size the build the bottom layer of pyramid + """ + if isinstance(factor, numbers.Number): + self.factor = round(sample_uniform(0, factor)) + elif isinstance(factor, (tuple, list)) and len(factor) == 2: + self.factor = round(sample_uniform(factor[0], factor[1])) + else: + raise Exception('factor must be number or list with length 2') + # assert factor is valid + self.base_h, self.base_w = base_size[:2] + + def __call__(self, img): + if self.factor == 0: return img + src_h, src_w = img.shape[:2] + cur_w, cur_h = self.base_w, self.base_h + scale_img = cv2.resize( + img, (cur_w, cur_h), interpolation=get_interpolation()) + for _ in range(self.factor): + scale_img = cv2.pyrDown(scale_img) + scale_img = cv2.resize( + scale_img, (src_w, src_h), interpolation=get_interpolation()) + return scale_img + + +class CVGaussianNoise(object): + def __init__(self, mean=0, var=20): + self.mean = mean + if isinstance(var, numbers.Number): + self.var = max(int(sample_asym(var)), 1) + elif isinstance(var, (tuple, list)) and len(var) == 2: + self.var = int(sample_uniform(var[0], var[1])) + else: + raise Exception('degree must be number or list with length 2') + + def __call__(self, img): + noise = np.random.normal(self.mean, self.var**0.5, img.shape) + img = np.clip(img + noise, 0, 255).astype(np.uint8) + return img + + +class CVMotionBlur(object): + def __init__(self, degrees=12, angle=90): + if isinstance(degrees, numbers.Number): + self.degree = max(int(sample_asym(degrees)), 1) + elif isinstance(degrees, (tuple, list)) and len(degrees) == 2: + self.degree = int(sample_uniform(degrees[0], degrees[1])) + else: + raise Exception('degree must be number or list with length 2') + self.angle = sample_uniform(-angle, angle) + + def __call__(self, img): + M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), + self.angle, 1) + motion_blur_kernel = np.zeros((self.degree, self.degree)) + motion_blur_kernel[self.degree // 2, :] = 1 + motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, + (self.degree, self.degree)) + motion_blur_kernel = motion_blur_kernel / self.degree + img = cv2.filter2D(img, -1, motion_blur_kernel) + img = np.clip(img, 0, 255).astype(np.uint8) + return img + + +class CVGeometry(object): + def __init__(self, + degrees=15, + translate=(0.3, 0.3), + scale=(0.5, 2.), + shear=(45, 15), + distortion=0.5, + p=0.5): + self.p = p + type_p = random.random() + if type_p < 0.33: + self.transforms = CVRandomRotation(degrees=degrees) + elif type_p < 0.66: + self.transforms = CVRandomAffine( + degrees=degrees, translate=translate, scale=scale, shear=shear) + else: + self.transforms = CVRandomPerspective(distortion=distortion) + + def __call__(self, img): + if random.random() < self.p: + return self.transforms(img) + else: + return img + + +class CVDeterioration(object): + def __init__(self, var, degrees, factor, p=0.5): + self.p = p + transforms = [] + if var is not None: + transforms.append(CVGaussianNoise(var=var)) + if degrees is not None: + transforms.append(CVMotionBlur(degrees=degrees)) + if factor is not None: + transforms.append(CVRescale(factor=factor)) + + random.shuffle(transforms) + transforms = Compose(transforms) + self.transforms = transforms + + def __call__(self, img): + if random.random() < self.p: + + return self.transforms(img) + else: + return img + + +class CVColorJitter(object): + def __init__(self, + brightness=0.5, + contrast=0.5, + saturation=0.5, + hue=0.1, + p=0.5): + self.p = p + self.transforms = ColorJitter( + brightness=brightness, + contrast=contrast, + saturation=saturation, + hue=hue) + + def __call__(self, img): + if random.random() < self.p: return self.transforms(img) + else: return img diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 874d9aa084..92f7382724 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -19,6 +19,8 @@ import random import copy from PIL import Image from .text_image_aug import tia_perspective, tia_stretch, tia_distort +from .abinet_aug import CVGeometry, CVDeterioration, CVColorJitter +from paddle.vision.transforms import Compose class RecAug(object): @@ -94,6 +96,31 @@ class BaseDataAugmentation(object): return data +class ABINetRecAug(object): + def __init__(self, **kwargs): + self.transforms = Compose([ + CVGeometry( + degrees=45, + translate=(0.0, 0.0), + scale=(0.5, 2.), + shear=(45, 15), + distortion=0.5, + p=0.5), CVDeterioration( + var=20, degrees=6, factor=4, p=0.25), CVColorJitter( + brightness=0.5, + contrast=0.5, + saturation=0.5, + hue=0.1, + p=0.25) + ]) + + def __call__(self, data): + img = data['image'] + img = self.transforms(img) + data['image'] = img + return data + + class RecConAug(object): def __init__(self, prob=0.5, diff --git a/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml b/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml index 306c49d683..5b5890e772 100644 --- a/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml +++ b/test_tipc/configs/rec_r45_abinet/rec_r45_abinet.yml @@ -68,6 +68,7 @@ Train: - DecodeImage: # load image img_mode: RGB channel_first: False + - ABINetRecAug: - ABINetLabelEncode: # Class handling label ignore_index: *ignore_index - ABINetRecResizeImg: From 1c5e7f2578b8ac1d8b4999782ac0349a0f49781d Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Fri, 24 Jun 2022 09:56:27 +0000 Subject: [PATCH 13/43] aug p infe --- applications/高精度中文识别模型.md | 11 ++++++----- configs/rec/rec_svtrnet_ch.yml | 8 ++++---- ppocr/data/imaug/rec_img_aug.py | 21 +++++++++++++-------- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/applications/高精度中文识别模型.md b/applications/高精度中文识别模型.md index e3d0693ece..3c31af42ee 100644 --- a/applications/高精度中文识别模型.md +++ b/applications/高精度中文识别模型.md @@ -2,7 +2,7 @@ ## 1. 简介 -PP-OCRv3是百度开源的超轻量级场景文本检测识别模型库，其中超轻量的场景中文识别模型SVTR_LCNet使用了SVTR算法结构。为了保证速度，SVTR_LCNet将SVTR模型的Local Blocks替换为LCNet，使用两层Global Blocks。在中文场景中，经过多次优化，SVTR_LCNet的最终精度为79.4%。具体的： +PP-OCRv3是百度开源的超轻量级场景文本检测识别模型库，其中超轻量的场景中文识别模型SVTR_LCNet使用了SVTR算法结构。为了保证速度，SVTR_LCNet将SVTR模型的Local Blocks替换为LCNet，使用两层Global Blocks。在中文场景中，PP-OCRv3识别主要使用如下优化策略： - GTC：Attention指导CTC训练策略； - TextConAug：挖掘文字上下文信息的数据增广策略； - TextRotNet：自监督的预训练模型； @@ -11,10 +11,11 @@ PP-OCRv3是百度开源的超轻量级场景文本检测识别模型库，其中 其中 *UIM：无标注数据挖掘方案* 使用了高精度的SVTR中文模型进行无标注文件的刷库，该模型在PP-OCRv3识别的数据集上训练，精度对比如下表。 -|中文识别算法|模型|精度| -| --- | --- | --- | -|PP-OCRv3|SVTR_LCNet|79.4%| -|SVTR|SVTR-Tiny|82.5%| +|中文识别算法|模型|UIM|精度| +| --- | --- | --- |--- | +|PP-OCRv3|SVTR_LCNet| w/o |78.4%| +|PP-OCRv3|SVTR_LCNet| w |79.4%| +|SVTR|SVTR-Tiny|-|82.5%| aistudio项目链接: [高精度中文场景文本识别模型SVTR](https://aistudio.baidu.com/aistudio/projectdetail/4263032) diff --git a/configs/rec/rec_svtrnet_ch.yml b/configs/rec/rec_svtrnet_ch.yml index b327edf499..c4ce4ccc16 100644 --- a/configs/rec/rec_svtrnet_ch.yml +++ b/configs/rec/rec_svtrnet_ch.yml @@ -92,9 +92,9 @@ Metric: Train: dataset: name: SimpleDataSet + data_dir: ./train_data label_file_list: - - /paddle/data/ocr_all/train_all_list.txt - data_dir: /paddle/data/ocr_all + - ./train_data/train_list.txt ext_op_transform_idx: 1 transforms: - DecodeImage: @@ -130,9 +130,9 @@ Train: Eval: dataset: name: SimpleDataSet - data_dir: /paddle/data/ocr_all + data_dir: ./train_data label_file_list: - - /paddle/data/ocr_all/val_list.txt + - ./train_data/val_list.txt transforms: - DecodeImage: img_mode: BGR diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 92f7382724..1de8e6ee3c 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -97,7 +97,11 @@ class BaseDataAugmentation(object): class ABINetRecAug(object): - def __init__(self, **kwargs): + def __init__(self, + geometry_p=0.5, + deterioration_p=0.25, + colorjitter_p=0.25, + **kwargs): self.transforms = Compose([ CVGeometry( degrees=45, @@ -105,13 +109,14 @@ class ABINetRecAug(object): scale=(0.5, 2.), shear=(45, 15), distortion=0.5, - p=0.5), CVDeterioration( - var=20, degrees=6, factor=4, p=0.25), CVColorJitter( - brightness=0.5, - contrast=0.5, - saturation=0.5, - hue=0.1, - p=0.25) + p=geometry_p), CVDeterioration( + var=20, degrees=6, factor=4, p=deterioration_p), + CVColorJitter( + brightness=0.5, + contrast=0.5, + saturation=0.5, + hue=0.1, + p=colorjitter_p) ]) def __call__(self, data): From 5f990d9cda3656ca8992f9e6473f25f13eee4ff3 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:11:31 +0000 Subject: [PATCH 14/43] add yml --- configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml b/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml index 3833cb0bad..df429314cd 100644 --- a/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml +++ b/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_cml.yml @@ -28,7 +28,7 @@ Architecture: algorithm: DB Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 18 Neck: name: DBFPN From 5ba74d56277debf930dc10dd6105e84f55247b56 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:16:10 +0000 Subject: [PATCH 15/43] add yml --- configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_distill.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_distill.yml b/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_distill.yml index b4644244f8..d24ee11f3d 100644 --- a/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_distill.yml +++ b/configs/det/ch_PP-OCRv2/ch_PP-OCRv2_det_distill.yml @@ -45,7 +45,7 @@ Architecture: algorithm: DB Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 18 Neck: name: DBFPN From bc29ec0ba8e72df1db81d0dfed57a337331660ba Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:17:49 +0000 Subject: [PATCH 16/43] add yml --- configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_cml.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_cml.yml b/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_cml.yml index 3e77577c17..ef58befd69 100644 --- a/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_cml.yml +++ b/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_cml.yml @@ -61,7 +61,7 @@ Architecture: model_type: det algorithm: DB Backbone: - name: ResNet + name: ResNet_vd in_channels: 3 layers: 50 Neck: From 961dca72cf3e4708d6ac57769a79fba1b56d5e2c Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:21:15 +0000 Subject: [PATCH 17/43] add yml --- configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_dml.yml | 4 ++-- configs/det/ch_ppocr_v2.0/ch_det_res18_db_v2.0.yml | 2 +- configs/det/det_r50_vd_db.yml | 2 +- configs/det/det_r50_vd_dcn_fce_ctw.yml | 2 +- configs/det/det_r50_vd_east.yml | 2 +- configs/det/det_r50_vd_pse.yml | 2 +- configs/det/det_res18_db_v2.0.yml | 2 +- 7 files changed, 8 insertions(+), 8 deletions(-) diff --git a/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_dml.yml b/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_dml.yml index 3f30ada13f..e1831f2639 100644 --- a/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_dml.yml +++ b/configs/det/ch_PP-OCRv3/ch_PP-OCRv3_det_dml.yml @@ -25,7 +25,7 @@ Architecture: model_type: det algorithm: DB Backbone: - name: ResNet + name: ResNet_vd in_channels: 3 layers: 50 Neck: @@ -40,7 +40,7 @@ Architecture: model_type: det algorithm: DB Backbone: - name: ResNet + name: ResNet_vd in_channels: 3 layers: 50 Neck: diff --git a/configs/det/ch_ppocr_v2.0/ch_det_res18_db_v2.0.yml b/configs/det/ch_ppocr_v2.0/ch_det_res18_db_v2.0.yml index 7b07ef9964..e983c221e2 100644 --- a/configs/det/ch_ppocr_v2.0/ch_det_res18_db_v2.0.yml +++ b/configs/det/ch_ppocr_v2.0/ch_det_res18_db_v2.0.yml @@ -20,7 +20,7 @@ Architecture: algorithm: DB Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 18 disable_se: True Neck: diff --git a/configs/det/det_r50_vd_db.yml b/configs/det/det_r50_vd_db.yml index ab67786ece..288dcc8c1a 100644 --- a/configs/det/det_r50_vd_db.yml +++ b/configs/det/det_r50_vd_db.yml @@ -20,7 +20,7 @@ Architecture: algorithm: DB Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 50 Neck: name: DBFPN diff --git a/configs/det/det_r50_vd_dcn_fce_ctw.yml b/configs/det/det_r50_vd_dcn_fce_ctw.yml index a9f7c4143d..3a4075b322 100755 --- a/configs/det/det_r50_vd_dcn_fce_ctw.yml +++ b/configs/det/det_r50_vd_dcn_fce_ctw.yml @@ -21,7 +21,7 @@ Architecture: algorithm: FCE Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 50 dcn_stage: [False, True, True, True] out_indices: [1,2,3] diff --git a/configs/det/det_r50_vd_east.yml b/configs/det/det_r50_vd_east.yml index e84a5fa7a7..af90ef0adb 100644 --- a/configs/det/det_r50_vd_east.yml +++ b/configs/det/det_r50_vd_east.yml @@ -20,7 +20,7 @@ Architecture: algorithm: EAST Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 50 Neck: name: EASTFPN diff --git a/configs/det/det_r50_vd_pse.yml b/configs/det/det_r50_vd_pse.yml index 8e77506c41..1a971564fd 100644 --- a/configs/det/det_r50_vd_pse.yml +++ b/configs/det/det_r50_vd_pse.yml @@ -20,7 +20,7 @@ Architecture: algorithm: PSE Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 50 Neck: name: FPN diff --git a/configs/det/det_res18_db_v2.0.yml b/configs/det/det_res18_db_v2.0.yml index 7b07ef9964..e983c221e2 100644 --- a/configs/det/det_res18_db_v2.0.yml +++ b/configs/det/det_res18_db_v2.0.yml @@ -20,7 +20,7 @@ Architecture: algorithm: DB Transform: Backbone: - name: ResNet + name: ResNet_vd layers: 18 disable_se: True Neck: From 142b5e9dfefe9f06a3b82067d8d0d50301d2f509 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:26:53 +0000 Subject: [PATCH 18/43] add db++ --- ppocr/data/imaug/operators.py | 7 +++ ppocr/modeling/backbones/__init__.py | 5 +- ppocr/modeling/backbones/det_resnet_vd.py | 12 +++-- ppocr/modeling/necks/db_fpn.py | 60 ++++++++++++++++++++++- ppocr/optimizer/learning_rate.py | 35 +++++++++++++ 5 files changed, 111 insertions(+), 8 deletions(-) diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index 09736515e7..a83dc64838 100644 --- a/ppocr/data/imaug/operators.py +++ b/ppocr/data/imaug/operators.py @@ -238,9 +238,12 @@ class DetResizeForTest(object): def __init__(self, **kwargs): super(DetResizeForTest, self).__init__() self.resize_type = 0 + self.keep_ratio = False if 'image_shape' in kwargs: self.image_shape = kwargs['image_shape'] self.resize_type = 1 + if 'keep_ratio' in kwargs: ###### + self.keep_ratio = kwargs['keep_ratio'] ####### elif 'limit_side_len' in kwargs: self.limit_side_len = kwargs['limit_side_len'] self.limit_type = kwargs.get('limit_type', 'min') @@ -270,6 +273,10 @@ class DetResizeForTest(object): def resize_image_type1(self, img): resize_h, resize_w = self.image_shape ori_h, ori_w = img.shape[:2] # (h, w, c) + if self.keep_ratio: ######## + resize_w = ori_w * resize_h / ori_h + N = math.ceil(resize_w / 32) + resize_w = N * 32 ratio_h = float(resize_h) / ori_h ratio_w = float(resize_w) / ori_w img = cv2.resize(img, (int(resize_w), int(resize_h))) diff --git a/ppocr/modeling/backbones/__init__.py b/ppocr/modeling/backbones/__init__.py index 072d6e0f84..0d8e60e93a 100755 --- a/ppocr/modeling/backbones/__init__.py +++ b/ppocr/modeling/backbones/__init__.py @@ -18,9 +18,10 @@ __all__ = ["build_backbone"] def build_backbone(config, model_type): if model_type == "det" or model_type == "table": from .det_mobilenet_v3 import MobileNetV3 - from .det_resnet_vd import ResNet + from .det_resnet import ResNet + from .det_resnet_vd import ResNet_vd from .det_resnet_vd_sast import ResNet_SAST - support_dict = ["MobileNetV3", "ResNet", "ResNet_SAST"] + support_dict = ["MobileNetV3", "ResNet", "ResNet_vd", "ResNet_SAST"] elif model_type == "rec" or model_type == "cls": from .rec_mobilenet_v3 import MobileNetV3 from .rec_resnet_vd import ResNet diff --git a/ppocr/modeling/backbones/det_resnet_vd.py b/ppocr/modeling/backbones/det_resnet_vd.py index 8c955a4af3..5337e14c40 100644 --- a/ppocr/modeling/backbones/det_resnet_vd.py +++ b/ppocr/modeling/backbones/det_resnet_vd.py @@ -25,7 +25,7 @@ from paddle.vision.ops import DeformConv2D from paddle.regularizer import L2Decay from paddle.nn.initializer import Normal, Constant, XavierUniform -__all__ = ["ResNet"] +__all__ = ["ResNet_vd", "ConvBNLayer", "DeformableConvV2"] class DeformableConvV2(nn.Layer): @@ -104,6 +104,7 @@ class ConvBNLayer(nn.Layer): kernel_size, stride=1, groups=1, + dcn_groups=1, is_vd_mode=False, act=None, is_dcn=False): @@ -128,7 +129,7 @@ class ConvBNLayer(nn.Layer): kernel_size=kernel_size, stride=stride, padding=(kernel_size - 1) // 2, - groups=2, #groups, + groups=dcn_groups, #groups, bias_attr=False) self._batch_norm = nn.BatchNorm(out_channels, act=act) @@ -162,7 +163,8 @@ class BottleneckBlock(nn.Layer): kernel_size=3, stride=stride, act='relu', - is_dcn=is_dcn) + is_dcn=is_dcn, + dcn_groups=2) self.conv2 = ConvBNLayer( in_channels=out_channels, out_channels=out_channels * 4, @@ -238,14 +240,14 @@ class BasicBlock(nn.Layer): return y -class ResNet(nn.Layer): +class ResNet_vd(nn.Layer): def __init__(self, in_channels=3, layers=50, dcn_stage=None, out_indices=None, **kwargs): - super(ResNet, self).__init__() + super(ResNet_vd, self).__init__() self.layers = layers supported_layers = [18, 34, 50, 101, 152, 200] diff --git a/ppocr/modeling/necks/db_fpn.py b/ppocr/modeling/necks/db_fpn.py index 93ed2dbfd1..b46d7c460d 100644 --- a/ppocr/modeling/necks/db_fpn.py +++ b/ppocr/modeling/necks/db_fpn.py @@ -105,9 +105,10 @@ class DSConv(nn.Layer): class DBFPN(nn.Layer): - def __init__(self, in_channels, out_channels, **kwargs): + def __init__(self, in_channels, out_channels, use_asf=None, **kwargs): super(DBFPN, self).__init__() self.out_channels = out_channels + self.use_asf = use_asf weight_attr = paddle.nn.initializer.KaimingUniform() self.in2_conv = nn.Conv2D( @@ -163,6 +164,9 @@ class DBFPN(nn.Layer): weight_attr=ParamAttr(initializer=weight_attr), bias_attr=False) + if self.use_asf: + self.asf = ASFBlock(self.out_channels, self.out_channels // 4) + def forward(self, x): c2, c3, c4, c5 = x @@ -187,6 +191,10 @@ class DBFPN(nn.Layer): p3 = F.upsample(p3, scale_factor=2, mode="nearest", align_mode=1) fuse = paddle.concat([p5, p4, p3, p2], axis=1) + + if self.use_asf: + fuse = self.asf(fuse, [p5, p4, p3, p2]) + return fuse @@ -356,3 +364,53 @@ class LKPAN(nn.Layer): fuse = paddle.concat([p5, p4, p3, p2], axis=1) return fuse + + +class ASFBlock(nn.Layer): + def __init__(self, in_channels, inter_channels, out_features_num=4): + super(ASFBlock, self).__init__() + weight_attr = paddle.nn.initializer.KaimingUniform() + self.in_channels = in_channels + self.inter_channels = inter_channels + self.out_features_num = out_features_num + self.conv = nn.Conv2D(in_channels, inter_channels, 3, padding=1) + + self.attention_block_1 = nn.Sequential( + #Nx1xHxW + nn.Conv2D( + 1, + 1, + 3, + bias_attr=False, + padding=1, + weight_attr=ParamAttr(initializer=weight_attr)), + nn.ReLU(), + nn.Conv2D( + 1, + 1, + 1, + bias_attr=False, + weight_attr=ParamAttr(initializer=weight_attr)), + nn.Sigmoid()) + + self.attention_block_2 = nn.Sequential( + nn.Conv2D( + inter_channels, + out_features_num, + 1, + bias_attr=False, + weight_attr=ParamAttr(initializer=weight_attr)), + nn.Sigmoid()) + + def forward(self, fuse_features, features_list): + fuse_features = self.conv(fuse_features) + attention_scores = self.attention_block_1( + paddle.mean( + fuse_features, axis=1, keepdim=True)) + fuse_features + attention_scores = self.attention_block_2(attention_scores) + assert len(features_list) == self.out_features_num + + out_list = [] + for i in range(self.out_features_num): + out_list.append(attention_scores[:, i:i + 1] * features_list[i]) + return paddle.concat(out_list, axis=1) diff --git a/ppocr/optimizer/learning_rate.py b/ppocr/optimizer/learning_rate.py index fe251f36e7..8e05a84011 100644 --- a/ppocr/optimizer/learning_rate.py +++ b/ppocr/optimizer/learning_rate.py @@ -308,3 +308,38 @@ class Const(object): end_lr=self.learning_rate, last_epoch=self.last_epoch) return learning_rate + + +class DecayLearningRate(object): + """ + DecayLearningRate learning rate decay + new_lr = (lr - end_lr) * (1 - epoch/decay_steps)**power + end_lr + Args: + learning_rate(float): initial learning rate + step_each_epoch(int): steps each epoch + epochs(int): total training epochs + factor(float): Power of polynomial, should greater than 0.0 to get learning rate decay. Default: 0.9 + end_lr(float): The minimum final learning rate. Default: 0.0. + """ + + def __init__(self, + learning_rate, + step_each_epoch, + epochs, + factor=0.9, + end_lr=0, + **kwargs): + super(DecayLearningRate, self).__init__() + self.learning_rate = learning_rate + self.epochs = epochs + 1 + self.factor = factor + self.end_lr = 0 + self.decay_steps = step_each_epoch * epochs + + def __call__(self): + learning_rate = lr.PolynomialDecay( + learning_rate=self.learning_rate, + decay_steps=self.decay_steps, + power=self.factor, + end_lr=self.end_lr) + return learning_rate From 26a89db76a71434dd1a759b1f4dcce36047b1fd0 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 08:27:26 +0000 Subject: [PATCH 19/43] add resnet --- configs/det/det_r50_db++_td_tr.yml | 166 +++++++++++++++++ ppocr/modeling/backbones/det_resnet.py | 237 +++++++++++++++++++++++++ 2 files changed, 403 insertions(+) create mode 100644 configs/det/det_r50_db++_td_tr.yml create mode 100644 ppocr/modeling/backbones/det_resnet.py diff --git a/configs/det/det_r50_db++_td_tr.yml b/configs/det/det_r50_db++_td_tr.yml new file mode 100644 index 0000000000..4954e1abb6 --- /dev/null +++ b/configs/det/det_r50_db++_td_tr.yml @@ -0,0 +1,166 @@ +Global: + debug: false + use_gpu: true + epoch_num: 1000 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/det_r50_td_tr/ + save_epoch_step: 200 + eval_batch_step: + - 0 + - 2000 + cal_metric_during_train: false + pretrained_model: ./pretrain_models/synthtext_pretrained_res50_dcn_asf_spatial + checkpoints: null + save_inference_dir: null + use_visualdl: false + infer_img: doc/imgs_en/img_10.jpg + save_res_path: ./checkpoints/det_db/predicts_db.txt +Architecture: + model_type: det + algorithm: DB + Transform: null + Backbone: + name: ResNet + layers: 50 + dcn_stage: [False, True, True, True] + Neck: + name: DBFPN + out_channels: 256 + use_asf: True + Head: + name: DBHead + k: 50 +Loss: + name: DBLoss + balance_loss: true + main_loss_type: BCELoss + alpha: 5 + beta: 10 + ohem_ratio: 3 +Optimizer: + name: Momentum + momentum: 0.9 + lr: + name: DecayLearningRate + learning_rate: 0.007 + epochs: 1000 + factor: 0.9 + end_lr: 0 + weight_decay: 0.0001 +PostProcess: + name: DBPostProcess + thresh: 0.3 + box_thresh: 0.5 + max_candidates: 1000 + unclip_ratio: 1.5 +Metric: + name: DetMetric + main_indicator: hmean +Train: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ + label_file_list: + - ./train_data/TD_TR/TD500/train_gt_labels.txt + - ./train_data/TD_TR/TR400/gt_labels.txt + ratio_list: + - 1.0 + - 1.0 + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - DetLabelEncode: null + - IaaAugment: + augmenter_args: + - type: Fliplr + args: + p: 0.5 + - type: Affine + args: + rotate: + - -10 + - 10 + - type: Resize + args: + size: + - 0.5 + - 3 + - EastRandomCropData: + size: + - 640 + - 640 + max_tries: 10 + keep_ratio: true + - MakeShrinkMap: + shrink_ratio: 0.4 + min_text_size: 8 + - MakeBorderMap: + shrink_ratio: 0.4 + thresh_min: 0.3 + thresh_max: 0.7 + - NormalizeImage: + scale: 1./255. + mean: + - 0.48109378172549 + - 0.45752457890196 + - 0.40787054090196 + std: + - 1.0 + - 1.0 + - 1.0 + order: hwc + - ToCHWImage: null + - KeepKeys: + keep_keys: + - image + - threshold_map + - threshold_mask + - shrink_map + - shrink_mask + loader: + shuffle: true + drop_last: false + batch_size_per_card: 4 + num_workers: 8 +Eval: + dataset: + name: SimpleDataSet + data_dir: ./train_data/ + label_file_list: + - ./train_data/TD_TR/TD500/test_gt_labels.txt + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - DetLabelEncode: null + - DetResizeForTest: + image_shape: + - 736 + - 736 + keep_ratio: True + - NormalizeImage: + scale: 1./255. + mean: + - 0.48109378172549 + - 0.45752457890196 + - 0.40787054090196 + std: + - 1.0 + - 1.0 + - 1.0 + order: hwc + - ToCHWImage: null + - KeepKeys: + keep_keys: + - image + - shape + - polys + - ignore_tags + loader: + shuffle: false + drop_last: false + batch_size_per_card: 1 + num_workers: 2 +profiler_options: null diff --git a/ppocr/modeling/backbones/det_resnet.py b/ppocr/modeling/backbones/det_resnet.py new file mode 100644 index 0000000000..4c75b663af --- /dev/null +++ b/ppocr/modeling/backbones/det_resnet.py @@ -0,0 +1,237 @@ +# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +import paddle +from paddle import ParamAttr +import paddle.nn as nn +import paddle.nn.functional as F +from paddle.nn import Conv2D, BatchNorm, Linear, Dropout +from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D +from paddle.nn.initializer import Uniform + +import math + +from paddle.vision.ops import DeformConv2D +from paddle.regularizer import L2Decay +from paddle.nn.initializer import Normal, Constant, XavierUniform +from .det_resnet_vd import DeformableConvV2, ConvBNLayer + + +class BottleneckBlock(nn.Layer): + def __init__(self, + num_channels, + num_filters, + stride, + shortcut=True, + is_dcn=False): + super(BottleneckBlock, self).__init__() + + self.conv0 = ConvBNLayer( + in_channels=num_channels, + out_channels=num_filters, + kernel_size=1, + act="relu", ) + self.conv1 = ConvBNLayer( + in_channels=num_filters, + out_channels=num_filters, + kernel_size=3, + stride=stride, + act="relu", + is_dcn=is_dcn, + dcn_groups=1, ) + self.conv2 = ConvBNLayer( + in_channels=num_filters, + out_channels=num_filters * 4, + kernel_size=1, + act=None, ) + + if not shortcut: + self.short = ConvBNLayer( + in_channels=num_channels, + out_channels=num_filters * 4, + kernel_size=1, + stride=stride, ) + + self.shortcut = shortcut + + self._num_channels_out = num_filters * 4 + + def forward(self, inputs): + y = self.conv0(inputs) + conv1 = self.conv1(y) + conv2 = self.conv2(conv1) + + if self.shortcut: + short = inputs + else: + short = self.short(inputs) + + y = paddle.add(x=short, y=conv2) + y = F.relu(y) + return y + + +class BasicBlock(nn.Layer): + def __init__(self, + num_channels, + num_filters, + stride, + shortcut=True, + name=None): + super(BasicBlock, self).__init__() + self.stride = stride + self.conv0 = ConvBNLayer( + in_channels=num_channels, + out_channels=num_filters, + kernel_size=3, + stride=stride, + act="relu") + self.conv1 = ConvBNLayer( + in_channels=num_filters, + out_channels=num_filters, + kernel_size=3, + act=None) + + if not shortcut: + self.short = ConvBNLayer( + in_channels=num_channels, + out_channels=num_filters, + kernel_size=1, + stride=stride) + + self.shortcut = shortcut + + def forward(self, inputs): + y = self.conv0(inputs) + conv1 = self.conv1(y) + + if self.shortcut: + short = inputs + else: + short = self.short(inputs) + y = paddle.add(x=short, y=conv1) + y = F.relu(y) + return y + + +class ResNet(nn.Layer): + def __init__(self, + in_channels=3, + layers=50, + out_indices=None, + dcn_stage=None): + super(ResNet, self).__init__() + + self.layers = layers + self.input_image_channel = in_channels + + supported_layers = [18, 34, 50, 101, 152] + assert layers in supported_layers, \ + "supported layers are {} but input layer is {}".format( + supported_layers, layers) + + if layers == 18: + depth = [2, 2, 2, 2] + elif layers == 34 or layers == 50: + depth = [3, 4, 6, 3] + elif layers == 101: + depth = [3, 4, 23, 3] + elif layers == 152: + depth = [3, 8, 36, 3] + num_channels = [64, 256, 512, + 1024] if layers >= 50 else [64, 64, 128, 256] + num_filters = [64, 128, 256, 512] + + self.dcn_stage = dcn_stage if dcn_stage is not None else [ + False, False, False, False + ] + self.out_indices = out_indices if out_indices is not None else [ + 0, 1, 2, 3 + ] + + self.conv = ConvBNLayer( + in_channels=self.input_image_channel, + out_channels=64, + kernel_size=7, + stride=2, + act="relu", ) + self.pool2d_max = MaxPool2D( + kernel_size=3, + stride=2, + padding=1, ) + + self.stages = [] + self.out_channels = [] + if layers >= 50: + for block in range(len(depth)): + shortcut = False + block_list = [] + is_dcn = self.dcn_stage[block] + for i in range(depth[block]): + if layers in [101, 152] and block == 2: + if i == 0: + conv_name = "res" + str(block + 2) + "a" + else: + conv_name = "res" + str(block + 2) + "b" + str(i) + else: + conv_name = "res" + str(block + 2) + chr(97 + i) + bottleneck_block = self.add_sublayer( + conv_name, + BottleneckBlock( + num_channels=num_channels[block] + if i == 0 else num_filters[block] * 4, + num_filters=num_filters[block], + stride=2 if i == 0 and block != 0 else 1, + shortcut=shortcut, + is_dcn=is_dcn)) + block_list.append(bottleneck_block) + shortcut = True + if block in self.out_indices: + self.out_channels.append(num_filters[block] * 4) + self.stages.append(nn.Sequential(*block_list)) + else: + for block in range(len(depth)): + shortcut = False + block_list = [] + # is_dcn = self.dcn_stage[block] + for i in range(depth[block]): + conv_name = "res" + str(block + 2) + chr(97 + i) + basic_block = self.add_sublayer( + conv_name, + BasicBlock( + num_channels=num_channels[block] + if i == 0 else num_filters[block], + num_filters=num_filters[block], + stride=2 if i == 0 and block != 0 else 1, + shortcut=shortcut)) + block_list.append(basic_block) + shortcut = True + if block in self.out_indices: + self.out_channels.append(num_filters[block]) + self.stages.append(nn.Sequential(*block_list)) + + def forward(self, inputs): + y = self.conv(inputs) + y = self.pool2d_max(y) + out = [] + for i, block in enumerate(self.stages): + y = block(y) + if i in self.out_indices: + out.append(y) + return out From 6e6d77432fb2f77af6b30e638f2d6eeef974ee12 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Mon, 27 Jun 2022 09:01:38 +0000 Subject: [PATCH 20/43] fix ci export bug --- configs/rec/rec_r45_abinet.yml | 2 -- configs/rec/rec_svtrnet.yml | 2 -- configs/rec/rec_svtrnet_ch.yml | 4 ---- ppocr/data/imaug/rec_img_aug.py | 17 ++--------------- test_tipc/configs/rec_svtrnet/rec_svtrnet.yml | 6 ++---- 5 files changed, 4 insertions(+), 27 deletions(-) diff --git a/configs/rec/rec_r45_abinet.yml b/configs/rec/rec_r45_abinet.yml index 3cf4cddb04..1e03db8b1b 100644 --- a/configs/rec/rec_r45_abinet.yml +++ b/configs/rec/rec_r45_abinet.yml @@ -72,7 +72,6 @@ Train: ignore_index: *ignore_index - ABINetRecResizeImg: image_shape: [3, 32, 128] - padding: False - KeepKeys: keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order loader: @@ -93,7 +92,6 @@ Eval: ignore_index: *ignore_index - ABINetRecResizeImg: image_shape: [3, 32, 128] - padding: False - KeepKeys: keep_keys: ['image', 'label', 'length'] # dataloader will return list in this order loader: diff --git a/configs/rec/rec_svtrnet.yml b/configs/rec/rec_svtrnet.yml index 5896b10515..c1f5cc380a 100644 --- a/configs/rec/rec_svtrnet.yml +++ b/configs/rec/rec_svtrnet.yml @@ -84,7 +84,6 @@ Train: channel_first: False - CTCLabelEncode: # Class handling label - SVTRRecResizeImg: - character_dict_path: image_shape: [3, 64, 256] padding: False - KeepKeys: @@ -105,7 +104,6 @@ Eval: channel_first: False - CTCLabelEncode: # Class handling label - SVTRRecResizeImg: - character_dict_path: image_shape: [3, 64, 256] padding: False - KeepKeys: diff --git a/configs/rec/rec_svtrnet_ch.yml b/configs/rec/rec_svtrnet_ch.yml index c4ce4ccc16..0d3f63d125 100644 --- a/configs/rec/rec_svtrnet_ch.yml +++ b/configs/rec/rec_svtrnet_ch.yml @@ -110,8 +110,6 @@ Train: - RecAug: null - CTCLabelEncode: null - SVTRRecResizeImg: - character_dict_path: null - infer_mode: False image_shape: - 3 - 32 @@ -139,8 +137,6 @@ Eval: channel_first: false - CTCLabelEncode: null - SVTRRecResizeImg: - character_dict_path: null - infer_mode: False image_shape: - 3 - 32 diff --git a/ppocr/data/imaug/rec_img_aug.py b/ppocr/data/imaug/rec_img_aug.py index 1de8e6ee3c..26773d0a51 100644 --- a/ppocr/data/imaug/rec_img_aug.py +++ b/ppocr/data/imaug/rec_img_aug.py @@ -312,14 +312,8 @@ class GrayRecResizeImg(object): class ABINetRecResizeImg(object): - def __init__(self, - image_shape, - infer_mode=False, - character_dict_path=None, - **kwargs): + def __init__(self, image_shape, **kwargs): self.image_shape = image_shape - self.infer_mode = infer_mode - self.character_dict_path = character_dict_path def __call__(self, data): img = data['image'] @@ -330,15 +324,8 @@ class ABINetRecResizeImg(object): class SVTRRecResizeImg(object): - def __init__(self, - image_shape, - infer_mode=False, - character_dict_path='./ppocr/utils/ppocr_keys_v1.txt', - padding=True, - **kwargs): + def __init__(self, image_shape, padding=True, **kwargs): self.image_shape = image_shape - self.infer_mode = infer_mode - self.character_dict_path = character_dict_path self.padding = padding def __call__(self, data): diff --git a/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml index bffd8ac09a..140b17e0e7 100644 --- a/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml +++ b/test_tipc/configs/rec_svtrnet/rec_svtrnet.yml @@ -84,8 +84,7 @@ Train: img_mode: BGR channel_first: False - CTCLabelEncode: # Class handling label - - RecResizeImg: - character_dict_path: + - SVTRRecResizeImg: image_shape: [3, 64, 256] padding: False - KeepKeys: @@ -106,8 +105,7 @@ Eval: img_mode: BGR channel_first: False - CTCLabelEncode: # Class handling label - - RecResizeImg: - character_dict_path: + - SVTRRecResizeImg: image_shape: [3, 64, 256] padding: False - KeepKeys: From 78064ad9560f479c3aa4def2c44cb0539c5a0bdc Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 09:43:14 +0000 Subject: [PATCH 21/43] add db++ --- ppocr/data/imaug/operators.py | 6 ++-- ppocr/modeling/backbones/det_resnet.py | 2 +- ppocr/modeling/necks/db_fpn.py | 47 ++++++++++++++++---------- 3 files changed, 33 insertions(+), 22 deletions(-) diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index a83dc64838..ee5492de79 100644 --- a/ppocr/data/imaug/operators.py +++ b/ppocr/data/imaug/operators.py @@ -242,8 +242,8 @@ class DetResizeForTest(object): if 'image_shape' in kwargs: self.image_shape = kwargs['image_shape'] self.resize_type = 1 - if 'keep_ratio' in kwargs: ###### - self.keep_ratio = kwargs['keep_ratio'] ####### + if 'keep_ratio' in kwargs: + self.keep_ratio = kwargs['keep_ratio'] elif 'limit_side_len' in kwargs: self.limit_side_len = kwargs['limit_side_len'] self.limit_type = kwargs.get('limit_type', 'min') @@ -273,7 +273,7 @@ class DetResizeForTest(object): def resize_image_type1(self, img): resize_h, resize_w = self.image_shape ori_h, ori_w = img.shape[:2] # (h, w, c) - if self.keep_ratio: ######## + if self.keep_ratio: resize_w = ori_w * resize_h / ori_h N = math.ceil(resize_w / 32) resize_w = N * 32 diff --git a/ppocr/modeling/backbones/det_resnet.py b/ppocr/modeling/backbones/det_resnet.py index 4c75b663af..3d8ce44508 100644 --- a/ppocr/modeling/backbones/det_resnet.py +++ b/ppocr/modeling/backbones/det_resnet.py @@ -1,4 +1,4 @@ -# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. +# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/ppocr/modeling/necks/db_fpn.py b/ppocr/modeling/necks/db_fpn.py index b46d7c460d..8c3f52a331 100644 --- a/ppocr/modeling/necks/db_fpn.py +++ b/ppocr/modeling/necks/db_fpn.py @@ -105,7 +105,7 @@ class DSConv(nn.Layer): class DBFPN(nn.Layer): - def __init__(self, in_channels, out_channels, use_asf=None, **kwargs): + def __init__(self, in_channels, out_channels, use_asf=False, **kwargs): super(DBFPN, self).__init__() self.out_channels = out_channels self.use_asf = use_asf @@ -164,7 +164,7 @@ class DBFPN(nn.Layer): weight_attr=ParamAttr(initializer=weight_attr), bias_attr=False) - if self.use_asf: + if self.use_asf is True: self.asf = ASFBlock(self.out_channels, self.out_channels // 4) def forward(self, x): @@ -192,7 +192,7 @@ class DBFPN(nn.Layer): fuse = paddle.concat([p5, p4, p3, p2], axis=1) - if self.use_asf: + if self.use_asf is True: fuse = self.asf(fuse, [p5, p4, p3, p2]) return fuse @@ -367,7 +367,19 @@ class LKPAN(nn.Layer): class ASFBlock(nn.Layer): + """ + This code is refered from: + https://github.com/MhLiao/DB/blob/master/decoders/feature_attention.py + """ + def __init__(self, in_channels, inter_channels, out_features_num=4): + """ + Adaptive Scale Fusion (ASF) block of DBNet++ + Args: + in_channels: the number of channels in the input data + inter_channels: the number of middle channels + out_features_num: the number of fused stages + """ super(ASFBlock, self).__init__() weight_attr = paddle.nn.initializer.KaimingUniform() self.in_channels = in_channels @@ -375,39 +387,38 @@ class ASFBlock(nn.Layer): self.out_features_num = out_features_num self.conv = nn.Conv2D(in_channels, inter_channels, 3, padding=1) - self.attention_block_1 = nn.Sequential( + self.spatial_scale = nn.Sequential( #Nx1xHxW nn.Conv2D( - 1, - 1, - 3, + in_channels=1, + out_channels=1, + kernel_size=3, bias_attr=False, padding=1, weight_attr=ParamAttr(initializer=weight_attr)), nn.ReLU(), nn.Conv2D( - 1, - 1, - 1, + in_channels=1, + out_channels=1, + kernel_size=1, bias_attr=False, weight_attr=ParamAttr(initializer=weight_attr)), nn.Sigmoid()) - self.attention_block_2 = nn.Sequential( + self.channel_scale = nn.Sequential( nn.Conv2D( - inter_channels, - out_features_num, - 1, + in_channels=inter_channels, + out_channels=out_features_num, + kernel_size=1, bias_attr=False, weight_attr=ParamAttr(initializer=weight_attr)), nn.Sigmoid()) def forward(self, fuse_features, features_list): fuse_features = self.conv(fuse_features) - attention_scores = self.attention_block_1( - paddle.mean( - fuse_features, axis=1, keepdim=True)) + fuse_features - attention_scores = self.attention_block_2(attention_scores) + spatial_x = paddle.mean(fuse_features, axis=1, keepdim=True) + attention_scores = self.spatial_scale(spatial_x) + fuse_features + attention_scores = self.channel_scale(attention_scores) assert len(features_list) == self.out_features_num out_list = [] From 679ec78bb2bc77bbc54c31b77c33074d34d100fe Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Mon, 27 Jun 2022 10:01:35 +0000 Subject: [PATCH 22/43] add db++ --- configs/det/{det_r50_db++_td_tr.yml => det_r50_db++.yml} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename configs/det/{det_r50_db++_td_tr.yml => det_r50_db++.yml} (100%) diff --git a/configs/det/det_r50_db++_td_tr.yml b/configs/det/det_r50_db++.yml similarity index 100% rename from configs/det/det_r50_db++_td_tr.yml rename to configs/det/det_r50_db++.yml From 0ae4f664612104227cf3e7ff13bc41b6d229dcfc Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 02:42:47 +0000 Subject: [PATCH 23/43] add db++ --- configs/det/det_r50_db++_ic15.yml | 163 ++++++++++++++++++ ...et_r50_db++.yml => det_r50_db++_td_tr.yml} | 0 2 files changed, 163 insertions(+) create mode 100644 configs/det/det_r50_db++_ic15.yml rename configs/det/{det_r50_db++.yml => det_r50_db++_td_tr.yml} (100%) diff --git a/configs/det/det_r50_db++_ic15.yml b/configs/det/det_r50_db++_ic15.yml new file mode 100644 index 0000000000..c1e8e86870 --- /dev/null +++ b/configs/det/det_r50_db++_ic15.yml @@ -0,0 +1,163 @@ +Global: + debug: false + use_gpu: true + epoch_num: 1000 + log_smooth_window: 20 + print_batch_step: 10 + save_model_dir: ./output/det_r50_icdar15/ + save_epoch_step: 200 + eval_batch_step: + - 0 + - 2000 + cal_metric_during_train: false + pretrained_model: ./pretrain_models/synthtext_pretrained_res50_dcn_asf_spatial + checkpoints: null + save_inference_dir: null + use_visualdl: false + infer_img: doc/imgs_en/img_10.jpg + save_res_path: ./checkpoints/det_db/predicts_db.txt +Architecture: + model_type: det + algorithm: DB + Transform: null + Backbone: + name: ResNet + layers: 50 + dcn_stage: [False, True, True, True] + Neck: + name: DBFPN + out_channels: 256 + use_asf: True + Head: + name: DBHead + k: 50 +Loss: + name: DBLoss + balance_loss: true + main_loss_type: BCELoss + alpha: 5 + beta: 10 + ohem_ratio: 3 +Optimizer: + name: Momentum + momentum: 0.9 + lr: + name: DecayLearningRate + learning_rate: 0.007 + epochs: 1000 + factor: 0.9 + end_lr: 0 + weight_decay: 0.0001 +PostProcess: + name: DBPostProcess + thresh: 0.3 + box_thresh: 0.6 + max_candidates: 1000 + unclip_ratio: 1.5 +Metric: + name: DetMetric + main_indicator: hmean +Train: + dataset: + name: SimpleDataSet + data_dir: ./train_data/icdar2015/text_localization/ + label_file_list: + - ./train_data/icdar2015/text_localization/train_icdar2015_label.txt + ratio_list: + - 1.0 + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - DetLabelEncode: null + - IaaAugment: + augmenter_args: + - type: Fliplr + args: + p: 0.5 + - type: Affine + args: + rotate: + - -10 + - 10 + - type: Resize + args: + size: + - 0.5 + - 3 + - EastRandomCropData: + size: + - 640 + - 640 + max_tries: 10 + keep_ratio: true + - MakeShrinkMap: + shrink_ratio: 0.4 + min_text_size: 8 + - MakeBorderMap: + shrink_ratio: 0.4 + thresh_min: 0.3 + thresh_max: 0.7 + - NormalizeImage: + scale: 1./255. + mean: + - 0.48109378172549 + - 0.45752457890196 + - 0.40787054090196 + std: + - 1.0 + - 1.0 + - 1.0 + order: hwc + - ToCHWImage: null + - KeepKeys: + keep_keys: + - image + - threshold_map + - threshold_mask + - shrink_map + - shrink_mask + loader: + shuffle: true + drop_last: false + batch_size_per_card: 4 + num_workers: 8 +Eval: + dataset: + name: SimpleDataSet + data_dir: ./train_data/icdar2015/text_localization + label_file_list: + - ./train_data/icdar2015/text_localization/test_icdar2015_label.txt + transforms: + - DecodeImage: + img_mode: BGR + channel_first: false + - DetLabelEncode: null + - DetResizeForTest: + image_shape: + - 1152 + - 2048 + - NormalizeImage: + scale: 1./255. + mean: + - 0.48109378172549 + - 0.45752457890196 + - 0.40787054090196 + std: + - 1.0 + - 1.0 + - 1.0 + order: hwc + - ToCHWImage: null + - KeepKeys: + keep_keys: + - image + - shape + - polys + - ignore_tags + loader: + shuffle: false + drop_last: false + batch_size_per_card: 1 + num_workers: 2 +profiler_options: null diff --git a/configs/det/det_r50_db++.yml b/configs/det/det_r50_db++_td_tr.yml similarity index 100% rename from configs/det/det_r50_db++.yml rename to configs/det/det_r50_db++_td_tr.yml From 03a6ef531f82eda961183ba38e7473b8cfc9c3fa Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Tue, 28 Jun 2022 03:27:48 +0000 Subject: [PATCH 24/43] fix abinet ci bug --- configs/rec/rec_r45_abinet.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/configs/rec/rec_r45_abinet.yml b/configs/rec/rec_r45_abinet.yml index 1e03db8b1b..e604bcd12d 100644 --- a/configs/rec/rec_r45_abinet.yml +++ b/configs/rec/rec_r45_abinet.yml @@ -8,7 +8,7 @@ Global: # evaluation is run every 2000 iterations eval_batch_step: [0, 2000] cal_metric_during_train: True - pretrained_model: ./rec_r45_abinet_train/abinet_vl_pretrained + pretrained_model: checkpoints: save_inference_dir: use_visualdl: False @@ -41,7 +41,6 @@ Architecture: Transform: Backbone: name: ResNet45 - Head: name: ABINetHead use_lang: True From a63f4414e2cdf3dfd42cddaeb9dcd3b0d412ffd0 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 03:28:53 +0000 Subject: [PATCH 25/43] add db++ --- ppocr/data/imaug/operators.py | 2 +- ppocr/modeling/backbones/det_resnet_vd.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ppocr/data/imaug/operators.py b/ppocr/data/imaug/operators.py index ee5492de79..070fb7afa0 100644 --- a/ppocr/data/imaug/operators.py +++ b/ppocr/data/imaug/operators.py @@ -273,7 +273,7 @@ class DetResizeForTest(object): def resize_image_type1(self, img): resize_h, resize_w = self.image_shape ori_h, ori_w = img.shape[:2] # (h, w, c) - if self.keep_ratio: + if self.keep_ratio is True: resize_w = ori_w * resize_h / ori_h N = math.ceil(resize_w / 32) resize_w = N * 32 diff --git a/ppocr/modeling/backbones/det_resnet_vd.py b/ppocr/modeling/backbones/det_resnet_vd.py index 5337e14c40..a421da0ab4 100644 --- a/ppocr/modeling/backbones/det_resnet_vd.py +++ b/ppocr/modeling/backbones/det_resnet_vd.py @@ -323,7 +323,6 @@ class ResNet_vd(nn.Layer): for block in range(len(depth)): block_list = [] shortcut = False - # is_dcn = self.dcn_stage[block] for i in range(depth[block]): basic_block = self.add_sublayer( 'bb_%d_%d' % (block, i), From 9422629cc05e0fdf6ead09d6b03a18424b131043 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 04:17:14 +0000 Subject: [PATCH 26/43] add db++ --- configs/det/det_r50_db++_ic15.yml | 2 +- configs/det/det_r50_db++_td_tr.yml | 2 +- doc/doc_ch/algorithm_det_db.md | 24 ++++++++++++++++++++++-- doc/doc_ch/dataset/ocr_datasets.md | 1 + 4 files changed, 25 insertions(+), 4 deletions(-) diff --git a/configs/det/det_r50_db++_ic15.yml b/configs/det/det_r50_db++_ic15.yml index c1e8e86870..a90c868cc3 100644 --- a/configs/det/det_r50_db++_ic15.yml +++ b/configs/det/det_r50_db++_ic15.yml @@ -10,7 +10,7 @@ Global: - 0 - 2000 cal_metric_during_train: false - pretrained_model: ./pretrain_models/synthtext_pretrained_res50_dcn_asf_spatial + pretrained_model: ./pretrain_models/ResNet50_dcn_asf_synthtext_pretrained checkpoints: null save_inference_dir: null use_visualdl: false diff --git a/configs/det/det_r50_db++_td_tr.yml b/configs/det/det_r50_db++_td_tr.yml index 4954e1abb6..5e26ddb0ac 100644 --- a/configs/det/det_r50_db++_td_tr.yml +++ b/configs/det/det_r50_db++_td_tr.yml @@ -10,7 +10,7 @@ Global: - 0 - 2000 cal_metric_during_train: false - pretrained_model: ./pretrain_models/synthtext_pretrained_res50_dcn_asf_spatial + pretrained_model: ./pretrain_models/ResNet50_dcn_asf_synthtext_pretrained checkpoints: null save_inference_dir: null use_visualdl: false diff --git a/doc/doc_ch/algorithm_det_db.md b/doc/doc_ch/algorithm_det_db.md index 90837c2ac1..8466565563 100644 --- a/doc/doc_ch/algorithm_det_db.md +++ b/doc/doc_ch/algorithm_det_db.md @@ -1,4 +1,4 @@ -# DB +# DB与DB++ - [1. 算法简介](#1) - [2. 环境配置](#2) @@ -21,12 +21,24 @@ > Liao, Minghui and Wan, Zhaoyi and Yao, Cong and Chen, Kai and Bai, Xiang > AAAI, 2020 +> [Real-Time Scene Text Detection with Differentiable Binarization and Adaptive Scale Fusion](https://arxiv.org/abs/2202.10304) +> Liao, Minghui and Zou, Zhisheng and Wan, Zhaoyi and Yao, Cong and Bai, Xiang +> TPAMI, 2022 + + 在ICDAR2015文本检测公开数据集上，算法复现效果如下： |模型|骨干网络|配置文件|precision|recall|Hmean|下载链接| | --- | --- | --- | --- | --- | --- | --- | |DB|ResNet50_vd|[configs/det/det_r50_vd_db.yml](../../configs/det/det_r50_vd_db.yml)|86.41%|78.72%|82.38%|[训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/det_r50_vd_db_v2.0_train.tar)| |DB|MobileNetV3|[configs/det/det_mv3_db.yml](../../configs/det/det_mv3_db.yml)|77.29%|73.08%|75.12%|[训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/det_mv3_db_v2.0_train.tar)| +|DB++|ResNet50|[configs/det/det_r50_db++_ic15.yml](../../configs/det/det_r50_db++_ic15.yml)|90.89%|82.66%|86.58%|[合成数据预训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/en_det/ResNet50_dcn_asf_synthtext_pretrained.pdparams)/[训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/en_det/det_r50_db%2B%2B_icdar15_train.tar)| + +在TD_TR文本检测公开数据集上，算法复现效果如下： + +|模型|骨干网络|配置文件|precision|recall|Hmean|下载链接| +| --- | --- | --- | --- | --- | --- | --- | +|DB++|ResNet50|[configs/det/det_r50_db++_td_tr.yml](../../configs/det/det_r50_db++_td_tr.yml)|92.92%|86.48%|89.58%|[合成数据预训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/en_det/ResNet50_dcn_asf_synthtext_pretrained.pdparams)/[训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/en_det/det_r50_db%2B%2B_td_tr_train.tar)| @@ -96,4 +108,12 @@ DB模型还支持以下推理部署方式： pages={11474--11481}, year={2020} } -``` \ No newline at end of file + +@article{liao2022real, + title={Real-Time Scene Text Detection with Differentiable Binarization and Adaptive Scale Fusion}, + author={Liao, Minghui and Zou, Zhisheng and Wan, Zhaoyi and Yao, Cong and Bai, Xiang}, + journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, + year={2022}, + publisher={IEEE} +} +``` diff --git a/doc/doc_ch/dataset/ocr_datasets.md b/doc/doc_ch/dataset/ocr_datasets.md index c6ff2e170f..b7666fd63e 100644 --- a/doc/doc_ch/dataset/ocr_datasets.md +++ b/doc/doc_ch/dataset/ocr_datasets.md @@ -34,6 +34,7 @@ json.dumps编码前的图像标注信息是包含多个字典的list，字典中 | ICDAR 2015 |https://rrc.cvc.uab.es/?ch=4&com=downloads| [train](https://paddleocr.bj.bcebos.com/dataset/train_icdar2015_label.txt) / [test](https://paddleocr.bj.bcebos.com/dataset/test_icdar2015_label.txt) | | ctw1500 |https://paddleocr.bj.bcebos.com/dataset/ctw1500.zip| 图片下载地址中已包含 | | total text |https://paddleocr.bj.bcebos.com/dataset/total_text.tar| 图片下载地址中已包含 | +| td tr |https://paddleocr.bj.bcebos.com/dataset/TD_TR.tar| 图片下载地址中已包含 | #### 1.2.1 ICDAR 2015 ICDAR 2015 数据集包含1000张训练图像和500张测试图像。ICDAR 2015 数据集可以从上表中链接下载，首次下载需注册。 From 1315cdfc8633ebb501f56c16a0b5dff3635eb7af Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 06:17:45 +0000 Subject: [PATCH 27/43] add db++ --- ppocr/modeling/backbones/det_resnet.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ppocr/modeling/backbones/det_resnet.py b/ppocr/modeling/backbones/det_resnet.py index 3d8ce44508..87eef11cf0 100644 --- a/ppocr/modeling/backbones/det_resnet.py +++ b/ppocr/modeling/backbones/det_resnet.py @@ -209,7 +209,6 @@ class ResNet(nn.Layer): for block in range(len(depth)): shortcut = False block_list = [] - # is_dcn = self.dcn_stage[block] for i in range(depth[block]): conv_name = "res" + str(block + 2) + chr(97 + i) basic_block = self.add_sublayer( From 04e710419468ad25afa5c8e478d9c88c23ea0d5c Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 11:45:21 +0000 Subject: [PATCH 28/43] add db++ --- configs/det/det_r50_db++_ic15.yml | 2 +- configs/det/det_r50_db++_td_tr.yml | 2 +- tools/infer/predict_det.py | 19 +++++++++++++++++- tools/program.py | 31 ++++++++++++++++++++---------- 4 files changed, 41 insertions(+), 13 deletions(-) diff --git a/configs/det/det_r50_db++_ic15.yml b/configs/det/det_r50_db++_ic15.yml index a90c868cc3..e0cd6012b6 100644 --- a/configs/det/det_r50_db++_ic15.yml +++ b/configs/det/det_r50_db++_ic15.yml @@ -18,7 +18,7 @@ Global: save_res_path: ./checkpoints/det_db/predicts_db.txt Architecture: model_type: det - algorithm: DB + algorithm: DB++ Transform: null Backbone: name: ResNet diff --git a/configs/det/det_r50_db++_td_tr.yml b/configs/det/det_r50_db++_td_tr.yml index 5e26ddb0ac..65021bb661 100644 --- a/configs/det/det_r50_db++_td_tr.yml +++ b/configs/det/det_r50_db++_td_tr.yml @@ -18,7 +18,7 @@ Global: save_res_path: ./checkpoints/det_db/predicts_db.txt Architecture: model_type: det - algorithm: DB + algorithm: DB++ Transform: null Backbone: name: ResNet diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py index 7b6bebf1fb..394a48948b 100755 --- a/tools/infer/predict_det.py +++ b/tools/infer/predict_det.py @@ -67,6 +67,23 @@ class TextDetector(object): postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio postprocess_params["use_dilation"] = args.use_dilation postprocess_params["score_mode"] = args.det_db_score_mode + elif self.det_algorithm == "DB++": + postprocess_params['name'] = 'DBPostProcess' + postprocess_params["thresh"] = args.det_db_thresh + postprocess_params["box_thresh"] = args.det_db_box_thresh + postprocess_params["max_candidates"] = 1000 + postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio + postprocess_params["use_dilation"] = args.use_dilation + postprocess_params["score_mode"] = args.det_db_score_mode + pre_process_list[1] = { + 'NormalizeImage': { + 'std': [1.0, 1.0, 1.0], + 'mean': + [0.48109378172549, 0.45752457890196, 0.40787054090196], + 'scale': '1./255.', + 'order': 'hwc' + } + } elif self.det_algorithm == "EAST": postprocess_params['name'] = 'EASTPostProcess' postprocess_params["score_thresh"] = args.det_east_score_thresh @@ -231,7 +248,7 @@ class TextDetector(object): preds['f_score'] = outputs[1] preds['f_tco'] = outputs[2] preds['f_tvo'] = outputs[3] - elif self.det_algorithm in ['DB', 'PSE']: + elif self.det_algorithm in ['DB', 'PSE', 'DB++']: preds['maps'] = outputs[0] elif self.det_algorithm == 'FCE': for i, output in enumerate(outputs): diff --git a/tools/program.py b/tools/program.py index aa0d2698cf..620c61f093 100755 --- a/tools/program.py +++ b/tools/program.py @@ -307,7 +307,8 @@ def train(config, train_stats.update(stats) if log_writer is not None and dist.get_rank() == 0: - log_writer.log_metrics(metrics=train_stats.get(), prefix="TRAIN", step=global_step) + log_writer.log_metrics( + metrics=train_stats.get(), prefix="TRAIN", step=global_step) if dist.get_rank() == 0 and ( (global_step > 0 and global_step % print_batch_step == 0) or @@ -354,7 +355,8 @@ def train(config, # logger metric if log_writer is not None: - log_writer.log_metrics(metrics=cur_metric, prefix="EVAL", step=global_step) + log_writer.log_metrics( + metrics=cur_metric, prefix="EVAL", step=global_step) if cur_metric[main_indicator] >= best_model_dict[ main_indicator]: @@ -377,11 +379,18 @@ def train(config, logger.info(best_str) # logger best metric if log_writer is not None: - log_writer.log_metrics(metrics={ - "best_{}".format(main_indicator): best_model_dict[main_indicator] - }, prefix="EVAL", step=global_step) - - log_writer.log_model(is_best=True, prefix="best_accuracy", metadata=best_model_dict) + log_writer.log_metrics( + metrics={ + "best_{}".format(main_indicator): + best_model_dict[main_indicator] + }, + prefix="EVAL", + step=global_step) + + log_writer.log_model( + is_best=True, + prefix="best_accuracy", + metadata=best_model_dict) reader_start = time.time() if dist.get_rank() == 0: @@ -413,7 +422,8 @@ def train(config, epoch=epoch, global_step=global_step) if log_writer is not None: - log_writer.log_model(is_best=False, prefix='iter_epoch_{}'.format(epoch)) + log_writer.log_model( + is_best=False, prefix='iter_epoch_{}'.format(epoch)) best_str = 'best metric, {}'.format(', '.join( ['{}: {}'.format(k, v) for k, v in best_model_dict.items()])) @@ -564,7 +574,7 @@ def preprocess(is_train=False): assert alg in [ 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', - 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR' + 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', 'DB++' ] if use_xpu: @@ -585,7 +595,8 @@ def preprocess(is_train=False): vdl_writer_path = '{}/vdl/'.format(save_model_dir) log_writer = VDLLogger(save_model_dir) loggers.append(log_writer) - if ('use_wandb' in config['Global'] and config['Global']['use_wandb']) or 'wandb' in config: + if ('use_wandb' in config['Global'] and + config['Global']['use_wandb']) or 'wandb' in config: save_dir = config['Global']['save_model_dir'] wandb_writer_path = "{}/wandb".format(save_dir) if "wandb" in config: From 9bb4ead8eae06b1be3366594a61e416b610adf51 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Tue, 28 Jun 2022 11:51:55 +0000 Subject: [PATCH 29/43] add db++ --- doc/doc_ch/algorithm_det_db.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/doc_ch/algorithm_det_db.md b/doc/doc_ch/algorithm_det_db.md index 8466565563..afdddb1a73 100644 --- a/doc/doc_ch/algorithm_det_db.md +++ b/doc/doc_ch/algorithm_det_db.md @@ -66,7 +66,7 @@ python3 tools/export_model.py -c configs/det/det_r50_vd_db.yml -o Global.pretrai DB文本检测模型推理，可以执行如下命令： ```shell -python3 tools/infer/predict_det.py --image_dir="./doc/imgs_en/img_10.jpg" --det_model_dir="./inference/det_db/" +python3 tools/infer/predict_det.py --image_dir="./doc/imgs_en/img_10.jpg" --det_model_dir="./inference/det_db/" --det_algorithm="DB" ``` 可视化文本检测结果默认保存到`./inference_results`文件夹里面，结果文件的名称前缀为'det_res'。结果示例如下： From 9da6038214ce90094ebb09c88f84fe73f4845092 Mon Sep 17 00:00:00 2001 From: Topdu <784990967@qq.com> Date: Wed, 29 Jun 2022 12:08:45 +0000 Subject: [PATCH 30/43] fix BaseAug --- doc/doc_ch/algorithm_overview.md | 4 ++-- doc/doc_en/algorithm_overview_en.md | 4 ++-- ppocr/data/imaug/__init__.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/doc/doc_ch/algorithm_overview.md b/doc/doc_ch/algorithm_overview.md index eb81e4cd6d..fbf0129ddd 100755 --- a/doc/doc_ch/algorithm_overview.md +++ b/doc/doc_ch/algorithm_overview.md @@ -86,8 +86,8 @@ |SAR|Resnet31| 87.20% | rec_r31_sar | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [训练模型](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [训练模型](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | -|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar) | -|ABINet|Resnet45| 90.75% | rec_r45_abinet_en | [训练模型](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce | [训练模型](https://paddleocr.bj.bcebos.com/rec_vitstr_none_ce_train.tar) | +|ABINet|Resnet45| 90.75% | rec_r45_abinet | [训练模型](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | diff --git a/doc/doc_en/algorithm_overview_en.md b/doc/doc_en/algorithm_overview_en.md index 28aca7c0d1..a3fa7a2d91 100755 --- a/doc/doc_en/algorithm_overview_en.md +++ b/doc/doc_en/algorithm_overview_en.md @@ -85,8 +85,8 @@ Refer to [DTRB](https://arxiv.org/abs/1904.01906), the training and evaluation r |SAR|Resnet31| 87.20% | rec_r31_sar | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_r31_sar_train.tar) | |SEED|Aster_Resnet| 85.35% | rec_resnet_stn_bilstm_att | [trained model](https://paddleocr.bj.bcebos.com/dygraph_v2.1/rec/rec_resnet_stn_bilstm_att.tar) | |SVTR|SVTR-Tiny| 89.25% | rec_svtr_tiny_none_ctc_en | [trained model](https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/rec_svtr_tiny_none_ctc_en_train.tar) | -|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce_en | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | -|ABINet|Resnet45| 90.75% | rec_r45_abinet_en | [trained model](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | +|ViTSTR|ViTSTR| 79.82% | rec_vitstr_none_ce | [trained model](https://paddleocr.bj.bcebos.com/rec_vitstr_none_none_train.tar) | +|ABINet|Resnet45| 90.75% | rec_r45_abinet | [trained model](https://paddleocr.bj.bcebos.com/rec_r45_abinet_train.tar) | diff --git a/ppocr/data/imaug/__init__.py b/ppocr/data/imaug/__init__.py index 63dfda91f8..c958c7ae4e 100644 --- a/ppocr/data/imaug/__init__.py +++ b/ppocr/data/imaug/__init__.py @@ -23,9 +23,9 @@ from .random_crop_data import EastRandomCropData, RandomCropImgMask from .make_pse_gt import MakePseGt -from .rec_img_aug import RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ - SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, \ - ABINetRecResizeImg, SVTRRecResizeImg, ABINetRecAug +from .rec_img_aug import BaseDataAugmentation, RecAug, RecConAug, RecResizeImg, ClsResizeImg, \ + SRNRecResizeImg, GrayRecResizeImg, SARRecResizeImg, PRENResizeImg, \ + ABINetRecResizeImg, SVTRRecResizeImg, ABINetRecAug from .ssl_img_aug import SSLRotateResize from .randaugment import RandAugment from .copy_paste import CopyPaste From d915416da9cca4b7af8642cdc7279f841343e9cc Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Thu, 30 Jun 2022 04:25:26 +0000 Subject: [PATCH 31/43] add db++ --- tools/program.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/program.py b/tools/program.py index 620c61f093..d86be5badb 100755 --- a/tools/program.py +++ b/tools/program.py @@ -574,7 +574,8 @@ def preprocess(is_train=False): assert alg in [ 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', - 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', 'DB++' + 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', + 'ViTSTR', 'ABINet', 'DB++' ] if use_xpu: From e13ec733a628006cf7295652600bbf8c1604755f Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Thu, 30 Jun 2022 15:23:31 +0800 Subject: [PATCH 32/43] unify kie and ser for vqa data format (#6704) * unify kie and ser for vqa data format * fix config and label ops * fix doc * add distort bbox --- configs/kie/kie_unet_sdmgr.yml | 4 +-- configs/vqa/ser/layoutlm.yml | 6 ++-- configs/vqa/ser/layoutlmv2.yml | 6 ++-- configs/vqa/ser/layoutxlm.yml | 6 ++-- ppocr/data/imaug/label_ops.py | 46 +++++++++++++++--------- ppocr/data/imaug/vqa/__init__.py | 7 +++- ppocr/data/imaug/vqa/augment.py | 37 +++++++++++++++++++ ppocr/utils/utility.py | 25 ++++++------- ppocr/utils/visual.py | 28 +++++++++++---- ppstructure/docs/kie.md | 2 +- ppstructure/docs/kie_en.md | 2 +- ppstructure/vqa/README.md | 4 +-- ppstructure/vqa/README_ch.md | 4 +-- ppstructure/vqa/labels/labels_ser.txt | 3 -- ppstructure/vqa/tools/trans_xfun_data.py | 18 ++++------ tools/infer_kie.py | 14 ++++---- tools/infer_vqa_token_ser.py | 30 ++++++++++++---- tools/program.py | 4 +-- 18 files changed, 161 insertions(+), 85 deletions(-) create mode 100644 ppocr/data/imaug/vqa/augment.py delete mode 100644 ppstructure/vqa/labels/labels_ser.txt diff --git a/configs/kie/kie_unet_sdmgr.yml b/configs/kie/kie_unet_sdmgr.yml index a6968aaa3a..da2e4fda50 100644 --- a/configs/kie/kie_unet_sdmgr.yml +++ b/configs/kie/kie_unet_sdmgr.yml @@ -17,7 +17,7 @@ Global: checkpoints: save_inference_dir: use_visualdl: False - class_path: ./train_data/wildreceipt/class_list.txt + class_path: &class_path ./train_data/wildreceipt/class_list.txt infer_img: ./train_data/wildreceipt/1.txt save_res_path: ./output/sdmgr_kie/predicts_kie.txt img_scale: [ 1024, 512 ] @@ -72,6 +72,7 @@ Train: order: 'hwc' - KieLabelEncode: # Class handling label character_dict_path: ./train_data/wildreceipt/dict.txt + class_path: *class_path - KieResize: - ToCHWImage: - KeepKeys: @@ -88,7 +89,6 @@ Eval: data_dir: ./train_data/wildreceipt label_file_list: - ./train_data/wildreceipt/wildreceipt_test.txt - # - /paddle/data/PaddleOCR/train_data/wildreceipt/1.txt transforms: - DecodeImage: # load image img_mode: RGB diff --git a/configs/vqa/ser/layoutlm.yml b/configs/vqa/ser/layoutlm.yml index 87131170c9..47ab093e1f 100644 --- a/configs/vqa/ser/layoutlm.yml +++ b/configs/vqa/ser/layoutlm.yml @@ -43,7 +43,7 @@ Optimizer: PostProcess: name: VQASerTokenLayoutLMPostProcess - class_path: &class_path ppstructure/vqa/labels/labels_ser.txt + class_path: &class_path train_data/XFUND/class_list_xfun.txt Metric: name: VQASerTokenMetric @@ -54,7 +54,7 @@ Train: name: SimpleDataSet data_dir: train_data/XFUND/zh_train/image label_file_list: - - train_data/XFUND/zh_train/xfun_normalize_train.json + - train_data/XFUND/zh_train/train.json transforms: - DecodeImage: # load image img_mode: RGB @@ -89,7 +89,7 @@ Eval: name: SimpleDataSet data_dir: train_data/XFUND/zh_val/image label_file_list: - - train_data/XFUND/zh_val/xfun_normalize_val.json + - train_data/XFUND/zh_val/val.json transforms: - DecodeImage: # load image img_mode: RGB diff --git a/configs/vqa/ser/layoutlmv2.yml b/configs/vqa/ser/layoutlmv2.yml index 33406252b3..d6a9c03e5e 100644 --- a/configs/vqa/ser/layoutlmv2.yml +++ b/configs/vqa/ser/layoutlmv2.yml @@ -44,7 +44,7 @@ Optimizer: PostProcess: name: VQASerTokenLayoutLMPostProcess - class_path: &class_path ppstructure/vqa/labels/labels_ser.txt + class_path: &class_path train_data/XFUND/class_list_xfun.txt Metric: name: VQASerTokenMetric @@ -55,7 +55,7 @@ Train: name: SimpleDataSet data_dir: train_data/XFUND/zh_train/image label_file_list: - - train_data/XFUND/zh_train/xfun_normalize_train.json + - train_data/XFUND/zh_train/train.json transforms: - DecodeImage: # load image img_mode: RGB @@ -90,7 +90,7 @@ Eval: name: SimpleDataSet data_dir: train_data/XFUND/zh_val/image label_file_list: - - train_data/XFUND/zh_val/xfun_normalize_val.json + - train_data/XFUND/zh_val/val.json transforms: - DecodeImage: # load image img_mode: RGB diff --git a/configs/vqa/ser/layoutxlm.yml b/configs/vqa/ser/layoutxlm.yml index eb1cca5a21..3686989ccf 100644 --- a/configs/vqa/ser/layoutxlm.yml +++ b/configs/vqa/ser/layoutxlm.yml @@ -11,7 +11,7 @@ Global: save_inference_dir: use_visualdl: False seed: 2022 - infer_img: doc/vqa/input/zh_val_42.jpg + infer_img: ppstructure/docs/vqa/input/zh_val_42.jpg save_res_path: ./output/ser Architecture: @@ -54,7 +54,7 @@ Train: name: SimpleDataSet data_dir: train_data/XFUND/zh_train/image label_file_list: - - train_data/XFUND/zh_train/xfun_normalize_train.json + - train_data/XFUND/zh_train/train.json ratio_list: [ 1.0 ] transforms: - DecodeImage: # load image @@ -90,7 +90,7 @@ Eval: name: SimpleDataSet data_dir: train_data/XFUND/zh_val/image label_file_list: - - train_data/XFUND/zh_val/xfun_normalize_val.json + - train_data/XFUND/zh_val/val.json transforms: - DecodeImage: # load image img_mode: RGB diff --git a/ppocr/data/imaug/label_ops.py b/ppocr/data/imaug/label_ops.py index 312d6dc9ad..c95b326224 100644 --- a/ppocr/data/imaug/label_ops.py +++ b/ppocr/data/imaug/label_ops.py @@ -259,15 +259,26 @@ class E2ELabelEncodeTrain(object): class KieLabelEncode(object): - def __init__(self, character_dict_path, norm=10, directed=False, **kwargs): + def __init__(self, + character_dict_path, + class_path, + norm=10, + directed=False, + **kwargs): super(KieLabelEncode, self).__init__() self.dict = dict({'': 0}) + self.label2classid_map = dict() with open(character_dict_path, 'r', encoding='utf-8') as fr: idx = 1 for line in fr: char = line.strip() self.dict[char] = idx idx += 1 + with open(class_path, "r") as fin: + lines = fin.readlines() + for idx, line in enumerate(lines): + line = line.strip("\n") + self.label2classid_map[line] = idx self.norm = norm self.directed = directed @@ -408,7 +419,7 @@ class KieLabelEncode(object): text_ind = [self.dict[c] for c in text if c in self.dict] text_inds.append(text_ind) if 'label' in ann.keys(): - labels.append(ann['label']) + labels.append(self.label2classid_map[ann['label']]) elif 'key_cls' in ann.keys(): labels.append(ann['key_cls']) else: @@ -876,15 +887,16 @@ class VQATokenLabelEncode(object): for info in ocr_info: if train_re: # for re - if len(info["text"]) == 0: + if len(info["transcription"]) == 0: empty_entity.add(info["id"]) continue id2label[info["id"]] = info["label"] relations.extend([tuple(sorted(l)) for l in info["linking"]]) # smooth_box + info["bbox"] = self.trans_poly_to_bbox(info["points"]) bbox = self._smooth_box(info["bbox"], height, width) - text = info["text"] + text = info["transcription"] encode_res = self.tokenizer.encode( text, pad_to_max_seq_len=False, return_attention_mask=True) @@ -900,7 +912,7 @@ class VQATokenLabelEncode(object): label = info['label'] gt_label = self._parse_label(label, encode_res) - # construct entities for re +# construct entities for re if train_re: if gt_label[0] != self.label2id_map["O"]: entity_id_to_index_map[info["id"]] = len(entities) @@ -944,29 +956,29 @@ class VQATokenLabelEncode(object): data['entity_id_to_index_map'] = entity_id_to_index_map return data - def _load_ocr_info(self, data): - def trans_poly_to_bbox(poly): - x1 = np.min([p[0] for p in poly]) - x2 = np.max([p[0] for p in poly]) - y1 = np.min([p[1] for p in poly]) - y2 = np.max([p[1] for p in poly]) - return [x1, y1, x2, y2] + def trans_poly_to_bbox(self, poly): + x1 = np.min([p[0] for p in poly]) + x2 = np.max([p[0] for p in poly]) + y1 = np.min([p[1] for p in poly]) + y2 = np.max([p[1] for p in poly]) + return [x1, y1, x2, y2] + def _load_ocr_info(self, data): if self.infer_mode: ocr_result = self.ocr_engine.ocr(data['image'], cls=False) ocr_info = [] for res in ocr_result: ocr_info.append({ - "text": res[1][0], - "bbox": trans_poly_to_bbox(res[0]), - "poly": res[0], + "transcription": res[1][0], + "bbox": self.trans_poly_to_bbox(res[0]), + "points": res[0], }) return ocr_info else: info = data['label'] # read text info info_dict = json.loads(info) - return info_dict["ocr_info"] + return info_dict def _smooth_box(self, bbox, height, width): bbox[0] = int(bbox[0] * 1000.0 / width) @@ -977,7 +989,7 @@ class VQATokenLabelEncode(object): def _parse_label(self, label, encode_res): gt_label = [] - if label.lower() == "other": + if label.lower() in ["other", "others", "ignore"]: gt_label.extend([0] * len(encode_res["input_ids"])) else: gt_label.append(self.label2id_map[("b-" + label).upper()]) diff --git a/ppocr/data/imaug/vqa/__init__.py b/ppocr/data/imaug/vqa/__init__.py index a5025e7985..bde1751155 100644 --- a/ppocr/data/imaug/vqa/__init__.py +++ b/ppocr/data/imaug/vqa/__init__.py @@ -13,7 +13,12 @@ # limitations under the License. from .token import VQATokenPad, VQASerTokenChunk, VQAReTokenChunk, VQAReTokenRelation +from .augment import DistortBBox __all__ = [ - 'VQATokenPad', 'VQASerTokenChunk', 'VQAReTokenChunk', 'VQAReTokenRelation' + 'VQATokenPad', + 'VQASerTokenChunk', + 'VQAReTokenChunk', + 'VQAReTokenRelation', + 'DistortBBox', ] diff --git a/ppocr/data/imaug/vqa/augment.py b/ppocr/data/imaug/vqa/augment.py new file mode 100644 index 0000000000..fcdc9685e9 --- /dev/null +++ b/ppocr/data/imaug/vqa/augment.py @@ -0,0 +1,37 @@ +# copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +import numpy as np +import random + + +class DistortBBox: + def __init__(self, prob=0.5, max_scale=1, **kwargs): + """Random distort bbox + """ + self.prob = prob + self.max_scale = max_scale + + def __call__(self, data): + if random.random() > self.prob: + return data + bbox = np.array(data['bbox']) + rnd_scale = (np.random.rand(*bbox.shape) - 0.5) * 2 * self.max_scale + bbox = np.round(bbox + rnd_scale).astype(bbox.dtype) + data['bbox'] = np.clip(data['bbox'], 0, 1000) + data['bbox'] = bbox.tolist() + sys.stdout.flush() + return data diff --git a/ppocr/utils/utility.py b/ppocr/utils/utility.py index 4a25ff8b2f..b881fcab20 100755 --- a/ppocr/utils/utility.py +++ b/ppocr/utils/utility.py @@ -91,18 +91,19 @@ def check_and_read_gif(img_path): def load_vqa_bio_label_maps(label_map_path): with open(label_map_path, "r", encoding='utf-8') as fin: lines = fin.readlines() - lines = [line.strip() for line in lines] - if "O" not in lines: - lines.insert(0, "O") - labels = [] - for line in lines: - if line == "O": - labels.append("O") - else: - labels.append("B-" + line) - labels.append("I-" + line) - label2id_map = {label: idx for idx, label in enumerate(labels)} - id2label_map = {idx: label for idx, label in enumerate(labels)} + old_lines = [line.strip() for line in lines] + lines = ["O"] + for line in old_lines: + # "O" has already been in lines + if line.upper() in ["OTHER", "OTHERS", "IGNORE"]: + continue + lines.append(line) + labels = ["O"] + for line in lines[1:]: + labels.append("B-" + line) + labels.append("I-" + line) + label2id_map = {label.upper(): idx for idx, label in enumerate(labels)} + id2label_map = {idx: label.upper() for idx, label in enumerate(labels)} return label2id_map, id2label_map diff --git a/ppocr/utils/visual.py b/ppocr/utils/visual.py index 7a8c1674a7..235eb572a3 100644 --- a/ppocr/utils/visual.py +++ b/ppocr/utils/visual.py @@ -19,7 +19,7 @@ from PIL import Image, ImageDraw, ImageFont def draw_ser_results(image, ocr_results, font_path="doc/fonts/simfang.ttf", - font_size=18): + font_size=14): np.random.seed(2021) color = (np.random.permutation(range(255)), np.random.permutation(range(255)), @@ -40,9 +40,15 @@ def draw_ser_results(image, if ocr_info["pred_id"] not in color_map: continue color = color_map[ocr_info["pred_id"]] - text = "{}: {}".format(ocr_info["pred"], ocr_info["text"]) + text = "{}: {}".format(ocr_info["pred"], ocr_info["transcription"]) - draw_box_txt(ocr_info["bbox"], text, draw, font, font_size, color) + if "bbox" in ocr_info: + # draw with ocr engine + bbox = ocr_info["bbox"] + else: + # draw with ocr groundtruth + bbox = trans_poly_to_bbox(ocr_info["points"]) + draw_box_txt(bbox, text, draw, font, font_size, color) img_new = Image.blend(image, img_new, 0.5) return np.array(img_new) @@ -62,6 +68,14 @@ def draw_box_txt(bbox, text, draw, font, font_size, color): draw.text((bbox[0][0] + 1, start_y), text, fill=(255, 255, 255), font=font) +def trans_poly_to_bbox(poly): + x1 = np.min([p[0] for p in poly]) + x2 = np.max([p[0] for p in poly]) + y1 = np.min([p[1] for p in poly]) + y2 = np.max([p[1] for p in poly]) + return [x1, y1, x2, y2] + + def draw_re_results(image, result, font_path="doc/fonts/simfang.ttf", @@ -80,10 +94,10 @@ def draw_re_results(image, color_line = (0, 255, 0) for ocr_info_head, ocr_info_tail in result: - draw_box_txt(ocr_info_head["bbox"], ocr_info_head["text"], draw, font, - font_size, color_head) - draw_box_txt(ocr_info_tail["bbox"], ocr_info_tail["text"], draw, font, - font_size, color_tail) + draw_box_txt(ocr_info_head["bbox"], ocr_info_head["transcription"], + draw, font, font_size, color_head) + draw_box_txt(ocr_info_tail["bbox"], ocr_info_tail["transcription"], + draw, font, font_size, color_tail) center_head = ( (ocr_info_head['bbox'][0] + ocr_info_head['bbox'][2]) // 2, diff --git a/ppstructure/docs/kie.md b/ppstructure/docs/kie.md index 35498b3347..315dd9f7ba 100644 --- a/ppstructure/docs/kie.md +++ b/ppstructure/docs/kie.md @@ -16,7 +16,7 @@ SDMGR是一个关键信息提取算法，将每个检测到的文本区域分类 训练和测试的数据采用wildreceipt数据集，通过如下指令下载数据集： ``` -wget https://paddleocr.bj.bcebos.com/dygraph_v2.1/kie/wildreceipt.tar && tar xf wildreceipt.tar +wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/wildreceipt.tar && tar xf wildreceipt.tar ``` 执行预测： diff --git a/ppstructure/docs/kie_en.md b/ppstructure/docs/kie_en.md index 1fe38b0b39..7b3752223d 100644 --- a/ppstructure/docs/kie_en.md +++ b/ppstructure/docs/kie_en.md @@ -15,7 +15,7 @@ This section provides a tutorial example on how to quickly use, train, and evalu [Wildreceipt dataset](https://paperswithcode.com/dataset/wildreceipt) is used for this tutorial. It contains 1765 photos, with 25 classes, and 50000 text boxes, which can be downloaded by wget: ```shell -wget https://paddleocr.bj.bcebos.com/dygraph_v2.1/kie/wildreceipt.tar && tar xf wildreceipt.tar +wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/wildreceipt.tar && tar xf wildreceipt.tar ``` Download the pretrained model and predict the result: diff --git a/ppstructure/vqa/README.md b/ppstructure/vqa/README.md index e3a10671dd..711ffa3138 100644 --- a/ppstructure/vqa/README.md +++ b/ppstructure/vqa/README.md @@ -125,13 +125,13 @@ If you want to experience the prediction process directly, you can download the * Download the processed dataset -The download address of the processed XFUND Chinese dataset: [https://paddleocr.bj.bcebos.com/dataset/XFUND.tar](https://paddleocr.bj.bcebos.com/dataset/XFUND.tar). +The download address of the processed XFUND Chinese dataset: [link](https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar). Download and unzip the dataset, and place the dataset in the current directory after unzipping. ```shell -wget https://paddleocr.bj.bcebos.com/dataset/XFUND.tar +wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar ```` * Convert the dataset diff --git a/ppstructure/vqa/README_ch.md b/ppstructure/vqa/README_ch.md index b677dc07bc..297ba64f82 100644 --- a/ppstructure/vqa/README_ch.md +++ b/ppstructure/vqa/README_ch.md @@ -122,13 +122,13 @@ python3 -m pip install -r ppstructure/vqa/requirements.txt * 下载处理好的数据集 -处理好的XFUND中文数据集下载地址：[https://paddleocr.bj.bcebos.com/dataset/XFUND.tar](https://paddleocr.bj.bcebos.com/dataset/XFUND.tar)。 +处理好的XFUND中文数据集下载地址：[链接](https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar)。 下载并解压该数据集，解压后将数据集放置在当前目录下。 ```shell -wget https://paddleocr.bj.bcebos.com/dataset/XFUND.tar +wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar ``` * 转换数据集 diff --git a/ppstructure/vqa/labels/labels_ser.txt b/ppstructure/vqa/labels/labels_ser.txt deleted file mode 100644 index 508e481124..0000000000 --- a/ppstructure/vqa/labels/labels_ser.txt +++ /dev/null @@ -1,3 +0,0 @@ -QUESTION -ANSWER -HEADER diff --git a/ppstructure/vqa/tools/trans_xfun_data.py b/ppstructure/vqa/tools/trans_xfun_data.py index 93ec98163c..11d221bea4 100644 --- a/ppstructure/vqa/tools/trans_xfun_data.py +++ b/ppstructure/vqa/tools/trans_xfun_data.py @@ -21,26 +21,22 @@ def transfer_xfun_data(json_path=None, output_file=None): json_info = json.loads(lines[0]) documents = json_info["documents"] - label_info = {} with open(output_file, "w", encoding='utf-8') as fout: for idx, document in enumerate(documents): + label_info = [] img_info = document["img"] document = document["document"] image_path = img_info["fname"] - label_info["height"] = img_info["height"] - label_info["width"] = img_info["width"] - - label_info["ocr_info"] = [] - for doc in document: - label_info["ocr_info"].append({ - "text": doc["text"], + x1, y1, x2, y2 = doc["box"] + points = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] + label_info.append({ + "transcription": doc["text"], "label": doc["label"], - "bbox": doc["box"], + "points": points, "id": doc["id"], - "linking": doc["linking"], - "words": doc["words"] + "linking": doc["linking"] }) fout.write(image_path + "\t" + json.dumps( diff --git a/tools/infer_kie.py b/tools/infer_kie.py index 0cb0b8702c..346e2e0aee 100755 --- a/tools/infer_kie.py +++ b/tools/infer_kie.py @@ -39,13 +39,12 @@ import time def read_class_list(filepath): - dict = {} + ret = {} with open(filepath, "r") as f: lines = f.readlines() - for line in lines: - key, value = line.split(" ") - dict[key] = value.rstrip() - return dict + for idx, line in enumerate(lines): + ret[idx] = line.strip("\n") + return ret def draw_kie_result(batch, node, idx_to_cls, count): @@ -71,7 +70,7 @@ def draw_kie_result(batch, node, idx_to_cls, count): x_min = int(min([point[0] for point in new_box])) y_min = int(min([point[1] for point in new_box])) - pred_label = str(node_pred_label[i]) + pred_label = node_pred_label[i] if pred_label in idx_to_cls: pred_label = idx_to_cls[pred_label] pred_score = '{:.2f}'.format(node_pred_score[i]) @@ -109,8 +108,7 @@ def main(): save_res_path = config['Global']['save_res_path'] class_path = config['Global']['class_path'] idx_to_cls = read_class_list(class_path) - if not os.path.exists(os.path.dirname(save_res_path)): - os.makedirs(os.path.dirname(save_res_path)) + os.makedirs(os.path.dirname(save_res_path), exist_ok=True) model.eval() diff --git a/tools/infer_vqa_token_ser.py b/tools/infer_vqa_token_ser.py index 83ed72b392..39ada64a99 100755 --- a/tools/infer_vqa_token_ser.py +++ b/tools/infer_vqa_token_ser.py @@ -86,15 +86,16 @@ class SerPredictor(object): ] transforms.append(op) - global_config['infer_mode'] = True + if config["Global"].get("infer_mode", None) is None: + global_config['infer_mode'] = True self.ops = create_operators(config['Eval']['dataset']['transforms'], global_config) self.model.eval() - def __call__(self, img_path): - with open(img_path, 'rb') as f: + def __call__(self, data): + with open(data["img_path"], 'rb') as f: img = f.read() - data = {'image': img} + data["image"] = img batch = transform(data, self.ops) batch = to_tensor(batch) preds = self.model(batch) @@ -112,20 +113,35 @@ if __name__ == '__main__': ser_engine = SerPredictor(config) - infer_imgs = get_image_file_list(config['Global']['infer_img']) + if config["Global"].get("infer_mode", None) is False: + data_dir = config['Eval']['dataset']['data_dir'] + with open(config['Global']['infer_img'], "rb") as f: + infer_imgs = f.readlines() + else: + infer_imgs = get_image_file_list(config['Global']['infer_img']) + with open( os.path.join(config['Global']['save_res_path'], "infer_results.txt"), "w", encoding='utf-8') as fout: - for idx, img_path in enumerate(infer_imgs): + for idx, info in enumerate(infer_imgs): + if config["Global"].get("infer_mode", None) is False: + data_line = info.decode('utf-8') + substr = data_line.strip("\n").split("\t") + img_path = os.path.join(data_dir, substr[0]) + data = {'img_path': img_path, 'label': substr[1]} + else: + img_path = info + data = {'img_path': img_path} + save_img_path = os.path.join( config['Global']['save_res_path'], os.path.splitext(os.path.basename(img_path))[0] + "_ser.jpg") logger.info("process: [{}/{}], save result to {}".format( idx, len(infer_imgs), save_img_path)) - result, _ = ser_engine(img_path) + result, _ = ser_engine(data) result = result[0] fout.write(img_path + "\t" + json.dumps( { diff --git a/tools/program.py b/tools/program.py index aa3ba82c44..f598feb7e8 100755 --- a/tools/program.py +++ b/tools/program.py @@ -576,8 +576,8 @@ def preprocess(is_train=False): assert alg in [ 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', - 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', - 'ViTSTR', 'ABINet' + 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE', + 'SVTR', 'ViTSTR', 'ABINet' ] if use_xpu: From 872d73daf5b8bcb4fd26afb194796fa42257b473 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Thu, 30 Jun 2022 08:59:05 +0000 Subject: [PATCH 33/43] add db++ --- tools/program.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/program.py b/tools/program.py index d86be5badb..e8e72fdd2d 100755 --- a/tools/program.py +++ b/tools/program.py @@ -574,8 +574,8 @@ def preprocess(is_train=False): assert alg in [ 'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN', 'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE', - 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'PREN', 'FCE', 'SVTR', - 'ViTSTR', 'ABINet', 'DB++' + 'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE', + 'SVTR', 'ViTSTR', 'ABINet', 'DB++' ] if use_xpu: From b1ee778fb367b960c7dd13a15cf6b90518ccc3c0 Mon Sep 17 00:00:00 2001 From: andyjpaddle Date: Thu, 30 Jun 2022 09:10:23 +0000 Subject: [PATCH 34/43] rm proxy --- test_tipc/test_serving_infer_cpp.sh | 3 +-- test_tipc/test_serving_infer_python.sh | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/test_tipc/test_serving_infer_cpp.sh b/test_tipc/test_serving_infer_cpp.sh index 6e313cd949..0be6a45adf 100644 --- a/test_tipc/test_serving_infer_cpp.sh +++ b/test_tipc/test_serving_infer_cpp.sh @@ -87,8 +87,7 @@ function func_serving(){ set_image_dir=$(func_set_params "${image_dir_key}" "${image_dir_value}") python_list=(${python_list}) cd ${serving_dir_value} - unset https_proxy - unset http_proxy + # cpp serving for gpu_id in ${gpu_value[*]}; do if [ ${gpu_id} = "null" ]; then diff --git a/test_tipc/test_serving_infer_python.sh b/test_tipc/test_serving_infer_python.sh index b14517265d..4ccccc06e2 100644 --- a/test_tipc/test_serving_infer_python.sh +++ b/test_tipc/test_serving_infer_python.sh @@ -112,8 +112,7 @@ function func_serving(){ cd ${serving_dir_value} python=${python_list[0]} - unset https_proxy - unset http_proxy + # python serving for use_gpu in ${web_use_gpu_list[*]}; do if [ ${use_gpu} = "null" ]; then From d4400fbf8cd577e11371b8deeda84e9913abd9f2 Mon Sep 17 00:00:00 2001 From: andyjpaddle Date: Fri, 1 Jul 2022 02:40:42 +0000 Subject: [PATCH 35/43] update log --- test_tipc/test_ptq_inference_python.sh | 2 +- test_tipc/test_train_inference_python.sh | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/test_tipc/test_ptq_inference_python.sh b/test_tipc/test_ptq_inference_python.sh index c1aa3daa6c..288e609896 100644 --- a/test_tipc/test_ptq_inference_python.sh +++ b/test_tipc/test_ptq_inference_python.sh @@ -139,7 +139,7 @@ if [ ${MODE} = "whole_infer" ]; then save_infer_dir="${infer_model}_klquant" set_export_weight=$(func_set_params "${export_weight}" "${infer_model}") set_save_infer_key=$(func_set_params "${save_infer_key}" "${save_infer_dir}") - export_log_path="${LOG_PATH}/_export_${Count}.log" + export_log_path="${LOG_PATH}_export_${Count}.log" export_cmd="${python} ${infer_run_exports[Count]} ${set_export_weight} ${set_save_infer_key} > ${export_log_path} 2>&1 " echo ${infer_run_exports[Count]} echo $export_cmd diff --git a/test_tipc/test_train_inference_python.sh b/test_tipc/test_train_inference_python.sh index fa68cb2632..907efcec90 100644 --- a/test_tipc/test_train_inference_python.sh +++ b/test_tipc/test_train_inference_python.sh @@ -193,7 +193,7 @@ if [ ${MODE} = "whole_infer" ]; then save_infer_dir="${infer_model}" set_export_weight=$(func_set_params "${export_weight}" "${infer_model}") set_save_infer_key=$(func_set_params "${save_infer_key}" "${save_infer_dir}") - export_log_path="${LOG_PATH}/_export_${Count}.log" + export_log_path="${LOG_PATH}_export_${Count}.log" export_cmd="${python} ${infer_run_exports[Count]} ${set_export_weight} ${set_save_infer_key} > ${export_log_path} 2>&1 " echo ${infer_run_exports[Count]} echo $export_cmd @@ -295,6 +295,7 @@ else fi # run train eval $cmd + eval "cat ${save_log}/train.log >> ${save_log}.log" status_check $? "${cmd}" "${status_log}" "${model_name}" set_eval_pretrain=$(func_set_params "${pretrain_model_key}" "${save_log}/${train_model_name}") From 8d46a1fbbe33d37fc858c53afd0e9fcd9cc185fa Mon Sep 17 00:00:00 2001 From: WenmuZhou <572459439@qq.com> Date: Fri, 1 Jul 2022 08:52:08 +0000 Subject: [PATCH 36/43] add dygraph2static support of layoutlm series SER model --- configs/vqa/re/layoutlmv2.yml | 14 +- configs/vqa/re/layoutxlm.yml | 12 +- configs/vqa/ser/layoutlm.yml | 6 +- configs/vqa/ser/layoutlmv2.yml | 6 +- configs/vqa/ser/layoutxlm.yml | 6 +- ppocr/data/imaug/label_ops.py | 2 +- ppocr/losses/vqa_token_layoutlm_loss.py | 4 +- ppocr/modeling/backbones/vqa_layoutlm.py | 50 +++--- .../vqa_token_ser_layoutlm_postprocess.py | 9 +- ppstructure/docs/models_list.md | 21 +-- ppstructure/docs/models_list_en.md | 21 +-- ppstructure/utility.py | 8 + ppstructure/vqa/README.md | 53 ++++-- ppstructure/vqa/README_ch.md | 39 ++++- ppstructure/vqa/predict_vqa_token_ser.py | 161 ++++++++++++++++++ ppstructure/vqa/requirements.txt | 5 +- tools/export_model.py | 18 +- tools/infer/utility.py | 11 +- tools/infer_vqa_token_ser.py | 18 +- tools/infer_vqa_token_ser_re.py | 16 +- 20 files changed, 368 insertions(+), 112 deletions(-) create mode 100644 ppstructure/vqa/predict_vqa_token_ser.py diff --git a/configs/vqa/re/layoutlmv2.yml b/configs/vqa/re/layoutlmv2.yml index 2fa5fd1165..737dbf6b60 100644 --- a/configs/vqa/re/layoutlmv2.yml +++ b/configs/vqa/re/layoutlmv2.yml @@ -11,7 +11,7 @@ Global: save_inference_dir: use_visualdl: False seed: 2048 - infer_img: doc/vqa/input/zh_val_21.jpg + infer_img: ppstructure/docs/vqa/input/zh_val_21.jpg save_res_path: ./output/re/ Architecture: @@ -21,7 +21,7 @@ Architecture: Backbone: name: LayoutLMv2ForRe pretrained: True - checkpoints: + checkpoints: Loss: name: LossFromOutput @@ -52,7 +52,7 @@ Train: name: SimpleDataSet data_dir: train_data/XFUND/zh_train/image label_file_list: - - train_data/XFUND/zh_train/xfun_normalize_train.json + - train_data/XFUND/zh_train/train.json ratio_list: [ 1.0 ] transforms: - DecodeImage: # load image @@ -61,7 +61,7 @@ Train: - VQATokenLabelEncode: # Class handling label contains_re: True algorithm: *algorithm - class_path: &class_path ppstructure/vqa/labels/labels_ser.txt + class_path: &class_path train_data/XFUND/class_list_xfun.txt - VQATokenPad: max_seq_len: &max_seq_len 512 return_attention_mask: True @@ -77,7 +77,7 @@ Train: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'bbox', 'image', 'attention_mask', 'token_type_ids','entities', 'relations'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids','image', 'entities', 'relations'] # dataloader will return list in this order loader: shuffle: True drop_last: False @@ -90,7 +90,7 @@ Eval: name: SimpleDataSet data_dir: train_data/XFUND/zh_val/image label_file_list: - - train_data/XFUND/zh_val/xfun_normalize_val.json + - train_data/XFUND/zh_val/val.json transforms: - DecodeImage: # load image img_mode: RGB @@ -114,7 +114,7 @@ Eval: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'bbox', 'image', 'attention_mask', 'token_type_ids','entities', 'relations'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image','entities', 'relations'] # dataloader will return list in this order loader: shuffle: False drop_last: False diff --git a/configs/vqa/re/layoutxlm.yml b/configs/vqa/re/layoutxlm.yml index ff16120ac1..d8585bb725 100644 --- a/configs/vqa/re/layoutxlm.yml +++ b/configs/vqa/re/layoutxlm.yml @@ -11,7 +11,7 @@ Global: save_inference_dir: use_visualdl: False seed: 2022 - infer_img: doc/vqa/input/zh_val_21.jpg + infer_img: ppstructure/docs/vqa/input/zh_val_21.jpg save_res_path: ./output/re/ Architecture: @@ -52,7 +52,7 @@ Train: name: SimpleDataSet data_dir: train_data/XFUND/zh_train/image label_file_list: - - train_data/XFUND/zh_train/xfun_normalize_train.json + - train_data/XFUND/zh_train/train.json ratio_list: [ 1.0 ] transforms: - DecodeImage: # load image @@ -61,7 +61,7 @@ Train: - VQATokenLabelEncode: # Class handling label contains_re: True algorithm: *algorithm - class_path: &class_path ppstructure/vqa/labels/labels_ser.txt + class_path: &class_path train_data/XFUND/class_list_xfun.txt - VQATokenPad: max_seq_len: &max_seq_len 512 return_attention_mask: True @@ -77,7 +77,7 @@ Train: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'bbox', 'image', 'attention_mask', 'token_type_ids','entities', 'relations'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox','attention_mask', 'token_type_ids', 'image', 'entities', 'relations'] # dataloader will return list in this order loader: shuffle: True drop_last: False @@ -90,7 +90,7 @@ Eval: name: SimpleDataSet data_dir: train_data/XFUND/zh_val/image label_file_list: - - train_data/XFUND/zh_val/xfun_normalize_val.json + - train_data/XFUND/zh_val/val.json transforms: - DecodeImage: # load image img_mode: RGB @@ -114,7 +114,7 @@ Eval: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'bbox', 'image', 'attention_mask', 'token_type_ids','entities', 'relations'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'entities', 'relations'] # dataloader will return list in this order loader: shuffle: False drop_last: False diff --git a/configs/vqa/ser/layoutlm.yml b/configs/vqa/ser/layoutlm.yml index 47ab093e1f..53e114defd 100644 --- a/configs/vqa/ser/layoutlm.yml +++ b/configs/vqa/ser/layoutlm.yml @@ -11,7 +11,7 @@ Global: save_inference_dir: use_visualdl: False seed: 2022 - infer_img: doc/vqa/input/zh_val_0.jpg + infer_img: ppstructure/docs/vqa/input/zh_val_42.jpg save_res_path: ./output/ser/ Architecture: @@ -77,7 +77,7 @@ Train: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids','labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: True drop_last: False @@ -112,7 +112,7 @@ Eval: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: False drop_last: False diff --git a/configs/vqa/ser/layoutlmv2.yml b/configs/vqa/ser/layoutlmv2.yml index d6a9c03e5e..e48c746956 100644 --- a/configs/vqa/ser/layoutlmv2.yml +++ b/configs/vqa/ser/layoutlmv2.yml @@ -11,7 +11,7 @@ Global: save_inference_dir: use_visualdl: False seed: 2022 - infer_img: doc/vqa/input/zh_val_0.jpg + infer_img: ppstructure/docs/vqa/input/zh_val_42.jpg save_res_path: ./output/ser/ Architecture: @@ -78,7 +78,7 @@ Train: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids','labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: True drop_last: False @@ -113,7 +113,7 @@ Eval: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: False drop_last: False diff --git a/configs/vqa/ser/layoutxlm.yml b/configs/vqa/ser/layoutxlm.yml index 3686989ccf..fa9df192af 100644 --- a/configs/vqa/ser/layoutxlm.yml +++ b/configs/vqa/ser/layoutxlm.yml @@ -43,7 +43,7 @@ Optimizer: PostProcess: name: VQASerTokenLayoutLMPostProcess - class_path: &class_path ppstructure/vqa/labels/labels_ser.txt + class_path: &class_path train_data/XFUND/class_list_xfun.txt Metric: name: VQASerTokenMetric @@ -78,7 +78,7 @@ Train: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids','labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: True drop_last: False @@ -113,7 +113,7 @@ Eval: order: 'hwc' - ToCHWImage: - KeepKeys: - keep_keys: [ 'input_ids', 'labels', 'bbox', 'image', 'attention_mask', 'token_type_ids'] # dataloader will return list in this order + keep_keys: [ 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels'] # dataloader will return list in this order loader: shuffle: False drop_last: False diff --git a/ppocr/data/imaug/label_ops.py b/ppocr/data/imaug/label_ops.py index c95b326224..0723e97ae7 100644 --- a/ppocr/data/imaug/label_ops.py +++ b/ppocr/data/imaug/label_ops.py @@ -912,7 +912,7 @@ class VQATokenLabelEncode(object): label = info['label'] gt_label = self._parse_label(label, encode_res) -# construct entities for re + # construct entities for re if train_re: if gt_label[0] != self.label2id_map["O"]: entity_id_to_index_map[info["id"]] = len(entities) diff --git a/ppocr/losses/vqa_token_layoutlm_loss.py b/ppocr/losses/vqa_token_layoutlm_loss.py index 244893d97d..f9cd463473 100755 --- a/ppocr/losses/vqa_token_layoutlm_loss.py +++ b/ppocr/losses/vqa_token_layoutlm_loss.py @@ -27,8 +27,8 @@ class VQASerTokenLayoutLMLoss(nn.Layer): self.ignore_index = self.loss_class.ignore_index def forward(self, predicts, batch): - labels = batch[1] - attention_mask = batch[4] + labels = batch[5] + attention_mask = batch[2] if attention_mask is not None: active_loss = attention_mask.reshape([-1, ]) == 1 active_outputs = predicts.reshape( diff --git a/ppocr/modeling/backbones/vqa_layoutlm.py b/ppocr/modeling/backbones/vqa_layoutlm.py index ede5b7a35a..1542519651 100644 --- a/ppocr/modeling/backbones/vqa_layoutlm.py +++ b/ppocr/modeling/backbones/vqa_layoutlm.py @@ -74,9 +74,9 @@ class LayoutLMForSer(NLPBaseModel): def forward(self, x): x = self.model( input_ids=x[0], - bbox=x[2], - attention_mask=x[4], - token_type_ids=x[5], + bbox=x[1], + attention_mask=x[2], + token_type_ids=x[3], position_ids=None, output_hidden_states=False) return x @@ -96,13 +96,15 @@ class LayoutLMv2ForSer(NLPBaseModel): def forward(self, x): x = self.model( input_ids=x[0], - bbox=x[2], - image=x[3], - attention_mask=x[4], - token_type_ids=x[5], + bbox=x[1], + attention_mask=x[2], + token_type_ids=x[3], + image=x[4], position_ids=None, head_mask=None, labels=None) + if not self.training: + return x return x[0] @@ -119,14 +121,16 @@ class LayoutXLMForSer(NLPBaseModel): def forward(self, x): x = self.model( - input_ids=x[0], - bbox=x[2], - image=x[3], - attention_mask=x[4], - token_type_ids=x[5], - position_ids=None, - head_mask=None, - labels=None) + input_ids=x[0], + bbox=x[1], + attention_mask=x[2], + token_type_ids=x[3], + image=x[4], + position_ids=None, + head_mask=None, + labels=None) + if not self.training: + return x return x[0] @@ -140,12 +144,12 @@ class LayoutLMv2ForRe(NLPBaseModel): x = self.model( input_ids=x[0], bbox=x[1], - labels=None, - image=x[2], - attention_mask=x[3], - token_type_ids=x[4], + attention_mask=x[2], + token_type_ids=x[3], + image=x[4], position_ids=None, head_mask=None, + labels=None, entities=x[5], relations=x[6]) return x @@ -161,12 +165,12 @@ class LayoutXLMForRe(NLPBaseModel): x = self.model( input_ids=x[0], bbox=x[1], - labels=None, - image=x[2], - attention_mask=x[3], - token_type_ids=x[4], + attention_mask=x[2], + token_type_ids=x[3], + image=x[4], position_ids=None, head_mask=None, + labels=None, entities=x[5], relations=x[6]) return x diff --git a/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py b/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py index 782cdea6c5..90bc52733d 100644 --- a/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py +++ b/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py @@ -41,11 +41,13 @@ class VQASerTokenLayoutLMPostProcess(object): self.id2label_map_for_show[val] = key def __call__(self, preds, batch=None, *args, **kwargs): + if isinstance(preds, tuple): + preds = preds[0] if isinstance(preds, paddle.Tensor): preds = preds.numpy() if batch is not None: - return self._metric(preds, batch[1]) + return self._metric(preds, batch[5]) else: return self._infer(preds, **kwargs) @@ -63,11 +65,10 @@ class VQASerTokenLayoutLMPostProcess(object): j]]) return decode_out_list, label_decode_out_list - def _infer(self, preds, attention_masks, segment_offset_ids, ocr_infos): + def _infer(self, preds, segment_offset_ids, ocr_infos): results = [] - for pred, attention_mask, segment_offset_id, ocr_info in zip( - preds, attention_masks, segment_offset_ids, ocr_infos): + for pred, segment_offset_id, ocr_info in zip(preds, segment_offset_ids, ocr_infos): pred = np.argmax(pred, axis=1) pred = [self.id2label_map[idx] for idx in pred] diff --git a/ppstructure/docs/models_list.md b/ppstructure/docs/models_list.md index c7dab999ff..71d0566bb0 100644 --- a/ppstructure/docs/models_list.md +++ b/ppstructure/docs/models_list.md @@ -1,11 +1,12 @@ # PP-Structure 系列模型列表 -- [1. 版面分析模型](#1) -- [2. OCR和表格识别模型](#2) - - [2.1 OCR](#21) - - [2.2 表格识别模型](#22) -- [3. VQA模型](#3) -- [4. KIE模型](#4) +- [PP-Structure 系列模型列表](#pp-structure-系列模型列表) + - [1. 版面分析模型](#1-版面分析模型) + - [2. OCR和表格识别模型](#2-ocr和表格识别模型) + - [2.1 OCR](#21-ocr) + - [2.2 表格识别模型](#22-表格识别模型) + - [3. VQA模型](#3-vqa模型) + - [4. KIE模型](#4-kie模型) @@ -42,11 +43,11 @@ |模型名称|模型简介|推理模型大小|下载地址| | --- | --- | --- | --- | -|ser_LayoutXLM_xfun_zh|基于LayoutXLM在xfun中文数据集上训练的SER模型|1.4G|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar) | -|re_LayoutXLM_xfun_zh|基于LayoutXLM在xfun中文数据集上训练的RE模型|1.4G|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar) | -|ser_LayoutLMv2_xfun_zh|基于LayoutLMv2在xfun中文数据集上训练的SER模型|778M|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh.tar) | +|ser_LayoutXLM_xfun_zh|基于LayoutXLM在xfun中文数据集上训练的SER模型|1.4G|[推理模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh_infer.tar) / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar) | +|re_LayoutXLM_xfun_zh|基于LayoutXLM在xfun中文数据集上训练的RE模型|1.4G|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar) | +|ser_LayoutLMv2_xfun_zh|基于LayoutLMv2在xfun中文数据集上训练的SER模型|778M|[推理模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh_infer.tar) / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh.tar) | |re_LayoutLMv2_xfun_zh|基于LayoutLMv2在xfun中文数据集上训练的RE模型|765M|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutLMv2_xfun_zh.tar) | -|ser_LayoutLM_xfun_zh|基于LayoutLM在xfun中文数据集上训练的SER模型|430M|[推理模型 coming soon]() / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh.tar) | +|ser_LayoutLM_xfun_zh|基于LayoutLM在xfun中文数据集上训练的SER模型|430M|[推理模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh_infer.tar) / [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh.tar) | ## 4. KIE模型 diff --git a/ppstructure/docs/models_list_en.md b/ppstructure/docs/models_list_en.md index b92c10c241..b3a0755505 100644 --- a/ppstructure/docs/models_list_en.md +++ b/ppstructure/docs/models_list_en.md @@ -1,11 +1,12 @@ # PP-Structure Model list -- [1. Layout Analysis](#1) -- [2. OCR and Table Recognition](#2) - - [2.1 OCR](#21) - - [2.2 Table Recognition](#22) -- [3. VQA](#3) -- [4. KIE](#4) +- [PP-Structure Model list](#pp-structure-model-list) + - [1. Layout Analysis](#1-layout-analysis) + - [2. OCR and Table Recognition](#2-ocr-and-table-recognition) + - [2.1 OCR](#21-ocr) + - [2.2 Table Recognition](#22-table-recognition) + - [3. VQA](#3-vqa) + - [4. KIE](#4-kie) @@ -42,11 +43,11 @@ If you need to use other OCR models, you can download the model in [PP-OCR model |model| description |inference model size|download| | --- |----------------------------------------------------------------| --- | --- | -|ser_LayoutXLM_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutXLM |1.4G|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar) | -|re_LayoutXLM_xfun_zh| Re model trained on xfun Chinese dataset based on LayoutXLM |1.4G|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar) | -|ser_LayoutLMv2_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutXLMv2 |778M|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh.tar) | +|ser_LayoutXLM_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutXLM |1.4G|[inference model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh_infer.tar) / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar) | +|re_LayoutXLM_xfun_zh| Re model trained on xfun Chinese dataset based on LayoutXLM |1.4G|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar) | +|ser_LayoutLMv2_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutXLMv2 |778M|[inference model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh_infer.tar) / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLMv2_xfun_zh.tar) | |re_LayoutLMv2_xfun_zh| Re model trained on xfun Chinese dataset based on LayoutXLMv2 |765M|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutLMv2_xfun_zh.tar) | -|ser_LayoutLM_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutLM |430M|[inference model coming soon]() / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh.tar) | +|ser_LayoutLM_xfun_zh| SER model trained on xfun Chinese dataset based on LayoutLM |430M|[inference model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh_infer.tar) / [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutLM_xfun_zh.tar) | ## 4. KIE diff --git a/ppstructure/utility.py b/ppstructure/utility.py index 1ad902e7e6..7ad3426740 100644 --- a/ppstructure/utility.py +++ b/ppstructure/utility.py @@ -40,6 +40,14 @@ def init_args(): type=ast.literal_eval, default=None, help='label map according to ppstructure/layout/README_ch.md') + + # params for vqa + parser.add_argument("--vqa_algorithm", type=str, default='LayoutXLM') + parser.add_argument("--ser_model_dir", type=str) + parser.add_argument( + "--ser_dict_path", + type=str, + default="../train_data/XFUND/class_list_xfun.txt") # params for inference parser.add_argument( "--mode", diff --git a/ppstructure/vqa/README.md b/ppstructure/vqa/README.md index 711ffa3138..cfb4210050 100644 --- a/ppstructure/vqa/README.md +++ b/ppstructure/vqa/README.md @@ -1,19 +1,16 @@ English | [简体中文](README_ch.md) -- [Document Visual Question Answering (Doc-VQA)](#Document-Visual-Question-Answering) - - [1. Introduction](#1-Introduction) +- [Document Visual Question Answering](#document-visual-question-answering) + - [1 Introduction](#1-introduction) - [2. Performance](#2-performance) - - [3. Effect demo](#3-Effect-demo) + - [3. Effect demo](#3-effect-demo) - [3.1 SER](#31-ser) - [3.2 RE](#32-re) - - [4. Install](#4-Install) - - [4.1 Installation dependencies](#41-Install-dependencies) - - [4.2 Install PaddleOCR](#42-Install-PaddleOCR) - - [5. Usage](#5-Usage) - - [5.1 Data and Model Preparation](#51-Data-and-Model-Preparation) - - [5.2 SER](#52-ser) + - [4. Install](#4-install) + - [4.1 Install dependencies](#41-install-dependencies) - [5.3 RE](#53-re) - - [6. Reference](#6-Reference-Links) + - [6. Reference Links](#6-reference-links) + - [License](#license) # Document Visual Question Answering @@ -187,17 +184,17 @@ CUDA_VISIBLE_DEVICES=0 python3 tools/eval.py -c configs/vqa/ser/layoutxlm.yml -o ```` Finally, `precision`, `recall`, `hmean` and other indicators will be printed -* Use `OCR engine + SER` tandem prediction +* `OCR + SER` tandem prediction based on training engine -Use the following command to complete the series prediction of `OCR engine + SER`, taking the pretrained SER model as an example: +Use the following command to complete the series prediction of `OCR engine + SER`, taking the SER model based on LayoutXLM as an example:: ```shell -CUDA_VISIBLE_DEVICES=0 python3 tools/infer_vqa_token_ser.py -c configs/vqa/ser/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/Global.infer_img=doc/vqa/input/zh_val_42.jpg +python3.7 tools/export_model.py -c configs/vqa/ser/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ Global.save_inference_dir=output/ser/infer ```` Finally, the prediction result visualization image and the prediction result text file will be saved in the directory configured by the `config.Global.save_res_path` field. The prediction result text file is named `infer_results.txt`. -* End-to-end evaluation of `OCR engine + SER` prediction system +* End-to-end evaluation of `OCR + SER` prediction system First use the `tools/infer_vqa_token_ser.py` script to complete the prediction of the dataset, then use the following command to evaluate. @@ -205,6 +202,24 @@ First use the `tools/infer_vqa_token_ser.py` script to complete the prediction o export CUDA_VISIBLE_DEVICES=0 python3 tools/eval_with_label_end2end.py --gt_json_path XFUND/zh_val/xfun_normalize_val.json --pred_json_path output_res/infer_results.txt ```` +* export model + +Use the following command to complete the model export of the SER model, taking the SER model based on LayoutXLM as an example: + +```shell +python3.7 tools/export_model.py -c configs/vqa/ser/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ Global.save_inference_dir=output/ser/infer +``` +The converted model will be stored in the directory specified by the `Global.save_inference_dir` field. + +* `OCR + SER` tandem prediction based on prediction engine + +Use the following command to complete the tandem prediction of `OCR + SER` based on the prediction engine, taking the SER model based on LayoutXLM as an example: + +```shell +cd ppstructure +CUDA_VISIBLE_DEVICES=0 python3.7 vqa/predict_vqa_token_ser.py --vqa_algorithm=LayoutXLM --ser_model_dir=../output/ser/infer --ser_dict_path=../train_data/XFUND/class_list_xfun.txt --image_dir=docs/vqa/input/zh_val_42.jpg --output=output +``` +After the prediction is successful, the visualization images and results will be saved in the directory specified by the `output` field ### 5.3 RE @@ -247,11 +262,19 @@ Finally, `precision`, `recall`, `hmean` and other indicators will be printed Use the following command to complete the series prediction of `OCR engine + SER + RE`, taking the pretrained SER and RE models as an example: ```shell export CUDA_VISIBLE_DEVICES=0 -python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/re_LayoutXLM_xfun_zh/Global.infer_img=doc/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm. yml -o_ser Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ +python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/re_LayoutXLM_xfun_zh/Global.infer_img=ppstructure/docs/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm. yml -o_ser Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ ```` Finally, the prediction result visualization image and the prediction result text file will be saved in the directory configured by the `config.Global.save_res_path` field. The prediction result text file is named `infer_results.txt`. +* export model + +cooming soon + +* `OCR + SER + RE` tandem prediction based on prediction engine + +cooming soon + ## 6. Reference Links - LayoutXLM: Multimodal Pre-training for Multilingual Visually-rich Document Understanding, https://arxiv.org/pdf/2104.08836.pdf diff --git a/ppstructure/vqa/README_ch.md b/ppstructure/vqa/README_ch.md index 297ba64f82..52113b80b5 100644 --- a/ppstructure/vqa/README_ch.md +++ b/ppstructure/vqa/README_ch.md @@ -14,6 +14,7 @@ - [5.2 SER](#52-ser) - [5.3 RE](#53-re) - [6. 参考链接](#6-参考链接) + - [License](#license) # 文档视觉问答（DOC-VQA） @@ -183,16 +184,16 @@ CUDA_VISIBLE_DEVICES=0 python3 tools/eval.py -c configs/vqa/ser/layoutxlm.yml -o ``` 最终会打印出`precision`, `recall`, `hmean`等指标 -* 使用`OCR引擎 + SER`串联预测 +* 基于训练引擎的`OCR + SER`串联预测 -使用如下命令即可完成`OCR引擎 + SER`的串联预测, 以SER预训练模型为例: +使用如下命令即可完成基于训练引擎的`OCR + SER`的串联预测, 以基于LayoutXLM的SER模型为例: ```shell CUDA_VISIBLE_DEVICES=0 python3 tools/infer_vqa_token_ser.py -c configs/vqa/ser/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ Global.infer_img=doc/vqa/input/zh_val_42.jpg ``` 最终会在`config.Global.save_res_path`字段所配置的目录下保存预测结果可视化图像以及预测结果文本文件，预测结果文本文件名为`infer_results.txt`。 -* 对`OCR引擎 + SER`预测系统进行端到端评估 +* 对`OCR + SER`预测系统进行端到端评估 首先使用 `tools/infer_vqa_token_ser.py` 脚本完成数据集的预测，然后使用下面的命令进行评估。 @@ -200,6 +201,24 @@ CUDA_VISIBLE_DEVICES=0 python3 tools/infer_vqa_token_ser.py -c configs/vqa/ser/l export CUDA_VISIBLE_DEVICES=0 python3 tools/eval_with_label_end2end.py --gt_json_path XFUND/zh_val/xfun_normalize_val.json --pred_json_path output_res/infer_results.txt ``` +* 模型导出 + +使用如下命令即可完成SER模型的模型导出, 以基于LayoutXLM的SER模型为例: + +```shell +python3.7 tools/export_model.py -c configs/vqa/ser/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ Global.save_inference_dir=output/ser/infer +``` +转换后的模型会存放在`Global.save_inference_dir`字段指定的目录下。 + +* 基于预测引擎的`OCR + SER`串联预测 + +使用如下命令即可完成基于预测引擎的`OCR + SER`的串联预测, 以基于LayoutXLM的SER模型为例: + +```shell +cd ppstructure +CUDA_VISIBLE_DEVICES=0 python3.7 vqa/predict_vqa_token_ser.py --vqa_algorithm=LayoutXLM --ser_model_dir=../output/ser/infer --ser_dict_path=../train_data/XFUND/class_list_xfun.txt --image_dir=docs/vqa/input/zh_val_42.jpg --output=output +``` +预测成功后，可视化图片和结果会保存在`output`字段指定的目录下 ### 5.3 RE @@ -236,16 +255,24 @@ CUDA_VISIBLE_DEVICES=0 python3 tools/eval.py -c configs/vqa/re/layoutxlm.yml -o ``` 最终会打印出`precision`, `recall`, `hmean`等指标 -* 使用`OCR引擎 + SER + RE`串联预测 +* 基于训练引擎的`OCR + SER + RE`串联预测 -使用如下命令即可完成`OCR引擎 + SER + RE`的串联预测, 以预训练SER和RE模型为例： +使用如下命令即可完成基于训练引擎的`OCR + SER + RE`串联预测, 以基于LayoutXLMSER和RE模型为例： ```shell export CUDA_VISIBLE_DEVICES=0 -python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/re_LayoutXLM_xfun_zh/ Global.infer_img=doc/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm.yml -o_ser Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ +python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=pretrain/re_LayoutXLM_xfun_zh/ Global.infer_img=ppstructure/docs/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm.yml -o_ser Architecture.Backbone.checkpoints=pretrain/ser_LayoutXLM_xfun_zh/ ``` 最终会在`config.Global.save_res_path`字段所配置的目录下保存预测结果可视化图像以及预测结果文本文件，预测结果文本文件名为`infer_results.txt`。 +* 模型导出 + +cooming soon + +* 基于预测引擎的`OCR + SER + RE`串联预测 + +cooming soon + ## 6. 参考链接 - LayoutXLM: Multimodal Pre-training for Multilingual Visually-rich Document Understanding, https://arxiv.org/pdf/2104.08836.pdf diff --git a/ppstructure/vqa/predict_vqa_token_ser.py b/ppstructure/vqa/predict_vqa_token_ser.py new file mode 100644 index 0000000000..f55c875713 --- /dev/null +++ b/ppstructure/vqa/predict_vqa_token_ser.py @@ -0,0 +1,161 @@ +# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import sys + +__dir__ = os.path.dirname(os.path.abspath(__file__)) +sys.path.append(__dir__) +sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) + +os.environ["FLAGS_allocator_strategy"] = 'auto_growth' + +import cv2 +import json +import numpy as np +import time + +import tools.infer.utility as utility +from ppocr.data import create_operators, transform +from ppocr.postprocess import build_post_process +from ppocr.utils.logging import get_logger +from ppocr.utils.visual import draw_ser_results +from ppocr.utils.utility import get_image_file_list, check_and_read_gif +from ppstructure.utility import parse_args + +from paddleocr import PaddleOCR + +logger = get_logger() + + +class SerPredictor(object): + def __init__(self, args): + self.ocr_engine = PaddleOCR(use_angle_cls=False, show_log=False) + + pre_process_list = [{ + 'VQATokenLabelEncode': { + 'algorithm': args.vqa_algorithm, + 'class_path': args.ser_dict_path, + 'contains_re': False, + 'ocr_engine': self.ocr_engine + } + }, { + 'VQATokenPad':{ + 'max_seq_len':512, + 'return_attention_mask': True + } + }, { + 'VQASerTokenChunk':{ + 'max_seq_len':512, + 'return_attention_mask': True + } + }, { + 'Resize':{ + 'size' : [224, 224] + } + }, { + 'NormalizeImage': { + 'std': [58.395, 57.12, 57.375], + 'mean': [123.675, 116.28, 103.53], + 'scale': '1', + 'order': 'hwc' + } + }, { + 'ToCHWImage': None + }, { + 'KeepKeys': { + 'keep_keys': [ + 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels', + 'segment_offset_id', 'ocr_info', + 'entities' + ] + } + }] + postprocess_params = { + 'name': 'VQASerTokenLayoutLMPostProcess', + "class_path": args.ser_dict_path, + } + + self.preprocess_op = create_operators(pre_process_list, {'infer_mode':True}) + self.postprocess_op = build_post_process(postprocess_params) + self.predictor, self.input_tensor, self.output_tensors, self.config = \ + utility.create_predictor(args, 'ser', logger) + + def __call__(self, img): + ori_im = img.copy() + data = {'image': img} + data = transform(data, self.preprocess_op) + img = data[0] + if img is None: + return None, 0 + img = np.expand_dims(img, axis=0) + img = img.copy() + starttime = time.time() + + for idx in range(len(self.input_tensor)): + expand_input = np.expand_dims(data[idx], axis=0) + self.input_tensor[idx].copy_from_cpu(expand_input) + + self.predictor.run() + + outputs = [] + for output_tensor in self.output_tensors: + output = output_tensor.copy_to_cpu() + outputs.append(output) + preds = outputs[0] + + post_result = self.postprocess_op( + preds, + segment_offset_ids=[data[6]], + ocr_infos=[data[7]]) + elapse = time.time() - starttime + return post_result, elapse + + +def main(args): + image_file_list = get_image_file_list(args.image_dir) + ser_predictor = SerPredictor(args) + count = 0 + total_time = 0 + + os.makedirs(args.output, exist_ok=True) + with open( + os.path.join(args.output, 'infer.txt'), mode='w', + encoding='utf-8') as f_w: + for image_file in image_file_list: + img, flag = check_and_read_gif(image_file) + if not flag: + img = cv2.imread(image_file) + img = img[:,:,::-1] + if img is None: + logger.info("error in loading image:{}".format(image_file)) + continue + ser_res, elapse = ser_predictor(img) + ser_res = ser_res[0] + + res_str = '{}\t{}\n'.format(image_file,json.dumps({"ocr_info": ser_res,}, ensure_ascii=False)) + f_w.write(res_str) + + img_res = draw_ser_results(image_file, ser_res, font_path="../doc/fonts/simfang.ttf",) + + img_save_path = os.path.join(args.output, + os.path.basename(image_file)) + cv2.imwrite(img_save_path, img_res) + logger.info("save vis result to {}".format(img_save_path)) + if count > 0: + total_time += elapse + count += 1 + logger.info("Predict time of {}: {}".format(image_file, elapse)) + +if __name__ == "__main__": + main(parse_args()) diff --git a/ppstructure/vqa/requirements.txt b/ppstructure/vqa/requirements.txt index 0042ec0bae..fcd882274c 100644 --- a/ppstructure/vqa/requirements.txt +++ b/ppstructure/vqa/requirements.txt @@ -1,4 +1,7 @@ sentencepiece yacs seqeval -paddlenlp>=2.2.1 \ No newline at end of file +paddlenlp>=2.2.1 +pypandoc +attrdict +python_docx \ No newline at end of file diff --git a/tools/export_model.py b/tools/export_model.py index b10d41d5b2..752732ff7e 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -97,6 +97,22 @@ def export_single_model(model, shape=[None, 1, 32, 100], dtype="float32"), ] model = to_static(model, input_spec=other_shape) + elif arch_config["algorithm"] in ["LayoutLM", "LayoutLMv2", "LayoutXLM"]: + input_spec=[ + paddle.static.InputSpec( + shape=[None, 512], dtype="int64"), # input_ids + paddle.static.InputSpec( + shape=[None, 512, 4], dtype="int64"), # bbox + paddle.static.InputSpec( + shape=[None, 512], dtype="int64"), # attention_mask + paddle.static.InputSpec( + shape=[None, 512], dtype="int64"), # token_type_ids + paddle.static.InputSpec( + shape=[None, 3, 224, 224], dtype="int64"), # image + ] + if arch_config["algorithm"] == "LayoutLM": + input_spec.pop(4) + model = to_static(model, input_spec=[input_spec]) else: infer_shape = [3, -1, -1] if arch_config["model_type"] == "rec": @@ -172,7 +188,7 @@ def main(): config["Architecture"]["Head"]["out_channels"] = char_num model = build_model(config["Architecture"]) - load_model(config, model) + load_model(config, model, model_type=config['Architecture']["model_type"]) model.eval() save_path = config["Global"]["save_inference_dir"] diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 366212f228..aa5d2371b6 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -153,6 +153,8 @@ def create_predictor(args, mode, logger): model_dir = args.rec_model_dir elif mode == 'table': model_dir = args.table_model_dir + elif mode == 'ser': + model_dir = args.ser_model_dir else: model_dir = args.e2e_model_dir @@ -316,8 +318,13 @@ def create_predictor(args, mode, logger): # create predictor predictor = inference.create_predictor(config) input_names = predictor.get_input_names() - for name in input_names: - input_tensor = predictor.get_input_handle(name) + if mode in ['ser','re']: + input_tensor = [] + for name in input_names: + input_tensor.append(predictor.get_input_handle(name)) + else: + for name in input_names: + input_tensor = predictor.get_input_handle(name) output_tensors = get_output_tensors(args, mode, predictor) return predictor, input_tensor, output_tensors, config diff --git a/tools/infer_vqa_token_ser.py b/tools/infer_vqa_token_ser.py index 39ada64a99..1e661ec2c7 100755 --- a/tools/infer_vqa_token_ser.py +++ b/tools/infer_vqa_token_ser.py @@ -44,6 +44,7 @@ def to_tensor(data): from collections import defaultdict data_dict = defaultdict(list) to_tensor_idxs = [] + for idx, v in enumerate(data): if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)): if idx not in to_tensor_idxs: @@ -57,6 +58,7 @@ def to_tensor(data): class SerPredictor(object): def __init__(self, config): global_config = config['Global'] + self.algorithm = config['Architecture']["algorithm"] # build post process self.post_process_class = build_post_process(config['PostProcess'], @@ -70,7 +72,7 @@ class SerPredictor(object): from paddleocr import PaddleOCR - self.ocr_engine = PaddleOCR(use_angle_cls=False, show_log=False) + self.ocr_engine = PaddleOCR(use_angle_cls=False, show_log=False, use_gpu=global_config['use_gpu']) # create data ops transforms = [] @@ -80,8 +82,8 @@ class SerPredictor(object): op[op_name]['ocr_engine'] = self.ocr_engine elif op_name == 'KeepKeys': op[op_name]['keep_keys'] = [ - 'input_ids', 'labels', 'bbox', 'image', 'attention_mask', - 'token_type_ids', 'segment_offset_id', 'ocr_info', + 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels', + 'segment_offset_id', 'ocr_info', 'entities' ] @@ -99,9 +101,11 @@ class SerPredictor(object): batch = transform(data, self.ops) batch = to_tensor(batch) preds = self.model(batch) + if self.algorithm in ['LayoutLMv2', 'LayoutXLM']: + preds = preds[0] + post_result = self.post_process_class( preds, - attention_masks=batch[4], segment_offset_ids=batch[6], ocr_infos=batch[7]) return post_result, batch @@ -138,8 +142,6 @@ if __name__ == '__main__': save_img_path = os.path.join( config['Global']['save_res_path'], os.path.splitext(os.path.basename(img_path))[0] + "_ser.jpg") - logger.info("process: [{}/{}], save result to {}".format( - idx, len(infer_imgs), save_img_path)) result, _ = ser_engine(data) result = result[0] @@ -149,3 +151,7 @@ if __name__ == '__main__': }, ensure_ascii=False) + "\n") img_res = draw_ser_results(img_path, result) cv2.imwrite(save_img_path, img_res) + + logger.info("process: [{}/{}], save result to {}".format( + idx, len(infer_imgs), save_img_path)) + diff --git a/tools/infer_vqa_token_ser_re.py b/tools/infer_vqa_token_ser_re.py index 6210f7f3c2..d5ae634cea 100755 --- a/tools/infer_vqa_token_ser_re.py +++ b/tools/infer_vqa_token_ser_re.py @@ -38,7 +38,7 @@ from ppocr.utils.save_load import load_model from ppocr.utils.visual import draw_re_results from ppocr.utils.logging import get_logger from ppocr.utils.utility import get_image_file_list, load_vqa_bio_label_maps, print_dict -from tools.program import ArgsParser, load_config, merge_config, check_gpu +from tools.program import ArgsParser, load_config, merge_config from tools.infer_vqa_token_ser import SerPredictor @@ -107,7 +107,7 @@ def make_input(ser_inputs, ser_results): # remove ocr_info segment_offset_id and label in ser input ser_inputs.pop(7) ser_inputs.pop(6) - ser_inputs.pop(1) + ser_inputs.pop(5) return ser_inputs, entity_idx_dict_batch @@ -131,9 +131,7 @@ class SerRePredictor(object): self.model.eval() def __call__(self, img_path): - ser_results, ser_inputs = self.ser_engine(img_path) - paddle.save(ser_inputs, 'ser_inputs.npy') - paddle.save(ser_results, 'ser_results.npy') + ser_results, ser_inputs = self.ser_engine({'img_path': img_path}) re_input, entity_idx_dict_batch = make_input(ser_inputs, ser_results) preds = self.model(re_input) post_result = self.post_process_class( @@ -155,7 +153,6 @@ def preprocess(): # check if set use_gpu=True in paddlepaddle cpu version use_gpu = config['Global']['use_gpu'] - check_gpu(use_gpu) device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu' device = paddle.set_device(device) @@ -185,9 +182,7 @@ if __name__ == '__main__': for idx, img_path in enumerate(infer_imgs): save_img_path = os.path.join( config['Global']['save_res_path'], - os.path.splitext(os.path.basename(img_path))[0] + "_ser.jpg") - logger.info("process: [{}/{}], save result to {}".format( - idx, len(infer_imgs), save_img_path)) + os.path.splitext(os.path.basename(img_path))[0] + "_ser_re.jpg") result = ser_re_engine(img_path) result = result[0] @@ -197,3 +192,6 @@ if __name__ == '__main__': }, ensure_ascii=False) + "\n") img_res = draw_re_results(img_path, result) cv2.imwrite(save_img_path, img_res) + + logger.info("process: [{}/{}], save result to {}".format( + idx, len(infer_imgs), save_img_path)) \ No newline at end of file From ce21ad8357b3da70999442e9ca492ea7e7978473 Mon Sep 17 00:00:00 2001 From: WenmuZhou <572459439@qq.com> Date: Fri, 1 Jul 2022 08:55:04 +0000 Subject: [PATCH 37/43] update docs --- ppstructure/docs/models_list.md | 13 ++++++------ ppstructure/docs/models_list_en.md | 13 ++++++------ ppstructure/infer.sh | 4 ++++ ppstructure/vqa/README.md | 25 +++++++++++----------- ppstructure/vqa/README_ch.md | 33 +++++++++++++++--------------- 5 files changed, 44 insertions(+), 44 deletions(-) create mode 100644 ppstructure/infer.sh diff --git a/ppstructure/docs/models_list.md b/ppstructure/docs/models_list.md index 71d0566bb0..dabce3a514 100644 --- a/ppstructure/docs/models_list.md +++ b/ppstructure/docs/models_list.md @@ -1,12 +1,11 @@ # PP-Structure 系列模型列表 -- [PP-Structure 系列模型列表](#pp-structure-系列模型列表) - - [1. 版面分析模型](#1-版面分析模型) - - [2. OCR和表格识别模型](#2-ocr和表格识别模型) - - [2.1 OCR](#21-ocr) - - [2.2 表格识别模型](#22-表格识别模型) - - [3. VQA模型](#3-vqa模型) - - [4. KIE模型](#4-kie模型) +- [1. 版面分析模型](#1-版面分析模型) +- [2. OCR和表格识别模型](#2-ocr和表格识别模型) + - [2.1 OCR](#21-ocr) + - [2.2 表格识别模型](#22-表格识别模型) +- [3. VQA模型](#3-vqa模型) +- [4. KIE模型](#4-kie模型) diff --git a/ppstructure/docs/models_list_en.md b/ppstructure/docs/models_list_en.md index b3a0755505..e133a0bb2a 100644 --- a/ppstructure/docs/models_list_en.md +++ b/ppstructure/docs/models_list_en.md @@ -1,12 +1,11 @@ # PP-Structure Model list -- [PP-Structure Model list](#pp-structure-model-list) - - [1. Layout Analysis](#1-layout-analysis) - - [2. OCR and Table Recognition](#2-ocr-and-table-recognition) - - [2.1 OCR](#21-ocr) - - [2.2 Table Recognition](#22-table-recognition) - - [3. VQA](#3-vqa) - - [4. KIE](#4-kie) +- [1. Layout Analysis](#1-layout-analysis) +- [2. OCR and Table Recognition](#2-ocr-and-table-recognition) + - [2.1 OCR](#21-ocr) + - [2.2 Table Recognition](#22-table-recognition) +- [3. VQA](#3-vqa) +- [4. KIE](#4-kie) diff --git a/ppstructure/infer.sh b/ppstructure/infer.sh new file mode 100644 index 0000000000..a08cbadf09 --- /dev/null +++ b/ppstructure/infer.sh @@ -0,0 +1,4 @@ +python3.7 vqa/predict_vqa_token_ser.py --vqa_algorithm=LayoutXLM --ser_model_dir=../models/ser_LayoutXLM_xfun_zh/infer --ser_dict_path=../train_data/XFUND/class_list_xfun.txt --image_dir=docs/vqa/input/zh_val_42.jpg + + +python3.7 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=models/re_LayoutXLM_xfun_zh/ Global.infer_img=ppstructure/docs/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm.yml -o_ser Architecture.Backbone.checkpoints=models/ser_LayoutXLM_xfun_zh/ \ No newline at end of file diff --git a/ppstructure/vqa/README.md b/ppstructure/vqa/README.md index cfb4210050..05635265b5 100644 --- a/ppstructure/vqa/README.md +++ b/ppstructure/vqa/README.md @@ -1,16 +1,15 @@ English | [简体中文](README_ch.md) -- [Document Visual Question Answering](#document-visual-question-answering) - - [1 Introduction](#1-introduction) - - [2. Performance](#2-performance) - - [3. Effect demo](#3-effect-demo) - - [3.1 SER](#31-ser) - - [3.2 RE](#32-re) - - [4. Install](#4-install) - - [4.1 Install dependencies](#41-install-dependencies) - - [5.3 RE](#53-re) - - [6. Reference Links](#6-reference-links) - - [License](#license) +- [1 Introduction](#1-introduction) +- [2. Performance](#2-performance) +- [3. Effect demo](#3-effect-demo) + - [3.1 SER](#31-ser) + - [3.2 RE](#32-re) +- [4. Install](#4-install) + - [4.1 Install dependencies](#41-install-dependencies) + - [5.3 RE](#53-re) +- [6. Reference Links](#6-reference-links) +- [License](#license) # Document Visual Question Answering @@ -212,7 +211,7 @@ python3.7 tools/export_model.py -c configs/vqa/ser/layoutxlm.yml -o Architecture The converted model will be stored in the directory specified by the `Global.save_inference_dir` field. * `OCR + SER` tandem prediction based on prediction engine - + Use the following command to complete the tandem prediction of `OCR + SER` based on the prediction engine, taking the SER model based on LayoutXLM as an example: ```shell @@ -268,7 +267,7 @@ python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Archi Finally, the prediction result visualization image and the prediction result text file will be saved in the directory configured by the `config.Global.save_res_path` field. The prediction result text file is named `infer_results.txt`. * export model - + cooming soon * `OCR + SER + RE` tandem prediction based on prediction engine diff --git a/ppstructure/vqa/README_ch.md b/ppstructure/vqa/README_ch.md index 52113b80b5..b421a82d3a 100644 --- a/ppstructure/vqa/README_ch.md +++ b/ppstructure/vqa/README_ch.md @@ -1,20 +1,19 @@ [English](README.md) | 简体中文 -- [文档视觉问答（DOC-VQA）](#文档视觉问答doc-vqa) - - [1. 简介](#1-简介) - - [2. 性能](#2-性能) - - [3. 效果演示](#3-效果演示) - - [3.1 SER](#31-ser) - - [3.2 RE](#32-re) - - [4. 安装](#4-安装) - - [4.1 安装依赖](#41-安装依赖) - - [4.2 安装PaddleOCR（包含 PP-OCR 和 VQA）](#42-安装paddleocr包含-pp-ocr-和-vqa) - - [5. 使用](#5-使用) - - [5.1 数据和预训练模型准备](#51-数据和预训练模型准备) - - [5.2 SER](#52-ser) - - [5.3 RE](#53-re) - - [6. 参考链接](#6-参考链接) - - [License](#license) +- [1. 简介](#1-简介) +- [2. 性能](#2-性能) +- [3. 效果演示](#3-效果演示) + - [3.1 SER](#31-ser) + - [3.2 RE](#32-re) +- [4. 安装](#4-安装) + - [4.1 安装依赖](#41-安装依赖) + - [4.2 安装PaddleOCR（包含 PP-OCR 和 VQA）](#42-安装paddleocr包含-pp-ocr-和-vqa) +- [5. 使用](#5-使用) + - [5.1 数据和预训练模型准备](#51-数据和预训练模型准备) + - [5.2 SER](#52-ser) + - [5.3 RE](#53-re) +- [6. 参考链接](#6-参考链接) +- [License](#license) # 文档视觉问答（DOC-VQA） @@ -211,7 +210,7 @@ python3.7 tools/export_model.py -c configs/vqa/ser/layoutxlm.yml -o Architecture 转换后的模型会存放在`Global.save_inference_dir`字段指定的目录下。 * 基于预测引擎的`OCR + SER`串联预测 - + 使用如下命令即可完成基于预测引擎的`OCR + SER`的串联预测, 以基于LayoutXLM的SER模型为例: ```shell @@ -266,7 +265,7 @@ python3 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Archi 最终会在`config.Global.save_res_path`字段所配置的目录下保存预测结果可视化图像以及预测结果文本文件，预测结果文本文件名为`infer_results.txt`。 * 模型导出 - + cooming soon * 基于预测引擎的`OCR + SER + RE`串联预测 From 907ebdfe05e6734ff2957e576c28668841ffd272 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Fri, 1 Jul 2022 09:17:55 +0000 Subject: [PATCH 38/43] add db++ tipc --- test_tipc/prepare.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test_tipc/prepare.sh b/test_tipc/prepare.sh index 2c9bd2901b..32df8e7865 100644 --- a/test_tipc/prepare.sh +++ b/test_tipc/prepare.sh @@ -59,6 +59,9 @@ if [ ${MODE} = "lite_train_lite_infer" ];then wget -nc -P ./inference/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar --no-check-certificate cd ./inference/ && tar xf en_ppocr_mobile_v2.0_table_det_infer.tar && tar xf en_ppocr_mobile_v2.0_table_rec_infer.tar && cd ../ fi + if [[ ${model_name} =~ "det_r50_db++" ]];then + wget -nc -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/dygraph_v2.1/en_det/ResNet50_dcn_asf_synthtext_pretrained.pdparams --no-check-certificate + fi cd ./pretrain_models/ && tar xf det_mv3_db_v2.0_train.tar && cd ../ rm -rf ./train_data/icdar2015 rm -rf ./train_data/ic15_data From 68b379e5c62b22fee0ee556c5de4ab6b910664f7 Mon Sep 17 00:00:00 2001 From: wangjingyeye <1025993141@qq.com> Date: Fri, 1 Jul 2022 09:37:18 +0000 Subject: [PATCH 39/43] add db++ tipc --- .../det_r50_db++/train_infer_python.txt | 59 +++++++++++++++++++ 1 file changed, 59 insertions(+) create mode 100644 test_tipc/configs/det_r50_db++/train_infer_python.txt diff --git a/test_tipc/configs/det_r50_db++/train_infer_python.txt b/test_tipc/configs/det_r50_db++/train_infer_python.txt new file mode 100644 index 0000000000..bcf393a52b --- /dev/null +++ b/test_tipc/configs/det_r50_db++/train_infer_python.txt @@ -0,0 +1,59 @@ +===========================train_params=========================== +model_name:det_r50_db++ +python:python3.7 +gpu_list:0|0,1 +Global.use_gpu:True|True +Global.auto_cast:null +Global.epoch_num:lite_train_lite_infer=1|whole_train_whole_infer=300 +Global.save_model_dir:./output/ +Train.loader.batch_size_per_card:lite_train_lite_infer=2|whole_train_whole_infer=4 +Global.pretrained_model:null +train_model_name:latest +train_infer_img_dir:./train_data/icdar2015/text_localization/ch4_test_images/ +null:null +## +trainer:norm_train +norm_train:tools/train.py -c configs/det/det_r50_db++_ic15.yml -o Global.pretrained_model=./pretrain_models/ResNet50_dcn_asf_synthtext_pretrained +pact_train:null +fpgm_train:null +distill_train:null +null:null +null:null +## +===========================eval_params=========================== +eval:null +null:null +## +===========================infer_params=========================== +Global.save_inference_dir:./output/ +Global.checkpoints: +norm_export:tools/export_model.py -c configs/det/det_r50_db++_ic15.yml -o +quant_export:null +fpgm_export:null +distill_export:null +export1:null +export2:null +inference_dir:null +train_model:./inference/det_r50_db++_train/best_accuracy +infer_export:tools/export_model.py -c configs/det/det_r50_db++_ic15.yml -o +infer_quant:False +inference:tools/infer/predict_det.py --det_algorithm="DB++" +--use_gpu:True|False +--enable_mkldnn:False +--cpu_threads:6 +--rec_batch_num:1 +--use_tensorrt:False +--precision:fp32 +--det_model_dir: +--image_dir:./inference/ch_det_data_50/all-sum-510/ +null:null +--benchmark:True +null:null +===========================infer_benchmark_params========================== +random_infer_input:[{float32,[3,640,640]}];[{float32,[3,960,960]}] +===========================train_benchmark_params========================== +batch_size:8|16 +fp_items:fp32|fp16 +epoch:2 +--profiler_options:batch_range=[10,20];state=GPU;tracer_option=Default;profile_path=model.profile +flags:FLAGS_eager_delete_tensor_gb=0.0;FLAGS_fraction_of_gpu_memory_to_use=0.98;FLAGS_conv_workspace_size_limit=4096 From dc7bfe8a8441167cf6303879800eee7e160e2679 Mon Sep 17 00:00:00 2001 From: WenmuZhou <572459439@qq.com> Date: Fri, 1 Jul 2022 09:42:27 +0000 Subject: [PATCH 40/43] fix --- ppstructure/infer.sh | 4 --- ppstructure/vqa/predict_vqa_token_ser.py | 42 ++++++++++++++---------- 2 files changed, 25 insertions(+), 21 deletions(-) delete mode 100644 ppstructure/infer.sh diff --git a/ppstructure/infer.sh b/ppstructure/infer.sh deleted file mode 100644 index a08cbadf09..0000000000 --- a/ppstructure/infer.sh +++ /dev/null @@ -1,4 +0,0 @@ -python3.7 vqa/predict_vqa_token_ser.py --vqa_algorithm=LayoutXLM --ser_model_dir=../models/ser_LayoutXLM_xfun_zh/infer --ser_dict_path=../train_data/XFUND/class_list_xfun.txt --image_dir=docs/vqa/input/zh_val_42.jpg - - -python3.7 tools/infer_vqa_token_ser_re.py -c configs/vqa/re/layoutxlm.yml -o Architecture.Backbone.checkpoints=models/re_LayoutXLM_xfun_zh/ Global.infer_img=ppstructure/docs/vqa/input/zh_val_21.jpg -c_ser configs/vqa/ser/layoutxlm.yml -o_ser Architecture.Backbone.checkpoints=models/ser_LayoutXLM_xfun_zh/ \ No newline at end of file diff --git a/ppstructure/vqa/predict_vqa_token_ser.py b/ppstructure/vqa/predict_vqa_token_ser.py index f55c875713..de0bbfe72d 100644 --- a/ppstructure/vqa/predict_vqa_token_ser.py +++ b/ppstructure/vqa/predict_vqa_token_ser.py @@ -16,7 +16,7 @@ import sys __dir__ = os.path.dirname(os.path.abspath(__file__)) sys.path.append(__dir__) -sys.path.append(os.path.abspath(os.path.join(__dir__, '../..'))) +sys.path.insert(0, os.path.abspath(os.path.join(__dir__, '../..'))) os.environ["FLAGS_allocator_strategy"] = 'auto_growth' @@ -50,18 +50,18 @@ class SerPredictor(object): 'ocr_engine': self.ocr_engine } }, { - 'VQATokenPad':{ - 'max_seq_len':512, + 'VQATokenPad': { + 'max_seq_len': 512, 'return_attention_mask': True } }, { - 'VQASerTokenChunk':{ - 'max_seq_len':512, + 'VQASerTokenChunk': { + 'max_seq_len': 512, 'return_attention_mask': True } }, { - 'Resize':{ - 'size' : [224, 224] + 'Resize': { + 'size': [224, 224] } }, { 'NormalizeImage': { @@ -75,8 +75,8 @@ class SerPredictor(object): }, { 'KeepKeys': { 'keep_keys': [ - 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels', - 'segment_offset_id', 'ocr_info', + 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', + 'image', 'labels', 'segment_offset_id', 'ocr_info', 'entities' ] } @@ -86,7 +86,8 @@ class SerPredictor(object): "class_path": args.ser_dict_path, } - self.preprocess_op = create_operators(pre_process_list, {'infer_mode':True}) + self.preprocess_op = create_operators(pre_process_list, + {'infer_mode': True}) self.postprocess_op = build_post_process(postprocess_params) self.predictor, self.input_tensor, self.output_tensors, self.config = \ utility.create_predictor(args, 'ser', logger) @@ -113,11 +114,9 @@ class SerPredictor(object): output = output_tensor.copy_to_cpu() outputs.append(output) preds = outputs[0] - + post_result = self.postprocess_op( - preds, - segment_offset_ids=[data[6]], - ocr_infos=[data[7]]) + preds, segment_offset_ids=[data[6]], ocr_infos=[data[7]]) elapse = time.time() - starttime return post_result, elapse @@ -136,17 +135,25 @@ def main(args): img, flag = check_and_read_gif(image_file) if not flag: img = cv2.imread(image_file) - img = img[:,:,::-1] + img = img[:, :, ::-1] if img is None: logger.info("error in loading image:{}".format(image_file)) continue ser_res, elapse = ser_predictor(img) ser_res = ser_res[0] - res_str = '{}\t{}\n'.format(image_file,json.dumps({"ocr_info": ser_res,}, ensure_ascii=False)) + res_str = '{}\t{}\n'.format( + image_file, + json.dumps( + { + "ocr_info": ser_res, + }, ensure_ascii=False)) f_w.write(res_str) - img_res = draw_ser_results(image_file, ser_res, font_path="../doc/fonts/simfang.ttf",) + img_res = draw_ser_results( + image_file, + ser_res, + font_path="../doc/fonts/simfang.ttf", ) img_save_path = os.path.join(args.output, os.path.basename(image_file)) @@ -157,5 +164,6 @@ def main(args): count += 1 logger.info("Predict time of {}: {}".format(image_file, elapse)) + if __name__ == "__main__": main(parse_args()) From 807dd106361ce97c0f0de59fda6ede5d15413c9d Mon Sep 17 00:00:00 2001 From: WenmuZhou <572459439@qq.com> Date: Fri, 1 Jul 2022 09:59:11 +0000 Subject: [PATCH 41/43] pre-commit --- ppocr/modeling/backbones/vqa_layoutlm.py | 16 ++++++++-------- .../vqa_token_ser_layoutlm_postprocess.py | 3 ++- ppstructure/utility.py | 3 +-- tools/export_model.py | 3 ++- tools/infer/utility.py | 2 +- tools/infer_vqa_token_ser.py | 18 +++++++++--------- tools/infer_vqa_token_ser_re.py | 4 ++-- 7 files changed, 25 insertions(+), 24 deletions(-) diff --git a/ppocr/modeling/backbones/vqa_layoutlm.py b/ppocr/modeling/backbones/vqa_layoutlm.py index 1542519651..2fd1b1b2a7 100644 --- a/ppocr/modeling/backbones/vqa_layoutlm.py +++ b/ppocr/modeling/backbones/vqa_layoutlm.py @@ -121,14 +121,14 @@ class LayoutXLMForSer(NLPBaseModel): def forward(self, x): x = self.model( - input_ids=x[0], - bbox=x[1], - attention_mask=x[2], - token_type_ids=x[3], - image=x[4], - position_ids=None, - head_mask=None, - labels=None) + input_ids=x[0], + bbox=x[1], + attention_mask=x[2], + token_type_ids=x[3], + image=x[4], + position_ids=None, + head_mask=None, + labels=None) if not self.training: return x return x[0] diff --git a/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py b/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py index 90bc52733d..8a6669f71f 100644 --- a/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py +++ b/ppocr/postprocess/vqa_token_ser_layoutlm_postprocess.py @@ -68,7 +68,8 @@ class VQASerTokenLayoutLMPostProcess(object): def _infer(self, preds, segment_offset_ids, ocr_infos): results = [] - for pred, segment_offset_id, ocr_info in zip(preds, segment_offset_ids, ocr_infos): + for pred, segment_offset_id, ocr_info in zip(preds, segment_offset_ids, + ocr_infos): pred = np.argmax(pred, axis=1) pred = [self.id2label_map[idx] for idx in pred] diff --git a/ppstructure/utility.py b/ppstructure/utility.py index 7ad3426740..4ae56099b8 100644 --- a/ppstructure/utility.py +++ b/ppstructure/utility.py @@ -40,7 +40,6 @@ def init_args(): type=ast.literal_eval, default=None, help='label map according to ppstructure/layout/README_ch.md') - # params for vqa parser.add_argument("--vqa_algorithm", type=str, default='LayoutXLM') parser.add_argument("--ser_model_dir", type=str) @@ -73,7 +72,7 @@ def init_args(): "--recovery", type=bool, default=False, - help='Whether to enable layout of recovery') + help='Whether to enable layout of recovery') return parser diff --git a/tools/export_model.py b/tools/export_model.py index 752732ff7e..503951a8a6 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -97,8 +97,9 @@ def export_single_model(model, shape=[None, 1, 32, 100], dtype="float32"), ] model = to_static(model, input_spec=other_shape) + elif arch_config["algorithm"] in ["LayoutLM", "LayoutLMv2", "LayoutXLM"]: - input_spec=[ + input_spec = [ paddle.static.InputSpec( shape=[None, 512], dtype="int64"), # input_ids paddle.static.InputSpec( diff --git a/tools/infer/utility.py b/tools/infer/utility.py index aa5d2371b6..7eb77dec74 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -318,7 +318,7 @@ def create_predictor(args, mode, logger): # create predictor predictor = inference.create_predictor(config) input_names = predictor.get_input_names() - if mode in ['ser','re']: + if mode in ['ser', 're']: input_tensor = [] for name in input_names: input_tensor.append(predictor.get_input_handle(name)) diff --git a/tools/infer_vqa_token_ser.py b/tools/infer_vqa_token_ser.py index 1e661ec2c7..0173a554ca 100755 --- a/tools/infer_vqa_token_ser.py +++ b/tools/infer_vqa_token_ser.py @@ -44,7 +44,7 @@ def to_tensor(data): from collections import defaultdict data_dict = defaultdict(list) to_tensor_idxs = [] - + for idx, v in enumerate(data): if isinstance(v, (np.ndarray, paddle.Tensor, numbers.Number)): if idx not in to_tensor_idxs: @@ -72,7 +72,10 @@ class SerPredictor(object): from paddleocr import PaddleOCR - self.ocr_engine = PaddleOCR(use_angle_cls=False, show_log=False, use_gpu=global_config['use_gpu']) + self.ocr_engine = PaddleOCR( + use_angle_cls=False, + show_log=False, + use_gpu=global_config['use_gpu']) # create data ops transforms = [] @@ -82,8 +85,8 @@ class SerPredictor(object): op[op_name]['ocr_engine'] = self.ocr_engine elif op_name == 'KeepKeys': op[op_name]['keep_keys'] = [ - 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', 'image', 'labels', - 'segment_offset_id', 'ocr_info', + 'input_ids', 'bbox', 'attention_mask', 'token_type_ids', + 'image', 'labels', 'segment_offset_id', 'ocr_info', 'entities' ] @@ -103,11 +106,9 @@ class SerPredictor(object): preds = self.model(batch) if self.algorithm in ['LayoutLMv2', 'LayoutXLM']: preds = preds[0] - + post_result = self.post_process_class( - preds, - segment_offset_ids=batch[6], - ocr_infos=batch[7]) + preds, segment_offset_ids=batch[6], ocr_infos=batch[7]) return post_result, batch @@ -154,4 +155,3 @@ if __name__ == '__main__': logger.info("process: [{}/{}], save result to {}".format( idx, len(infer_imgs), save_img_path)) - diff --git a/tools/infer_vqa_token_ser_re.py b/tools/infer_vqa_token_ser_re.py index d5ae634cea..20ab1fe176 100755 --- a/tools/infer_vqa_token_ser_re.py +++ b/tools/infer_vqa_token_ser_re.py @@ -192,6 +192,6 @@ if __name__ == '__main__': }, ensure_ascii=False) + "\n") img_res = draw_re_results(img_path, result) cv2.imwrite(save_img_path, img_res) - + logger.info("process: [{}/{}], save result to {}".format( - idx, len(infer_imgs), save_img_path)) \ No newline at end of file + idx, len(infer_imgs), save_img_path)) From 8d84144029ac920f3f9b037bec10fc2646031b13 Mon Sep 17 00:00:00 2001 From: WenmuZhou <572459439@qq.com> Date: Fri, 1 Jul 2022 10:00:10 +0000 Subject: [PATCH 42/43] pre-commit --- tools/export_model.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/export_model.py b/tools/export_model.py index 503951a8a6..65573cf46a 100755 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -97,7 +97,6 @@ def export_single_model(model, shape=[None, 1, 32, 100], dtype="float32"), ] model = to_static(model, input_spec=other_shape) - elif arch_config["algorithm"] in ["LayoutLM", "LayoutLMv2", "LayoutXLM"]: input_spec = [ paddle.static.InputSpec( From a3c96e6fadcc45ba5a3a9e31579c15300ff55c56 Mon Sep 17 00:00:00 2001 From: andyjpaddle Date: Mon, 4 Jul 2022 07:08:45 +0000 Subject: [PATCH 43/43] fix east r50 hang --- test_tipc/configs/det_r50_vd_east_v2_0/train_infer_python.txt | 4 ++-- test_tipc/prepare.sh | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/test_tipc/configs/det_r50_vd_east_v2_0/train_infer_python.txt b/test_tipc/configs/det_r50_vd_east_v2_0/train_infer_python.txt index 8477a4fa74..24e4d760c3 100644 --- a/test_tipc/configs/det_r50_vd_east_v2_0/train_infer_python.txt +++ b/test_tipc/configs/det_r50_vd_east_v2_0/train_infer_python.txt @@ -1,13 +1,13 @@ ===========================train_params=========================== model_name:det_r50_vd_east_v2_0 python:python3.7 -gpu_list:0 +gpu_list:0|0,1 Global.use_gpu:True|True Global.auto_cast:fp32 Global.epoch_num:lite_train_lite_infer=1|whole_train_whole_infer=500 Global.save_model_dir:./output/ Train.loader.batch_size_per_card:lite_train_lite_infer=2|whole_train_whole_infer=4 -Global.pretrained_model:null +Global.pretrained_model:./pretrain_models/det_r50_vd_east_v2.0_train/best_accuracy train_model_name:latest train_infer_img_dir:./train_data/icdar2015/text_localization/ch4_test_images/ null:null diff --git a/test_tipc/prepare.sh b/test_tipc/prepare.sh index 2c9bd2901b..073e1c6d1b 100644 --- a/test_tipc/prepare.sh +++ b/test_tipc/prepare.sh @@ -120,6 +120,10 @@ if [ ${MODE} = "lite_train_lite_infer" ];then wget -nc -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/det_mv3_east_v2.0_train.tar --no-check-certificate cd ./pretrain_models/ && tar xf det_mv3_east_v2.0_train.tar && cd ../ fi + if [ ${model_name} == "det_r50_vd_east_v2_0" ]; then + wget -nc -P ./pretrain_models/ https://paddleocr.bj.bcebos.com/dygraph_v2.0/en/det_r50_vd_east_v2.0_train.tar --no-check-certificate + cd ./pretrain_models/ && tar xf det_r50_vd_east_v2.0_train.tar && cd ../ + fi elif [ ${MODE} = "whole_train_whole_infer" ];then wget -nc -P ./pretrain_models/ https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileNetV3_large_x0_5_pretrained.pdparams --no-check-certificate