This commit is contained in:
3
ppstructure/README.md
Normal file
3
ppstructure/README.md
Normal file
@@ -0,0 +1,3 @@
|
||||
See [Docs](https://paddlepaddle.github.io/PaddleOCR/latest/en/ppstructure/overview.html) for details.
|
||||
|
||||
请移步[Docs](https://paddlepaddle.github.io/PaddleOCR/latest/ppstructure/overview.html)查看。
|
||||
13
ppstructure/__init__.py
Normal file
13
ppstructure/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
289
ppstructure/kie/README.md
Normal file
289
ppstructure/kie/README.md
Normal file
@@ -0,0 +1,289 @@
|
||||
English | [简体中文](README_ch.md)
|
||||
|
||||
# Key Information Extraction (KIE)
|
||||
|
||||
- [1. Introduction](#1-introduction)
|
||||
- [2. Performance](#2-performance)
|
||||
- [3. Visualization](#3-visualization)
|
||||
- [3.1 SER](#31-ser)
|
||||
- [3.2 RE](#32-re)
|
||||
- [4. Usage](#4-usage)
|
||||
- [4.1 Prepare for the environment](#41-prepare-for-the-environment)
|
||||
- [4.2 Quick start](#42-quick-start)
|
||||
- [4.3 More](#43-more)
|
||||
- [5. Reference](#5-reference)
|
||||
- [6. License](#6-license)
|
||||
|
||||
|
||||
## 1. Introduction
|
||||
|
||||
Key information extraction (KIE) refers to extracting key information from text or images. As downstream task of OCR, the key information extraction task of document image has many practical application scenarios, such as form recognition, ticket information extraction, ID card information extraction, etc.
|
||||
|
||||
PP-Structure conducts research based on the LayoutXLM multi-modal, and proposes the VI-LayoutXLM, which gets rid of visual features when finetuning the downstream tasks. An textline sorting method is also utilized to fit in reading order. What's more, UDML knowledge distillation is used for higher accuracy. Finally, the accuracy and inference speed of VI-LayoutXLM surpass those of LayoutXLM.
|
||||
|
||||
The main features of the key information extraction module in PP-Structure are as follows.
|
||||
|
||||
|
||||
- Integrate multi-modal methods such as [LayoutXLM](https://arxiv.org/pdf/2104.08836.pdf), VI-LayoutXLM, and PP-OCR inference engine.
|
||||
- Supports Semantic Entity Recognition (SER) and Relation Extraction (RE) tasks based on multimodal methods. Based on the SER task, the text recognition and classification in the image can be completed; based on the RE task, the relationship extraction of the text content in the image can be completed, such as judging the problem pair (pair).
|
||||
- Supports custom training for SER tasks and RE tasks.
|
||||
- Supports end-to-end system prediction and evaluation of OCR+SER.
|
||||
- Supports end-to-end system prediction of OCR+SER+RE.
|
||||
- Support SER model export and inference using PaddleInference.
|
||||
|
||||
|
||||
## 2. Performance
|
||||
|
||||
We evaluate the methods on the Chinese dataset of [XFUND](https://github.com/doc-analysis/XFUND), and the performance is as follows
|
||||
|
||||
|Model | Backbone | Task | Config file | Hmean | Inference time (ms) | Download link|
|
||||
| --- | --- | --- | --- | --- | --- | --- |
|
||||
|VI-LayoutXLM| VI-LayoutXLM-base | SER | [ser_vi_layoutxlm_xfund_zh_udml.yml](../../configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml)|**93.19%**| 15.49|[trained model](https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_pretrained.tar)|
|
||||
|LayoutXLM| LayoutXLM-base | SER | [ser_layoutxlm_xfund_zh.yml](../../configs/kie/layoutlm_series/ser_layoutxlm_xfund_zh.yml)|90.38%| 19.49 | [trained model](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar)|
|
||||
|VI-LayoutXLM| VI-LayoutXLM-base | RE | [re_vi_layoutxlm_xfund_zh_udml.yml](../../configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh_udml.yml)|**83.92%**| 15.49|[trained model](https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_pretrained.tar)|
|
||||
|LayoutXLM| LayoutXLM-base | RE | [re_layoutxlm_xfund_zh.yml](../../configs/kie/layoutlm_series/re_layoutxlm_xfund_zh.yml)|74.83%| 19.49|[trained model](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar)|
|
||||
|
||||
|
||||
* Note：Inference environment：V100 GPU + cuda10.2 + cudnn8.1.1 + TensorRT 7.2.3.4，tested using fp16.
|
||||
|
||||
For more KIE models in PaddleOCR, please refer to [KIE model zoo](../../doc/doc_en/algorithm_overview_en.md).
|
||||
|
||||
|
||||
## 3. Visualization
|
||||
|
||||
There are two main solutions to the key information extraction task based on VI-LayoutXLM series model.
|
||||
|
||||
(1) Text detection + text recognition + semantic entity recognition (SER)
|
||||
|
||||
(2) Text detection + text recognition + semantic entity recognition (SER) + relationship extraction (RE)
|
||||
|
||||
|
||||
The following images are demo results of the SER and RE models. For more detailed introduction to the above solutions, please refer to [KIE Guide](./how_to_do_kie.md).
|
||||
|
||||
### 3.1 SER
|
||||
|
||||
Demo results for SER task are as follows.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539141-68e71c75-5cf7-4529-b2ca-219d29fa5f68.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185310636-6ce02f7c-790d-479f-b163-ea97a5a04808.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539517-ccf2372a-f026-4a7c-ad28-c741c770f60a.png" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539735-37b5c2ef-629d-43fe-9abb-44bb717ef7ee.jpg" width="600">
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
**Note:** test pictures are from [xfund dataset](https://github.com/doc-analysis/XFUND), [invoice dataset](https://aistudio.baidu.com/aistudio/datasetdetail/165561) and a composite ID card dataset.
|
||||
|
||||
|
||||
Boxes of different colors in the image represent different categories.
|
||||
|
||||
The invoice and application form images have three categories: `request`, `answer` and `header`. The `question` and `answer` can be used to extract the relationship.
|
||||
|
||||
For the ID card image, the model can directly identify the key information such as `name`, `gender`, `nationality`, so that the subsequent relationship extraction process is not required, and the key information extraction task can be completed using only one model.
|
||||
|
||||
### 3.2 RE
|
||||
|
||||
Demo results for RE task are as follows.
|
||||
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185393805-c67ff571-cf7e-4217-a4b0-8b396c4f22bb.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185540080-0431e006-9235-4b6d-b63d-0b3c6e1de48f.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185540291-f64e5daf-6d42-4e7c-bbbb-471e3fac4fcc.png" width="600">
|
||||
</div>
|
||||
|
||||
Red boxes are questions, blue boxes are answers. The green lines means the two connected objects are a pair.
|
||||
|
||||
|
||||
## 4. Usage
|
||||
|
||||
### 4.1 Prepare for the environment
|
||||
|
||||
|
||||
Use the following command to install KIE dependencies.
|
||||
|
||||
|
||||
```bash
|
||||
git clone https://github.com/PaddlePaddle/PaddleOCR.git
|
||||
cd PaddleOCR
|
||||
pip install -r requirements.txt
|
||||
pip install -r ppstructure/kie/requirements.txt
|
||||
# 安装PaddleOCR引擎用于预测
|
||||
pip install paddleocr -U
|
||||
```
|
||||
|
||||
NOTE: For KIE tasks, it is necessary to downgrade the Paddle framework version (Paddle<2.6) and the PaddleNLP version (PaddleNLP<2.6).
|
||||
|
||||
The visualized results of SER are saved in the `./output` folder by default. Examples of results are as follows.
|
||||
|
||||
|
||||
<div align="center">
|
||||
<img src="../../ppstructure/docs/kie/result_ser/zh_val_42_ser.jpg" width="800">
|
||||
</div>
|
||||
|
||||
|
||||
### 4.2 Quick start
|
||||
|
||||
Here we use XFUND dataset to quickly experience the SER model and RE model.
|
||||
|
||||
|
||||
#### 4.2.1 Prepare for the dataset
|
||||
|
||||
```bash
|
||||
mkdir train_data
|
||||
cd train_data
|
||||
# download and uncompress the dataset
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar && tar -xf XFUND.tar
|
||||
cd ..
|
||||
```
|
||||
|
||||
#### 4.2.2 Predict images using the trained model
|
||||
|
||||
Use the following command to download the models.
|
||||
|
||||
```bash
|
||||
mkdir pretrained_model
|
||||
cd pretrained_model
|
||||
# download and uncompress the SER trained model
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_pretrained.tar && tar -xf ser_vi_layoutxlm_xfund_pretrained.tar
|
||||
|
||||
# download and uncompress the RE trained model
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_pretrained.tar && tar -xf re_vi_layoutxlm_xfund_pretrained.tar
|
||||
```
|
||||
|
||||
|
||||
If you want to use OCR engine to obtain end-to-end prediction results, you can use the following command to predict.
|
||||
|
||||
```bash
|
||||
# just predict using SER trained model
|
||||
python3 tools/infer_kie_token_ser.py \
|
||||
-c configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./ppstructure/docs/kie/input/zh_val_42.jpg
|
||||
|
||||
# predict using SER and RE trained model at the same time
|
||||
python3 ./tools/infer_kie_token_ser_re.py \
|
||||
-c configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/re_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/image/zh_val_42.jpg \
|
||||
-c_ser configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o_ser Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy
|
||||
```
|
||||
|
||||
The visual result images and the predicted text file will be saved in the `Global.save_res_path` directory.
|
||||
|
||||
If you want to use a custom ocr model, you can set it through the following fields
|
||||
- `Global.kie_det_model_dir`: the detection inference model path
|
||||
- `Global.kie_rec_model_dir`: the recognition inference model path
|
||||
|
||||
|
||||
If you want to load the text detection and recognition results collected before, you can use the following command to predict.
|
||||
|
||||
```bash
|
||||
# just predict using SER trained model
|
||||
python3 tools/infer_kie_token_ser.py \
|
||||
-c configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/val.json \
|
||||
Global.infer_mode=False
|
||||
|
||||
# predict using SER and RE trained model at the same time
|
||||
python3 ./tools/infer_kie_token_ser_re.py \
|
||||
-c configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/re_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/val.json \
|
||||
Global.infer_mode=False \
|
||||
-c_ser configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o_ser Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy
|
||||
```
|
||||
|
||||
#### 4.2.3 Inference using PaddleInference
|
||||
|
||||
Firstly, download the inference SER inference model.
|
||||
|
||||
```bash
|
||||
mkdir inference
|
||||
cd inference
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_infer.tar && tar -xf ser_vi_layoutxlm_xfund_infer.tar
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_infer.tar && tar -xf re_vi_layoutxlm_xfund_infer.tar
|
||||
cd ..
|
||||
```
|
||||
|
||||
- SER
|
||||
|
||||
Use the following command for inference.
|
||||
|
||||
|
||||
```bash
|
||||
cd ppstructure
|
||||
python3 kie/predict_kie_token_ser.py \
|
||||
--kie_algorithm=LayoutXLM \
|
||||
--ser_model_dir=../inference/ser_vi_layoutxlm_xfund_infer \
|
||||
--image_dir=./docs/kie/input/zh_val_42.jpg \
|
||||
--ser_dict_path=../train_data/XFUND/class_list_xfun.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--ocr_order_method="tb-yx"
|
||||
```
|
||||
|
||||
The visual results and text file will be saved in directory `output`.
|
||||
|
||||
- RE
|
||||
|
||||
Use the following command for inference.
|
||||
|
||||
|
||||
```bash
|
||||
cd ppstructure
|
||||
python3 kie/predict_kie_token_ser_re.py \
|
||||
--kie_algorithm=LayoutXLM \
|
||||
--re_model_dir=../inference/re_vi_layoutxlm_xfund_infer \
|
||||
--ser_model_dir=../inference/ser_vi_layoutxlm_xfund_infer \
|
||||
--use_visual_backbone=False \
|
||||
--image_dir=./docs/kie/input/zh_val_42.jpg \
|
||||
--ser_dict_path=../train_data/XFUND/class_list_xfun.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--ocr_order_method="tb-yx"
|
||||
```
|
||||
|
||||
The visual results and text file will be saved in directory `output`.
|
||||
|
||||
If you want to use a custom ocr model, you can set it through the following fields
|
||||
- `--det_model_dir`: the detection inference model path
|
||||
- `--rec_model_dir`: the recognition inference model path
|
||||
|
||||
### 4.3 More
|
||||
|
||||
For training, evaluation and inference tutorial for KIE models, please refer to [KIE doc](../../doc/doc_en/kie_en.md).
|
||||
|
||||
For training, evaluation and inference tutorial for text detection models, please refer to [text detection doc](../../doc/doc_en/detection_en.md).
|
||||
|
||||
For training, evaluation and inference tutorial for text recognition models, please refer to [text recognition doc](../../doc/doc_en/recognition_en.md).
|
||||
|
||||
To complete the key information extraction task in your own scenario from data preparation to model selection, please refer to: [Guide to End-to-end KIE](./how_to_do_kie_en.md)。
|
||||
|
||||
|
||||
## 5. Reference
|
||||
|
||||
- LayoutXLM: Multimodal Pre-training for Multilingual Visually-rich Document Understanding, https://arxiv.org/pdf/2104.08836.pdf
|
||||
- microsoft/unilm/layoutxlm, https://github.com/microsoft/unilm/tree/master/layoutxlm
|
||||
- XFUND dataset, https://github.com/doc-analysis/XFUND
|
||||
|
||||
## 6. License
|
||||
|
||||
The content of this project itself is licensed under the [Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/)
|
||||
271
ppstructure/kie/README_ch.md
Normal file
271
ppstructure/kie/README_ch.md
Normal file
@@ -0,0 +1,271 @@
|
||||
[English](README.md) | 简体中文
|
||||
|
||||
# 关键信息抽取
|
||||
|
||||
- [1. 简介](#1-简介)
|
||||
- [2. 精度与性能](#2-精度与性能)
|
||||
- [3. 效果演示](#3-效果演示)
|
||||
- [3.1 SER](#31-ser)
|
||||
- [3.2 RE](#32-re)
|
||||
- [4. 使用](#4-使用)
|
||||
- [4.1 准备环境](#41-准备环境)
|
||||
- [4.2 快速开始](#42-快速开始)
|
||||
- [4.3 更多](#43-更多)
|
||||
- [5. 参考链接](#5-参考链接)
|
||||
- [6. License](#6-License)
|
||||
|
||||
|
||||
## 1. 简介
|
||||
|
||||
关键信息抽取 (Key Information Extraction, KIE)指的是是从文本或者图像中，抽取出关键的信息。针对文档图像的关键信息抽取任务作为OCR的下游任务，存在非常多的实际应用场景，如表单识别、车票信息抽取、身份证信息抽取等。
|
||||
|
||||
PP-Structure 基于 LayoutXLM 文档多模态系列方法进行研究与优化，设计了视觉特征无关的多模态模型结构VI-LayoutXLM，同时引入符合阅读顺序的文本行排序方法以及UDML联合互学习蒸馏方法，最终在精度与速度均超越LayoutXLM。
|
||||
|
||||
PP-Structure中关键信息抽取模块的主要特性如下：
|
||||
|
||||
- 集成[LayoutXLM](https://arxiv.org/pdf/2104.08836.pdf)、VI-LayoutXLM等多模态模型以及PP-OCR预测引擎。
|
||||
- 支持基于多模态方法的语义实体识别 (Semantic Entity Recognition, SER) 以及关系抽取 (Relation Extraction, RE) 任务。基于 SER 任务，可以完成对图像中的文本识别与分类；基于 RE 任务，可以完成对图象中的文本内容的关系提取，如判断问题对(pair)。
|
||||
- 支持SER任务和RE任务的自定义训练。
|
||||
- 支持OCR+SER的端到端系统预测与评估。
|
||||
- 支持OCR+SER+RE的端到端系统预测。
|
||||
- 支持SER模型的动转静导出与基于PaddleInfernece的模型推理。
|
||||
|
||||
|
||||
## 2. 精度与性能
|
||||
|
||||
|
||||
我们在 [XFUND](https://github.com/doc-analysis/XFUND) 的中文数据集上对算法进行了评估，SER与RE上的任务性能如下
|
||||
|
||||
|模型|骨干网络|任务|配置文件|hmean|预测耗时(ms)|下载链接|
|
||||
| --- | --- | --- | --- | --- | --- | --- |
|
||||
|VI-LayoutXLM| VI-LayoutXLM-base | SER | [ser_vi_layoutxlm_xfund_zh_udml.yml](../../configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh_udml.yml)|**93.19%**| 15.49|[训练模型](https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_pretrained.tar)|
|
||||
|LayoutXLM| LayoutXLM-base | SER | [ser_layoutxlm_xfund_zh.yml](../../configs/kie/layoutlm_series/ser_layoutxlm_xfund_zh.yml)|90.38%| 19.49 | [训练模型](https://paddleocr.bj.bcebos.com/pplayout/ser_LayoutXLM_xfun_zh.tar)|
|
||||
|VI-LayoutXLM| VI-LayoutXLM-base | RE | [re_vi_layoutxlm_xfund_zh_udml.yml](../../configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh_udml.yml)|**83.92%**| 15.49|[训练模型](https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_pretrained.tar)|
|
||||
|LayoutXLM| LayoutXLM-base | RE | [re_layoutxlm_xfund_zh.yml](../../configs/kie/layoutlm_series/re_layoutxlm_xfund_zh.yml)|74.83%| 19.49|[训练模型](https://paddleocr.bj.bcebos.com/pplayout/re_LayoutXLM_xfun_zh.tar)|
|
||||
|
||||
|
||||
* 注：预测耗时测试条件：V100 GPU + cuda10.2 + cudnn8.1.1 + TensorRT 7.2.3.4，使用FP16进行测试。
|
||||
|
||||
更多关于PaddleOCR中关键信息抽取模型的介绍，请参考[关键信息抽取模型库](../../doc/doc_ch/algorithm_overview.md)。
|
||||
|
||||
|
||||
## 3. 效果演示
|
||||
|
||||
基于多模态模型的关键信息抽取任务有2种主要的解决方案。
|
||||
|
||||
（1）文本检测 + 文本识别 + 语义实体识别(SER)
|
||||
（2）文本检测 + 文本识别 + 语义实体识别(SER) + 关系抽取(RE)
|
||||
|
||||
下面给出SER与RE任务的示例效果，关于上述解决方案的详细介绍，请参考[关键信息抽取全流程指南](./how_to_do_kie.md)。
|
||||
|
||||
### 3.1 SER
|
||||
|
||||
对于SER任务，效果如下所示。
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539141-68e71c75-5cf7-4529-b2ca-219d29fa5f68.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185310636-6ce02f7c-790d-479f-b163-ea97a5a04808.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539517-ccf2372a-f026-4a7c-ad28-c741c770f60a.png" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185539735-37b5c2ef-629d-43fe-9abb-44bb717ef7ee.jpg" width="600">
|
||||
</div>
|
||||
|
||||
**注意：** 测试图片来源于[XFUND数据集](https://github.com/doc-analysis/XFUND)、[发票数据集](https://aistudio.baidu.com/aistudio/datasetdetail/165561)以及合成的身份证数据集。
|
||||
|
||||
|
||||
图中不同颜色的框表示不同的类别。
|
||||
|
||||
图中的发票以及申请表图像，有`QUESTION`, `ANSWER`, `HEADER` 3种类别，识别的`QUESTION`, `ANSWER`可以用于后续的问题与答案的关系抽取。
|
||||
|
||||
图中的身份证图像，则直接识别出其中的`姓名`、`性别`、`民族`等关键信息，这样就无需后续的关系抽取过程，一个模型即可完成关键信息抽取。
|
||||
|
||||
|
||||
### 3.2 RE
|
||||
|
||||
对于RE任务，效果如下所示。
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185393805-c67ff571-cf7e-4217-a4b0-8b396c4f22bb.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185540080-0431e006-9235-4b6d-b63d-0b3c6e1de48f.jpg" width="600">
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185540291-f64e5daf-6d42-4e7c-bbbb-471e3fac4fcc.png" width="600">
|
||||
</div>
|
||||
|
||||
|
||||
红色框是问题，蓝色框是答案。绿色线条表示连接的两端为一个key-value的pair。
|
||||
|
||||
## 4. 使用
|
||||
|
||||
### 4.1 准备环境
|
||||
|
||||
使用下面的命令安装运行SER与RE关键信息抽取的依赖。
|
||||
|
||||
```bash
|
||||
git clone https://github.com/PaddlePaddle/PaddleOCR.git
|
||||
cd PaddleOCR
|
||||
pip install -r requirements.txt
|
||||
pip install -r ppstructure/kie/requirements.txt
|
||||
# 安装PaddleOCR引擎用于预测
|
||||
pip install paddleocr -U
|
||||
```
|
||||
|
||||
NOTE: 对于KIE任务需要降低Paddle框架版本（Paddle<2.6），和PaddleNLP版本（PaddleNLP<2.6）。
|
||||
|
||||
### 4.2 快速开始
|
||||
|
||||
下面XFUND数据集，快速体验SER模型与RE模型。
|
||||
|
||||
#### 4.2.1 准备数据
|
||||
|
||||
```bash
|
||||
mkdir train_data
|
||||
cd train_data
|
||||
# 下载与解压数据
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/dataset/XFUND.tar && tar -xf XFUND.tar
|
||||
cd ..
|
||||
```
|
||||
|
||||
#### 4.2.2 基于动态图的预测
|
||||
|
||||
首先下载模型。
|
||||
|
||||
```bash
|
||||
mkdir pretrained_model
|
||||
cd pretrained_model
|
||||
# 下载并解压SER预训练模型
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_pretrained.tar && tar -xf ser_vi_layoutxlm_xfund_pretrained.tar
|
||||
|
||||
# 下载并解压RE预训练模型
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_pretrained.tar && tar -xf re_vi_layoutxlm_xfund_pretrained.tar
|
||||
```
|
||||
|
||||
如果希望使用OCR引擎，获取端到端的预测结果，可以使用下面的命令进行预测。
|
||||
|
||||
```bash
|
||||
# 仅预测SER模型
|
||||
python3 tools/infer_kie_token_ser.py \
|
||||
-c configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./ppstructure/docs/kie/input/zh_val_42.jpg
|
||||
|
||||
# SER + RE模型串联
|
||||
python3 ./tools/infer_kie_token_ser_re.py \
|
||||
-c configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/re_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/image/zh_val_42.jpg \
|
||||
-c_ser configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o_ser Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy
|
||||
```
|
||||
|
||||
`Global.save_res_path`目录中会保存可视化的结果图像以及预测的文本文件。
|
||||
|
||||
如果想使用自定义OCR模型，可通过如下字段进行设置
|
||||
- `Global.kie_det_model_dir`: 设置检测inference模型地址
|
||||
- `Global.kie_rec_model_dir`: 设置识别inference模型地址
|
||||
|
||||
|
||||
如果希望加载标注好的文本检测与识别结果，仅预测可以使用下面的命令进行预测。
|
||||
|
||||
```bash
|
||||
# 仅预测SER模型
|
||||
python3 tools/infer_kie_token_ser.py \
|
||||
-c configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/val.json \
|
||||
Global.infer_mode=False
|
||||
|
||||
# SER + RE模型串联
|
||||
python3 ./tools/infer_kie_token_ser_re.py \
|
||||
-c configs/kie/vi_layoutxlm/re_vi_layoutxlm_xfund_zh.yml \
|
||||
-o Architecture.Backbone.checkpoints=./pretrained_model/re_vi_layoutxlm_xfund_pretrained/best_accuracy \
|
||||
Global.infer_img=./train_data/XFUND/zh_val/val.json \
|
||||
Global.infer_mode=False \
|
||||
-c_ser configs/kie/vi_layoutxlm/ser_vi_layoutxlm_xfund_zh.yml \
|
||||
-o_ser Architecture.Backbone.checkpoints=./pretrained_model/ser_vi_layoutxlm_xfund_pretrained/best_accuracy
|
||||
```
|
||||
|
||||
#### 4.2.3 基于PaddleInference的预测
|
||||
|
||||
首先下载SER和RE的推理模型。
|
||||
|
||||
```bash
|
||||
mkdir inference
|
||||
cd inference
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/ser_vi_layoutxlm_xfund_infer.tar && tar -xf ser_vi_layoutxlm_xfund_infer.tar
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/vi_layoutxlm/re_vi_layoutxlm_xfund_infer.tar && tar -xf re_vi_layoutxlm_xfund_infer.tar
|
||||
cd ..
|
||||
```
|
||||
|
||||
- SER
|
||||
|
||||
执行下面的命令进行预测。
|
||||
|
||||
```bash
|
||||
cd ppstructure
|
||||
python3 kie/predict_kie_token_ser.py \
|
||||
--kie_algorithm=LayoutXLM \
|
||||
--ser_model_dir=../inference/ser_vi_layoutxlm_xfund_infer \
|
||||
--image_dir=./docs/kie/input/zh_val_42.jpg \
|
||||
--ser_dict_path=../train_data/XFUND/class_list_xfun.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--ocr_order_method="tb-yx"
|
||||
```
|
||||
|
||||
可视化结果保存在`output`目录下。
|
||||
|
||||
- RE
|
||||
|
||||
执行下面的命令进行预测。
|
||||
|
||||
```bash
|
||||
cd ppstructure
|
||||
python3 kie/predict_kie_token_ser_re.py \
|
||||
--kie_algorithm=LayoutXLM \
|
||||
--re_model_dir=../inference/re_vi_layoutxlm_xfund_infer \
|
||||
--ser_model_dir=../inference/ser_vi_layoutxlm_xfund_infer \
|
||||
--use_visual_backbone=False \
|
||||
--image_dir=./docs/kie/input/zh_val_42.jpg \
|
||||
--ser_dict_path=../train_data/XFUND/class_list_xfun.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--ocr_order_method="tb-yx"
|
||||
```
|
||||
|
||||
可视化结果保存在`output`目录下。
|
||||
|
||||
如果想使用自定义OCR模型，可通过如下字段进行设置
|
||||
- `--det_model_dir`: 设置检测inference模型地址
|
||||
- `--rec_model_dir`: 设置识别inference模型地址
|
||||
|
||||
### 4.3 更多
|
||||
|
||||
关于KIE模型的训练评估与推理，请参考：[关键信息抽取教程](../../doc/doc_ch/kie.md)。
|
||||
|
||||
关于文本检测模型的训练评估与推理，请参考：[文本检测教程](../../doc/doc_ch/detection.md)。
|
||||
|
||||
关于文本识别模型的训练评估与推理，请参考：[文本识别教程](../../doc/doc_ch/recognition.md)。
|
||||
|
||||
关于怎样在自己的场景中完成关键信息抽取任务，请参考：[关键信息抽取全流程指南](./how_to_do_kie.md)。
|
||||
|
||||
|
||||
## 5. 参考链接
|
||||
|
||||
- LayoutXLM: Multimodal Pre-training for Multilingual Visually-rich Document Understanding, https://arxiv.org/pdf/2104.08836.pdf
|
||||
- microsoft/unilm/layoutxlm, https://github.com/microsoft/unilm/tree/master/layoutxlm
|
||||
- XFUND dataset, https://github.com/doc-analysis/XFUND
|
||||
|
||||
## 6. License
|
||||
|
||||
The content of this project itself is licensed under the [Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/)
|
||||
168
ppstructure/kie/how_to_do_kie.md
Normal file
168
ppstructure/kie/how_to_do_kie.md
Normal file
@@ -0,0 +1,168 @@
|
||||
|
||||
# 怎样完成基于图像数据的信息抽取任务
|
||||
|
||||
- [1. 简介](#1-简介)
|
||||
- [1.1 背景](#11-背景)
|
||||
- [1.2 主流方法](#12-主流方法)
|
||||
- [2. 关键信息抽取任务流程](#2-关键信息抽取任务流程)
|
||||
- [2.1 训练OCR模型](#21-训练OCR模型)
|
||||
- [2.2 训练KIE模型](#22-训练KIE模型)
|
||||
- [3. 参考文献](#3-参考文献)
|
||||
|
||||
|
||||
## 1. 简介
|
||||
|
||||
### 1.1 背景
|
||||
|
||||
关键信息抽取 (Key Information Extraction, KIE)指的是是从文本或者图像中，抽取出关键的信息。针对文档图像的关键信息抽取任务作为OCR的下游任务，存在非常多的实际应用场景，如表单识别、车票信息抽取、身份证信息抽取等。然而，使用人力从这些文档图像中提取或者收集关键信息耗时费力，怎样自动化融合图像中的视觉、布局、文字等特征并完成关键信息抽取是一个价值与挑战并存的问题。
|
||||
|
||||
对于特定场景的文档图像，其中的关键信息位置、版式等较为固定，因此在研究早期有很多基于模板匹配的方法进行关键信息的抽取，考虑到其流程较为简单，该方法仍然被广泛应用在目前的很多场景中。但是这种基于模板匹配的方法在应用到不同的场景中时，需要耗费大量精力去调整与适配模板，迁移成本较高。
|
||||
|
||||
文档图像中的KIE一般包含2个子任务，示意图如下图所示。
|
||||
|
||||
* （1）SER: 语义实体识别 (Semantic Entity Recognition)，对每一个检测到的文本进行分类，如将其分为姓名，身份证。如下图中的黑色框和红色框。
|
||||
* （2）RE: 关系抽取 (Relation Extraction)，对每一个检测到的文本进行分类，如将其分为问题 (key) 和答案 (value) 。然后对每一个问题找到对应的答案，相当于完成key-value的匹配过程。如下图中的红色框和黑色框分别代表问题和答案，黄色线代表问题和答案之间的对应关系。
|
||||
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/184588654-d87f54f3-13ab-42c4-afc0-da79bead3f14.png" width="800">
|
||||
</div>
|
||||
|
||||
|
||||
### 1.2 基于深度学习的主流方法
|
||||
|
||||
一般的KIE方法基于命名实体识别(Named Entity Recognition,NER)来展开研究，但是此类方法仅使用了文本信息而忽略了位置与视觉特征信息，因此精度受限。近几年大多学者开始融合多个模态的输入信息，进行特征融合，并对多模态信息进行处理，从而提升KIE的精度。主要方法有以下几种
|
||||
|
||||
* （1）基于Grid的方法：此类方法主要关注图像层面多模态信息的融合，文本大多大多为字符粒度，对文本与结构结构信息的嵌入方式较为简单，如Chargrid[1]等算法。
|
||||
* （2）基于Token的方法：此类方法参考NLP中的BERT等方法，将位置、视觉等特征信息共同编码到多模态模型中，并且在大规模数据集上进行预训练，从而在下游任务中，仅需要少量的标注数据便可以获得很好的效果。如LayoutLM[2], LayoutLMv2[3], LayoutXLM[4], StrucText[5]等算法。
|
||||
* （3）基于GCN的方法：此类方法尝试学习图像、文字之间的结构信息，从而可以解决开集信息抽取的问题（训练集中没有见过的模板），如GCN[6]、SDMGR[7]等算法。
|
||||
* （4）基于End-to-end的方法：此类方法将现有的OCR文字识别以及KIE信息抽取2个任务放在一个统一的网络中进行共同学习，并在学习过程中相互加强。如Trie[8]等算法。
|
||||
|
||||
更多关于该系列算法的详细介绍，请参考“动手学OCR·十讲”课程的课节六部分：[文档分析理论与实践](https://aistudio.baidu.com/aistudio/education/group/info/25207)。
|
||||
|
||||
## 2. 关键信息抽取任务流程
|
||||
|
||||
PaddleOCR中实现了LayoutXLM等算法（基于Token），同时，在PP-StructureV2中，对LayoutXLM多模态预训练模型的网络结构进行简化，去除了其中的Visual backbone部分，设计了视觉无关的VI-LayoutXLM模型，同时引入符合人类阅读顺序的排序逻辑以及UDML知识蒸馏策略，最终同时提升了关键信息抽取模型的精度与推理速度。
|
||||
|
||||
下面介绍怎样基于PaddleOCR完成关键信息抽取任务。
|
||||
|
||||
在非End-to-end的KIE方法中，完成关键信息抽取，至少需要**2个步骤**：首先使用OCR模型，完成文字位置与内容的提取，然后使用KIE模型，根据图像、文字位置以及文字内容，提取出其中的关键信息。
|
||||
|
||||
### 2.1 训练OCR模型
|
||||
|
||||
#### 2.1.1 文本检测
|
||||
|
||||
**（1）数据**
|
||||
|
||||
PaddleOCR中提供的模型大多数为通用模型，在进行文本检测的过程中，相邻文本行的检测一般是根据位置的远近进行区分，如上图，使用PP-OCRv3通用中英文检测模型进行文本检测时，容易将”民族“与“汉”这2个代表不同的字段检测到一起，从而增加后续KIE任务的难度。因此建议在做KIE任务的过程中，首先训练一个针对该文档数据集的检测模型。
|
||||
|
||||
在数据标注时，关键信息的标注需要隔开，比上图中的 “民族汉” 3个字相隔较近，此时需要将”民族“与”汉“标注为2个文本检测框，否则会增加后续KIE任务的难度。
|
||||
|
||||
对于下游任务，一般来说，`200~300`张的文本训练数据即可保证基本的训练效果，如果没有太多的先验知识，可以先标注 **`200~300`** 张图片，进行后续文本检测模型的训练。
|
||||
|
||||
|
||||
**（2）模型**
|
||||
|
||||
在模型选择方面，推荐使用PP-OCRv3_det，关于更多关于检测模型的训练方法介绍，请参考：[OCR文本检测模型训练教程](../../doc/doc_ch/detection.md)与[PP-OCRv3 文本检测模型训练教程](../../doc/doc_ch/PPOCRv3_det_train.md)。
|
||||
|
||||
#### 2.1.2 文本识别
|
||||
|
||||
相对自然场景，文档图像中的文本内容识别难度一般相对较低（背景相对不太复杂），因此**优先建议**尝试PaddleOCR中提供的PP-OCRv3通用文本识别模型([PP-OCRv3模型库链接](../../doc/doc_ch/models_list.md))。
|
||||
|
||||
**（1）数据**
|
||||
|
||||
然而，在部分文档场景中也会存在一些挑战，如身份证场景中存在着罕见字，在发票等场景中的字体比较特殊，这些问题都会增加文本识别的难度，此时如果希望保证或者进一步提升模型的精度，建议基于特定文档场景的文本识别数据集，加载PP-OCRv3模型进行微调。
|
||||
|
||||
在模型微调的过程中，建议准备至少`5000`张垂类场景的文本识别图像，可以保证基本的模型微调效果。如果希望提升模型的精度与泛化能力，可以合成更多与该场景类似的文本识别数据，从公开数据集中收集通用真实文本识别数据，一并添加到该场景的文本识别训练任务过程中。在训练过程中，建议每个epoch的真实垂类数据、合成数据、通用数据比例在`1:1:1`左右，这可以通过设置不同数据源的采样比例进行控制。如有3个训练文本文件，分别包含1W、2W、5W条数据，那么可以在配置文件中设置数据如下：
|
||||
|
||||
```yml
|
||||
Train:
|
||||
dataset:
|
||||
name: SimpleDataSet
|
||||
data_dir: ./train_data/
|
||||
label_file_list:
|
||||
- ./train_data/train_list_1W.txt
|
||||
- ./train_data/train_list_2W.txt
|
||||
- ./train_data/train_list_5W.txt
|
||||
ratio_list: [1.0, 0.5, 0.2]
|
||||
...
|
||||
```
|
||||
|
||||
**（2）模型**
|
||||
|
||||
在模型选择方面，推荐使用通用中英文文本识别模型PP-OCRv3_rec，关于更多关于文本识别模型的训练方法介绍，请参考：[OCR文本识别模型训练教程](../../doc/doc_ch/recognition.md)与[PP-OCRv3文本识别模型库与配置文件](../../doc/doc_ch/models_list.md)。
|
||||
|
||||
### 2.2 训练KIE模型
|
||||
|
||||
对于识别得到的文字进行关键信息抽取，有2种主要的方法。
|
||||
|
||||
（1）直接使用SER，获取关键信息的类别：如身份证场景中，将“姓名“与”张三“分别标记为`name_key`与`name_value`。最终识别得到的类别为`name_value`对应的**文本字段**即为我们所需要的关键信息。
|
||||
|
||||
（2）联合SER与RE进行使用：这种方法中，首先使用SER，获取图像文字内容中所有的key与value，然后使用RE方法，对所有的key与value进行配对，找到映射关系，从而完成关键信息的抽取。
|
||||
|
||||
#### 2.2.1 SER
|
||||
|
||||
以身份证场景为例， 关键信息一般包含`姓名`、`性别`、`民族`等，我们直接将对应的字段标注为特定的类别即可，如下图所示。
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/184526682-8b810397-5a93-4395-93da-37b8b8494c41.png" width="500">
|
||||
</div>
|
||||
|
||||
**注意：**
|
||||
|
||||
- 标注过程中，对于无关于KIE关键信息的文本内容，均需要将其标注为`other`类别，相当于背景信息。如在身份证场景中，如果我们不关注性别信息，那么可以将“性别”与“男”这2个字段的类别均标注为`other`。
|
||||
- 标注过程中，需要以**文本行**为单位进行标注，无需标注单个字符的位置信息。
|
||||
|
||||
数据量方面，一般来说，对于比较固定的场景，**50张**左右的训练图片即可达到可以接受的效果，可以使用[PPOCRLabel](https://github.com/PFCCLab/PPOCRLabel/blob/main/README_ch.md)完成KIE的标注过程。
|
||||
|
||||
模型方面，推荐使用PP-StructureV2中提出的VI-LayoutXLM模型，它基于LayoutXLM模型进行改进，去除其中的视觉特征提取模块，在精度基本无损的情况下，进一步提升了模型推理速度。更多教程请参考：[VI-LayoutXLM算法介绍](../../doc/doc_ch/algorithm_kie_vi_layoutxlm.md)与[KIE关键信息抽取使用教程](../../doc/doc_ch/kie.md)。
|
||||
|
||||
|
||||
#### 2.2.2 SER + RE
|
||||
|
||||
该过程主要包含SER与RE 2个过程。SER阶段主要用于识别出文档图像中的所有key与value，RE阶段主要用于对所有的key与value进行匹配。
|
||||
|
||||
以身份证场景为例， 关键信息一般包含`姓名`、`性别`、`民族`等关键信息，在SER阶段，我们需要识别所有的question (key) 与answer (value) 。标注如下所示。每个字段的类别信息（`label`字段）可以是question、answer或者other（与待抽取的关键信息无关的字段）
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/184526785-c3d2d310-cd57-4d31-b933-912716b29856.jpg" width="500">
|
||||
</div>
|
||||
|
||||
|
||||
在RE阶段，需要标注每个字段的的id与连接信息，如下图所示。
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/184528728-626f77eb-fd9f-4709-a7dc-5411cc417dab.jpg" width="500">
|
||||
</div>
|
||||
|
||||
每个文本行字段中，需要添加`id`与`linking`字段信息，`id`记录该文本行的唯一标识，同一张图片中的不同文本内容不能重复，`linking`是一个列表，记录了不同文本之间的连接信息。如字段“出生”的id为0，字段“1996年1月11日”的id为1，那么它们均有[[0, 1]]的`linking`标注，表示该id=0与id=1的字段构成key-value的关系（姓名、性别等字段类似，此处不再一一赘述）。
|
||||
|
||||
|
||||
**注意：**
|
||||
|
||||
- 标注过程中，如果value是多个字符，那么linking中可以新增一个key-value对，如`[[0, 1], [0, 2]]`
|
||||
|
||||
|
||||
数据量方面，一般来说，对于比较固定的场景，**50张**左右的训练图片即可达到可以接受的效果，可以使用PPOCRLabel完成KIE的标注过程。
|
||||
|
||||
模型方面，推荐使用PP-StructureV2中提出的VI-LayoutXLM模型，它基于LayoutXLM模型进行改进，去除其中的视觉特征提取模块，在精度基本无损的情况下，进一步提升了模型推理速度。更多教程请参考：[VI-LayoutXLM算法介绍](../../doc/doc_ch/algorithm_kie_vi_layoutxlm.md)与[KIE关键信息抽取使用教程](../../doc/doc_ch/kie.md)。
|
||||
|
||||
|
||||
## 3. 参考文献
|
||||
|
||||
|
||||
[1] Katti A R, Reisswig C, Guder C, et al. Chargrid: Towards understanding 2d documents[J]. arXiv preprint arXiv:1809.08799, 2018.
|
||||
|
||||
[2] Xu Y, Li M, Cui L, et al. Layoutlm: Pre-training of text and layout for document image understanding[C]//Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020: 1192-1200.
|
||||
|
||||
[3] Xu Y, Xu Y, Lv T, et al. LayoutLMv2: Multi-modal pre-training for visually-rich document understanding[J]. arXiv preprint arXiv:2012.14740, 2020.
|
||||
|
||||
[4]: Xu Y, Lv T, Cui L, et al. Layoutxlm: Multimodal pre-training for multilingual visually-rich document understanding[J]. arXiv preprint arXiv:2104.08836, 2021.
|
||||
|
||||
[5] Li Y, Qian Y, Yu Y, et al. StrucTexT: Structured Text Understanding with Multi-Modal Transformers[C]//Proceedings of the 29th ACM International Conference on Multimedia. 2021: 1912-1920.
|
||||
|
||||
[6] Liu X, Gao F, Zhang Q, et al. Graph convolution for multimodal information extraction from visually rich documents[J]. arXiv preprint arXiv:1903.11279, 2019.
|
||||
|
||||
[7] Sun H, Kuang Z, Yue X, et al. Spatial Dual-Modality Graph Reasoning for Key Information Extraction[J]. arXiv preprint arXiv:2103.14470, 2021.
|
||||
|
||||
[8] Zhang P, Xu Y, Cheng Z, et al. Trie: End-to-end text reading and information extraction for document understanding[C]//Proceedings of the 28th ACM International Conference on Multimedia. 2020: 1413-1422.
|
||||
179
ppstructure/kie/how_to_do_kie_en.md
Normal file
179
ppstructure/kie/how_to_do_kie_en.md
Normal file
@@ -0,0 +1,179 @@
|
||||
|
||||
# Key Information Extraction Pipeline
|
||||
|
||||
- [1. Introduction](#1-Introduction)
|
||||
- [1.1 Background](#11-Background)
|
||||
- [1.2 Mainstream Deep-learning Solutions](#12-Mainstream-Deep-learning-Solutions)
|
||||
- [2. KIE Pipeline](#2-KIE-Pipeline)
|
||||
- [2.1 Train OCR Models](#21-Train-OCR-Models)
|
||||
- [2.2 Train KIE Models](#22-Train-KIE-Models)
|
||||
- [3. Reference](#3-Reference)
|
||||
|
||||
|
||||
## 1. Introduction
|
||||
|
||||
### 1.1 Background
|
||||
|
||||
Key information extraction (KIE) refers to extracting key information from text or images. As the downstream task of OCR, KIE of document image has many practical application scenarios, such as form recognition, ticket information extraction, ID card information extraction, etc. However, it is time-consuming and laborious to extract key information from these document images by manpower. It's challengable but also valuable to combine multi-modal features (visual, layout, text, etc) together and complete KIE tasks.
|
||||
|
||||
For the document images in a specific scene, the position and layout of the key information are relatively fixed. Therefore, in the early stage of the research, there are many methods based on template matching to extract the key information. This method is still widely used in many simple scenarios at present. However, it takes long time to adjut the template for different scenarios.
|
||||
|
||||
|
||||
The KIE in the document image generally contains 2 subtasks, which is as shown follows.
|
||||
|
||||
* (1) SER: semantic entity recognition, which classifies each detected textline, such as dividing it into name and ID No. As shown in the red boxes in the following figure.
|
||||
|
||||
* (2) RE: relationship extraction, which matches the question and answer based on SER results. As shown in the figure below, the yellow arrows match the question and answer.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185726510-faba470d-2c79-4784-b8da-6c1aa5af9572.png" width="800">
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
### 1.2 Mainstream Deep-learning Solutions
|
||||
|
||||
General KIE methods are based on Named Entity Recognition (NER), but such methods only use text information and ignore location and visual feature information, which leads to limited accuracy. In recent years, most scholars have started to combine mutil-modal features to improve the accuracy of KIE model. The main methods are as follows:
|
||||
|
||||
* (1) Grid based methods. These methods mainly focus on the fusion of multi-modal information at the image level. Most texts are of character granularity. The text and structure information embedding method is simple, such as the algorithm of chargrid [1].
|
||||
|
||||
* (2) Token based methods. These methods refer to the NLP methods such as Bert, which encode the position, vision and other feature information into the multi-modal model, and conduct pre-training on large-scale datasets, so that in downstream tasks, only a small amount of annotation data is required to obtain excellent results. The representative algorithms are layoutlm [2], layoutlmv2 [3], layoutxlm [4], structext [5], etc.
|
||||
|
||||
* (3) GCN based methods. These methods try to learn the structural information between images and characters, so as to solve the problem of extracting open set information (templates not seen in the training set), such as GCN [6], SDMGR [7] and other algorithms.
|
||||
|
||||
* (4) End to end based methods: these methods put the existing OCR character recognition and KIE information extraction tasks into a unified network for common learning, and strengthen each other in the learning process. Such as TRIE [8].
|
||||
|
||||
|
||||
For more detailed introduction of the algorithms, please refer to Chapter 6 of [Diving into OCR](https://aistudio.baidu.com/aistudio/education/group/info/25207).
|
||||
|
||||
## 2. KIE Pipeline
|
||||
|
||||
Token based methods such as LayoutXLM are implemented in PaddleOCR. What's more, in PP-StructureV2, we simplify the LayoutXLM model and proposed VI-LayoutXLM, in which the visual feature extraction module is removed for speed-up. The textline sorting strategy conforming to the human reading order and UDML knowledge distillation strategy are utilized for higher model accuracy.
|
||||
|
||||
|
||||
In the non end-to-end KIE method, KIE needs at least **2 steps**. Firstly, the OCR model is used to extract the text and its position. Secondly, the KIE model is used to extract the key information according to the image, text position and text content.
|
||||
|
||||
|
||||
### 2.1 Train OCR Models
|
||||
|
||||
#### 2.1.1 Text Detection
|
||||
|
||||
**(1) Data**
|
||||
|
||||
Most of the models provided in PaddleOCR are general models. In the process of text detection, the detection of adjacent text lines is generally based on the distance of the position. As shown in the figure above, when using PP-OCRv3 general English detection model for text detection, it is easy to detect the two fields representing different properties as one. Therefore, it is suggested to finetune a detection model according to your scenario firstly during the KIE task.
|
||||
|
||||
|
||||
During data annotation, the different key information needs to be separated. Otherwise, it will increase the difficulty of subsequent KIE tasks.
|
||||
|
||||
For downstream tasks, generally speaking, `200~300` training images can guarantee the basic training effect. If there is not too much prior knowledge, **`200~300`** images can be labeled firstly for subsequent text detection model training.
|
||||
|
||||
**(2) Model**
|
||||
|
||||
In terms of model selection, PP-OCRv3 detection model is recommended. For more information about the training methods of the detection model, please refer to: [Text detection tutorial](../../doc/doc_en/detection_en.md) and [PP-OCRv3 detection model tutorial](../../doc/doc_ch/PPOCRv3_det_train.md).
|
||||
|
||||
#### 2.1.2 Text recognition
|
||||
|
||||
|
||||
Compared with the natural scene, the text recognition in the document image is generally relatively easier (the background is not too complex), so **it is suggested to** try the PP-OCRv3 general text recognition model provided in PaddleOCR ([PP-OCRv3 model list](../../doc/doc_en/models_list_en.md))
|
||||
|
||||
|
||||
**(1) Data**
|
||||
|
||||
However, there are also some challenges in some document scenarios, such as rare words in ID card scenarios and special fonts in invoice and other scenarios. These problems will increase the difficulty of text recognition. At this time, if you want to ensure or further improve the model accuracy, it is recommended to load PP-OCRv3 model based on the text recognition dataset of specific document scenarios for finetuning.
|
||||
|
||||
In the process of model finetuning, it is recommended to prepare at least `5000` vertical scene text recognition images to ensure the basic model fine-tuning effect. If you want to improve the accuracy and generalization ability of the model, you can synthesize more text recognition images similar to the scene, collect general real text recognition data from the public data set, and add them to the text recognition training process. In the training process, it is suggested that the ratio of real data, synthetic data and general data of each epoch should be around `1:1:1`, which can be controlled by setting the sampling ratio of different data sources. If there are 3 training text files, including 10k, 20k and 50k pieces of data respectively, the data can be set in the configuration file as follows:
|
||||
|
||||
```yml
|
||||
Train:
|
||||
dataset:
|
||||
name: SimpleDataSet
|
||||
data_dir: ./train_data/
|
||||
label_file_list:
|
||||
- ./train_data/train_list_10k.txt
|
||||
- ./train_data/train_list_10k.txt
|
||||
- ./train_data/train_list_50k.txt
|
||||
ratio_list: [1.0, 0.5, 0.2]
|
||||
...
|
||||
```
|
||||
|
||||
**(2) Model**
|
||||
|
||||
In terms of model selection, PP-OCRv3 recognition model is recommended. For more information about the training methods of the recognition model, please refer to: [Text recognition tutorial](../../doc/doc_en/recognition_en.md) and [PP-OCRv3 model list](../../doc/doc_en/models_list_en.md).
|
||||
|
||||
|
||||
### 2.2 Train KIE Models
|
||||
|
||||
There are two main methods to extract the key information from the recognized texts.
|
||||
|
||||
(1) Directly use SER model to obtain the key information category. For example, in the ID card scenario, we mark "name" and "Geoff Sample" as "name_key" and "name_value", respectively. The **text field** corresponding to the category "name_value" finally identified is the key information we need.
|
||||
|
||||
(2) Joint use SER and RE models. For this case, we firstly use SER model to obtain all questions (keys) and questions (values) for the image text, and then use RE model to match all keys and values to find the relationship, so as to complete the extraction of key information.
|
||||
|
||||
#### 2.2.1 SER
|
||||
|
||||
Take the ID card scenario as an example. The key information generally includes `name`, `DOB`, etc. We can directly mark the corresponding fields as specific categories, as shown in the following figure.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185728456-dc396f47-0880-4279-9c7c-c99601bf16a7.png" width="500">
|
||||
</div>
|
||||
|
||||
**Note:**
|
||||
|
||||
- In the labeling process, text content without key information about KIE shall be labeled as`other`, which is equivalent to background information. For example, in the ID card scenario, if we do not pay attention to `DOB` information, we can mark the categories of `DOB` and `Area manager` as `other`.
|
||||
- In the annotation process of, it is required to annotate the **textline** position rather than the character.
|
||||
|
||||
|
||||
In terms of data, generally speaking, for relatively fixed scenes, **50** training images can achieve acceptable effects. You can refer to [PPOCRLabel](https://github.com/PFCCLab/PPOCRLabel/blob/main/README.md) for finish the labeling process.
|
||||
|
||||
In terms of model, it is recommended to use the VI-layoutXLM model proposed in PP-StructureV2. It is improved based on the LayoutXLM model, removing the visual feature extraction module, and further improving the model inference speed without the significant reduction on model accuracy. For more tutorials, please refer to [VI-LayoutXLM introduction](../../doc/doc_en/algorithm_kie_vi_layoutxlm_en.md) and [KIE tutorial](../../doc/doc_en/kie_en.md).
|
||||
|
||||
|
||||
#### 2.2.2 SER + RE
|
||||
|
||||
The SER model is mainly used to identify all keys and values in the document image, and the RE model is mainly used to match all keys and values.
|
||||
|
||||
Taking the ID card scenario as an example, the key information generally includes key information such as `name`, `DOB`, etc. in the SER stage, we need to identify all questions (keys) and answers (values). The demo annotation is as follows. All keys can be annotated as `question`, and all values can be annotated as `answer`.
|
||||
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185728881-b6055e01-034c-4584-aaa6-97c9c25fb61b.png" width="500">
|
||||
</div>
|
||||
|
||||
|
||||
In the RE stage, the ID and connection information of each field need to be marked, as shown in the following figure.
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/14270174/185728948-a4208013-5038-4025-9a93-0c6d51447488.png" width="500">
|
||||
</div>
|
||||
|
||||
For each textline, you need to add 'ID' and 'linking' field information. The 'ID' records the unique identifier of the textline. Different text contents in the same images cannot be repeated. The 'linking' is a list that records the connection information between different texts. If the ID of the field "name" is 0 and the ID of the field "Geoff Sample" is 1, then they all have [[0, 1]] 'linking' marks, indicating that the fields with `id=0` and `id=1` form a key value relationship (the fields such as DOB and Expires are similar, and will not be repeated here).
|
||||
|
||||
|
||||
**Note:**
|
||||
|
||||
During annotation, if value is multiple text lines, a key-value pair can be added in linking, such as `[[0, 1], [0, 2]]`.
|
||||
|
||||
In terms of data, generally speaking, for relatively fixed scenes, about **50** training images can achieve acceptable effects.
|
||||
|
||||
In terms of model, it is recommended to use the VI-layoutXLM model proposed in PP-StructureV2. It is improved based on the LayoutXLM model, removing the visual feature extraction module, and further improving the model inference speed without the significant reduction on model accuracy. For more tutorials, please refer to [VI-LayoutXLM introduction](../../doc/doc_en/algorithm_kie_vi_layoutxlm_en.md) and [KIE tutorial](../../doc/doc_en/kie_en.md).
|
||||
|
||||
|
||||
|
||||
## 3. Reference
|
||||
|
||||
|
||||
[1] Katti A R, Reisswig C, Guder C, et al. Chargrid: Towards understanding 2d documents[J]. arXiv preprint arXiv:1809.08799, 2018.
|
||||
|
||||
[2] Xu Y, Li M, Cui L, et al. Layoutlm: Pre-training of text and layout for document image understanding[C]//Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2020: 1192-1200.
|
||||
|
||||
[3] Xu Y, Xu Y, Lv T, et al. LayoutLMv2: Multi-modal pre-training for visually-rich document understanding[J]. arXiv preprint arXiv:2012.14740, 2020.
|
||||
|
||||
[4]: Xu Y, Lv T, Cui L, et al. Layoutxlm: Multimodal pre-training for multilingual visually-rich document understanding[J]. arXiv preprint arXiv:2104.08836, 2021.
|
||||
|
||||
[5] Li Y, Qian Y, Yu Y, et al. StrucTexT: Structured Text Understanding with Multi-Modal Transformers[C]//Proceedings of the 29th ACM International Conference on Multimedia. 2021: 1912-1920.
|
||||
|
||||
[6] Liu X, Gao F, Zhang Q, et al. Graph convolution for multimodal information extraction from visually rich documents[J]. arXiv preprint arXiv:1903.11279, 2019.
|
||||
|
||||
[7] Sun H, Kuang Z, Yue X, et al. Spatial Dual-Modality Graph Reasoning for Key Information Extraction[J]. arXiv preprint arXiv:2103.14470, 2021.
|
||||
|
||||
[8] Zhang P, Xu Y, Cheng Z, et al. Trie: End-to-end text reading and information extraction for document understanding[C]//Proceedings of the 28th ACM International Conference on Multimedia. 2020: 1413-1422.
|
||||
191
ppstructure/kie/predict_kie_token_ser.py
Normal file
191
ppstructure/kie/predict_kie_token_ser.py
Normal file
@@ -0,0 +1,191 @@
|
||||
# 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.insert(0, 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
|
||||
from ppstructure.utility import parse_args
|
||||
|
||||
from paddleocr import PaddleOCR
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class SerPredictor(object):
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.ocr_engine = PaddleOCR(
|
||||
use_angle_cls=args.use_angle_cls,
|
||||
det_model_dir=args.det_model_dir,
|
||||
rec_model_dir=args.rec_model_dir,
|
||||
show_log=False,
|
||||
use_gpu=args.use_gpu,
|
||||
)
|
||||
|
||||
pre_process_list = [
|
||||
{
|
||||
"VQATokenLabelEncode": {
|
||||
"algorithm": args.kie_algorithm,
|
||||
"class_path": args.ser_dict_path,
|
||||
"contains_re": False,
|
||||
"ocr_engine": self.ocr_engine,
|
||||
"order_method": args.ocr_order_method,
|
||||
}
|
||||
},
|
||||
{"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)
|
||||
if data[0] is None:
|
||||
return None, 0
|
||||
starttime = time.time()
|
||||
|
||||
for idx in range(len(data)):
|
||||
if isinstance(data[idx], np.ndarray):
|
||||
data[idx] = np.expand_dims(data[idx], axis=0)
|
||||
else:
|
||||
data[idx] = [data[idx]]
|
||||
if self.args.use_onnx:
|
||||
input_tensor = {
|
||||
name: data[idx] for idx, name in enumerate(self.input_tensor)
|
||||
}
|
||||
self.output_tensors = self.predictor.run(None, input_tensor)
|
||||
else:
|
||||
for idx in range(len(self.input_tensor)):
|
||||
self.input_tensor[idx].copy_from_cpu(data[idx])
|
||||
|
||||
self.predictor.run()
|
||||
|
||||
outputs = []
|
||||
for output_tensor in self.output_tensors:
|
||||
output = (
|
||||
output_tensor if self.args.use_onnx else 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, data, 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(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=args.vis_font_path,
|
||||
)
|
||||
|
||||
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())
|
||||
144
ppstructure/kie/predict_kie_token_ser_re.py
Normal file
144
ppstructure/kie/predict_kie_token_ser_re.py
Normal file
@@ -0,0 +1,144 @@
|
||||
# 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.insert(0, 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 tools.infer_kie_token_ser_re import make_input
|
||||
from ppocr.postprocess import build_post_process
|
||||
from ppocr.utils.logging import get_logger
|
||||
from ppocr.utils.visual import draw_ser_results, draw_re_results
|
||||
from ppocr.utils.utility import get_image_file_list, check_and_read
|
||||
from ppstructure.utility import parse_args
|
||||
from ppstructure.kie.predict_kie_token_ser import SerPredictor
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class SerRePredictor(object):
|
||||
def __init__(self, args):
|
||||
self.use_visual_backbone = args.use_visual_backbone
|
||||
self.ser_engine = SerPredictor(args)
|
||||
if args.re_model_dir is not None:
|
||||
postprocess_params = {"name": "VQAReTokenLayoutLMPostProcess"}
|
||||
self.postprocess_op = build_post_process(postprocess_params)
|
||||
(
|
||||
self.predictor,
|
||||
self.input_tensor,
|
||||
self.output_tensors,
|
||||
self.config,
|
||||
) = utility.create_predictor(args, "re", logger)
|
||||
else:
|
||||
self.predictor = None
|
||||
|
||||
def __call__(self, img):
|
||||
starttime = time.time()
|
||||
ser_results, ser_inputs, ser_elapse = self.ser_engine(img)
|
||||
if self.predictor is None:
|
||||
return ser_results, ser_elapse
|
||||
|
||||
re_input, entity_idx_dict_batch = make_input(ser_inputs, ser_results)
|
||||
if self.use_visual_backbone == False:
|
||||
re_input.pop(4)
|
||||
for idx in range(len(self.input_tensor)):
|
||||
self.input_tensor[idx].copy_from_cpu(re_input[idx])
|
||||
|
||||
self.predictor.run()
|
||||
outputs = []
|
||||
for output_tensor in self.output_tensors:
|
||||
output = output_tensor.copy_to_cpu()
|
||||
outputs.append(output)
|
||||
preds = dict(
|
||||
loss=outputs[1],
|
||||
pred_relations=outputs[2],
|
||||
hidden_states=outputs[0],
|
||||
)
|
||||
|
||||
post_result = self.postprocess_op(
|
||||
preds, ser_results=ser_results, entity_idx_dict_batch=entity_idx_dict_batch
|
||||
)
|
||||
|
||||
elapse = time.time() - starttime
|
||||
return post_result, elapse
|
||||
|
||||
|
||||
def main(args):
|
||||
image_file_list = get_image_file_list(args.image_dir)
|
||||
ser_re_predictor = SerRePredictor(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(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
|
||||
re_res, elapse = ser_re_predictor(img)
|
||||
re_res = re_res[0]
|
||||
|
||||
res_str = "{}\t{}\n".format(
|
||||
image_file,
|
||||
json.dumps(
|
||||
{
|
||||
"ocr_info": re_res,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
)
|
||||
f_w.write(res_str)
|
||||
if ser_re_predictor.predictor is not None:
|
||||
img_res = draw_re_results(
|
||||
image_file, re_res, font_path=args.vis_font_path
|
||||
)
|
||||
img_save_path = os.path.join(
|
||||
args.output,
|
||||
os.path.splitext(os.path.basename(image_file))[0] + "_ser_re.jpg",
|
||||
)
|
||||
else:
|
||||
img_res = draw_ser_results(
|
||||
image_file, re_res, font_path=args.vis_font_path
|
||||
)
|
||||
img_save_path = os.path.join(
|
||||
args.output,
|
||||
os.path.splitext(os.path.basename(image_file))[0] + "_ser.jpg",
|
||||
)
|
||||
|
||||
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())
|
||||
7
ppstructure/kie/requirements.txt
Normal file
7
ppstructure/kie/requirements.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
sentencepiece
|
||||
yacs
|
||||
seqeval
|
||||
pypandoc
|
||||
attrdict3
|
||||
python_docx
|
||||
paddlenlp==2.5.2
|
||||
259
ppstructure/kie/tools/eval_with_label_end2end.py
Normal file
259
ppstructure/kie/tools/eval_with_label_end2end.py
Normal file
@@ -0,0 +1,259 @@
|
||||
# Copyright (c) 2021 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 re
|
||||
import sys
|
||||
import shapely
|
||||
from shapely.geometry import Polygon
|
||||
import numpy as np
|
||||
from collections import defaultdict
|
||||
import operator
|
||||
from rapidfuzz.distance import Levenshtein
|
||||
import argparse
|
||||
import json
|
||||
import copy
|
||||
|
||||
|
||||
def parse_ser_results_fp(fp, fp_type="gt", ignore_background=True):
|
||||
# img/zh_val_0.jpg {
|
||||
# "height": 3508,
|
||||
# "width": 2480,
|
||||
# "ocr_info": [
|
||||
# {"text": "Maribyrnong", "label": "other", "bbox": [1958, 144, 2184, 198]},
|
||||
# {"text": "CITYCOUNCIL", "label": "other", "bbox": [2052, 183, 2171, 214]},
|
||||
# ]
|
||||
assert fp_type in ["gt", "pred"]
|
||||
key = "label" if fp_type == "gt" else "pred"
|
||||
res_dict = dict()
|
||||
with open(fp, "r", encoding="utf-8") as fin:
|
||||
lines = fin.readlines()
|
||||
|
||||
for _, line in enumerate(lines):
|
||||
img_path, info = line.strip().split("\t")
|
||||
# get key
|
||||
image_name = os.path.basename(img_path)
|
||||
res_dict[image_name] = []
|
||||
# get infos
|
||||
json_info = json.loads(info)
|
||||
for single_ocr_info in json_info["ocr_info"]:
|
||||
label = single_ocr_info[key].upper()
|
||||
if label in ["O", "OTHERS", "OTHER"]:
|
||||
label = "O"
|
||||
if ignore_background and label == "O":
|
||||
continue
|
||||
single_ocr_info["label"] = label
|
||||
res_dict[image_name].append(copy.deepcopy(single_ocr_info))
|
||||
return res_dict
|
||||
|
||||
|
||||
def polygon_from_str(polygon_points):
|
||||
"""
|
||||
Create a shapely polygon object from gt or dt line.
|
||||
"""
|
||||
polygon_points = np.array(polygon_points).reshape(4, 2)
|
||||
polygon = Polygon(polygon_points).convex_hull
|
||||
return polygon
|
||||
|
||||
|
||||
def polygon_iou(poly1, poly2):
|
||||
"""
|
||||
Intersection over union between two shapely polygons.
|
||||
"""
|
||||
if not poly1.intersects(poly2): # this test is fast and can accelerate calculation
|
||||
iou = 0
|
||||
else:
|
||||
try:
|
||||
inter_area = poly1.intersection(poly2).area
|
||||
union_area = poly1.area + poly2.area - inter_area
|
||||
iou = float(inter_area) / union_area
|
||||
except shapely.geos.TopologicalError:
|
||||
# except Exception as e:
|
||||
# print(e)
|
||||
print("shapely.geos.TopologicalError occurred, iou set to 0")
|
||||
iou = 0
|
||||
return iou
|
||||
|
||||
|
||||
def ed(args, str1, str2):
|
||||
if args.ignore_space:
|
||||
str1 = str1.replace(" ", "")
|
||||
str2 = str2.replace(" ", "")
|
||||
if args.ignore_case:
|
||||
str1 = str1.lower()
|
||||
str2 = str2.lower()
|
||||
return Levenshtein.distance(str1, str2)
|
||||
|
||||
|
||||
def convert_bbox_to_polygon(bbox):
|
||||
"""
|
||||
bbox : [x1, y1, x2, y2]
|
||||
output: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]
|
||||
"""
|
||||
xmin, ymin, xmax, ymax = bbox
|
||||
poly = [[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]
|
||||
return poly
|
||||
|
||||
|
||||
def eval_e2e(args):
|
||||
# gt
|
||||
gt_results = parse_ser_results_fp(args.gt_json_path, "gt", args.ignore_background)
|
||||
# pred
|
||||
dt_results = parse_ser_results_fp(
|
||||
args.pred_json_path, "pred", args.ignore_background
|
||||
)
|
||||
iou_thresh = args.iou_thres
|
||||
num_gt_chars = 0
|
||||
gt_count = 0
|
||||
dt_count = 0
|
||||
hit = 0
|
||||
ed_sum = 0
|
||||
|
||||
for img_name in dt_results:
|
||||
gt_info = gt_results[img_name]
|
||||
gt_count += len(gt_info)
|
||||
|
||||
dt_info = dt_results[img_name]
|
||||
dt_count += len(dt_info)
|
||||
|
||||
dt_match = [False] * len(dt_info)
|
||||
gt_match = [False] * len(gt_info)
|
||||
|
||||
all_ious = defaultdict(tuple)
|
||||
# gt: {text, label, bbox or poly}
|
||||
for index_gt, gt in enumerate(gt_info):
|
||||
if "poly" not in gt:
|
||||
gt["poly"] = convert_bbox_to_polygon(gt["bbox"])
|
||||
gt_poly = polygon_from_str(gt["poly"])
|
||||
for index_dt, dt in enumerate(dt_info):
|
||||
if "poly" not in dt:
|
||||
dt["poly"] = convert_bbox_to_polygon(dt["bbox"])
|
||||
dt_poly = polygon_from_str(dt["poly"])
|
||||
iou = polygon_iou(dt_poly, gt_poly)
|
||||
if iou >= iou_thresh:
|
||||
all_ious[(index_gt, index_dt)] = iou
|
||||
sorted_ious = sorted(all_ious.items(), key=operator.itemgetter(1), reverse=True)
|
||||
sorted_gt_dt_pairs = [item[0] for item in sorted_ious]
|
||||
|
||||
# matched gt and dt
|
||||
for gt_dt_pair in sorted_gt_dt_pairs:
|
||||
index_gt, index_dt = gt_dt_pair
|
||||
if gt_match[index_gt] == False and dt_match[index_dt] == False:
|
||||
gt_match[index_gt] = True
|
||||
dt_match[index_dt] = True
|
||||
# ocr rec results
|
||||
gt_text = gt_info[index_gt]["text"]
|
||||
dt_text = dt_info[index_dt]["text"]
|
||||
|
||||
# ser results
|
||||
gt_label = gt_info[index_gt]["label"]
|
||||
dt_label = dt_info[index_dt]["pred"]
|
||||
|
||||
if True: # ignore_masks[index_gt] == '0':
|
||||
ed_sum += ed(args, gt_text, dt_text)
|
||||
num_gt_chars += len(gt_text)
|
||||
if gt_text == dt_text:
|
||||
if args.ignore_ser_prediction or gt_label == dt_label:
|
||||
hit += 1
|
||||
|
||||
# unmatched dt
|
||||
for tindex, dt_match_flag in enumerate(dt_match):
|
||||
if dt_match_flag == False:
|
||||
dt_text = dt_info[tindex]["text"]
|
||||
gt_text = ""
|
||||
ed_sum += ed(args, dt_text, gt_text)
|
||||
|
||||
# unmatched gt
|
||||
for tindex, gt_match_flag in enumerate(gt_match):
|
||||
if gt_match_flag == False:
|
||||
dt_text = ""
|
||||
gt_text = gt_info[tindex]["text"]
|
||||
ed_sum += ed(args, gt_text, dt_text)
|
||||
num_gt_chars += len(gt_text)
|
||||
|
||||
eps = 1e-9
|
||||
print("config: ", args)
|
||||
print("hit, dt_count, gt_count", hit, dt_count, gt_count)
|
||||
precision = hit / (dt_count + eps)
|
||||
recall = hit / (gt_count + eps)
|
||||
fmeasure = 2.0 * precision * recall / (precision + recall + eps)
|
||||
avg_edit_dist_img = ed_sum / len(gt_results)
|
||||
avg_edit_dist_field = ed_sum / (gt_count + eps)
|
||||
character_acc = 1 - ed_sum / (num_gt_chars + eps)
|
||||
|
||||
print("character_acc: %.2f" % (character_acc * 100) + "%")
|
||||
print("avg_edit_dist_field: %.2f" % (avg_edit_dist_field))
|
||||
print("avg_edit_dist_img: %.2f" % (avg_edit_dist_img))
|
||||
print("precision: %.2f" % (precision * 100) + "%")
|
||||
print("recall: %.2f" % (recall * 100) + "%")
|
||||
print("fmeasure: %.2f" % (fmeasure * 100) + "%")
|
||||
|
||||
return
|
||||
|
||||
|
||||
def parse_args():
|
||||
""" """
|
||||
|
||||
def str2bool(v):
|
||||
return v.lower() in ("true", "t", "1")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
## Required parameters
|
||||
parser.add_argument(
|
||||
"--gt_json_path",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pred_json_path",
|
||||
default=None,
|
||||
type=str,
|
||||
required=True,
|
||||
)
|
||||
|
||||
parser.add_argument("--iou_thres", default=0.5, type=float)
|
||||
|
||||
parser.add_argument(
|
||||
"--ignore_case",
|
||||
default=False,
|
||||
type=str2bool,
|
||||
help="whether to do lower case for the strs",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--ignore_space", default=True, type=str2bool, help="whether to ignore space"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--ignore_background",
|
||||
default=True,
|
||||
type=str2bool,
|
||||
help="whether to ignore other label",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--ignore_ser_prediction",
|
||||
default=False,
|
||||
type=str2bool,
|
||||
help="whether to ignore ocr pred results",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
eval_e2e(args)
|
||||
166
ppstructure/kie/tools/trans_funsd_label.py
Normal file
166
ppstructure/kie/tools/trans_funsd_label.py
Normal file
@@ -0,0 +1,166 @@
|
||||
# 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 json
|
||||
import os
|
||||
import sys
|
||||
import cv2
|
||||
import numpy as np
|
||||
from copy import deepcopy
|
||||
|
||||
|
||||
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 get_outer_poly(bbox_list):
|
||||
x1 = min([bbox[0] for bbox in bbox_list])
|
||||
y1 = min([bbox[1] for bbox in bbox_list])
|
||||
x2 = max([bbox[2] for bbox in bbox_list])
|
||||
y2 = max([bbox[3] for bbox in bbox_list])
|
||||
return [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
|
||||
|
||||
|
||||
def load_funsd_label(image_dir, anno_dir):
|
||||
imgs = os.listdir(image_dir)
|
||||
annos = os.listdir(anno_dir)
|
||||
|
||||
imgs = [img.replace(".png", "") for img in imgs]
|
||||
annos = [anno.replace(".json", "") for anno in annos]
|
||||
|
||||
fn_info_map = dict()
|
||||
for anno_fn in annos:
|
||||
res = []
|
||||
with open(os.path.join(anno_dir, anno_fn + ".json"), "r") as fin:
|
||||
infos = json.load(fin)
|
||||
infos = infos["form"]
|
||||
old_id2new_id_map = dict()
|
||||
global_new_id = 0
|
||||
for info in infos:
|
||||
if info["text"] is None:
|
||||
continue
|
||||
words = info["words"]
|
||||
if len(words) <= 0:
|
||||
continue
|
||||
word_idx = 1
|
||||
curr_bboxes = [words[0]["box"]]
|
||||
curr_texts = [words[0]["text"]]
|
||||
while word_idx < len(words):
|
||||
# switch to a new link
|
||||
if words[word_idx]["box"][0] + 10 <= words[word_idx - 1]["box"][2]:
|
||||
if len("".join(curr_texts[0])) > 0:
|
||||
res.append(
|
||||
{
|
||||
"transcription": " ".join(curr_texts),
|
||||
"label": info["label"],
|
||||
"points": get_outer_poly(curr_bboxes),
|
||||
"linking": info["linking"],
|
||||
"id": global_new_id,
|
||||
}
|
||||
)
|
||||
if info["id"] not in old_id2new_id_map:
|
||||
old_id2new_id_map[info["id"]] = []
|
||||
old_id2new_id_map[info["id"]].append(global_new_id)
|
||||
global_new_id += 1
|
||||
curr_bboxes = [words[word_idx]["box"]]
|
||||
curr_texts = [words[word_idx]["text"]]
|
||||
else:
|
||||
curr_bboxes.append(words[word_idx]["box"])
|
||||
curr_texts.append(words[word_idx]["text"])
|
||||
word_idx += 1
|
||||
if len("".join(curr_texts[0])) > 0:
|
||||
res.append(
|
||||
{
|
||||
"transcription": " ".join(curr_texts),
|
||||
"label": info["label"],
|
||||
"points": get_outer_poly(curr_bboxes),
|
||||
"linking": info["linking"],
|
||||
"id": global_new_id,
|
||||
}
|
||||
)
|
||||
if info["id"] not in old_id2new_id_map:
|
||||
old_id2new_id_map[info["id"]] = []
|
||||
old_id2new_id_map[info["id"]].append(global_new_id)
|
||||
global_new_id += 1
|
||||
res = sorted(res, key=lambda r: (r["points"][0][1], r["points"][0][0]))
|
||||
for i in range(len(res) - 1):
|
||||
for j in range(i, 0, -1):
|
||||
if abs(
|
||||
res[j + 1]["points"][0][1] - res[j]["points"][0][1]
|
||||
) < 20 and (res[j + 1]["points"][0][0] < res[j]["points"][0][0]):
|
||||
tmp = deepcopy(res[j])
|
||||
res[j] = deepcopy(res[j + 1])
|
||||
res[j + 1] = deepcopy(tmp)
|
||||
else:
|
||||
break
|
||||
# re-generate unique ids
|
||||
for idx, r in enumerate(res):
|
||||
new_links = []
|
||||
for link in r["linking"]:
|
||||
# illegal links will be removed
|
||||
if (
|
||||
link[0] not in old_id2new_id_map
|
||||
or link[1] not in old_id2new_id_map
|
||||
):
|
||||
continue
|
||||
for src in old_id2new_id_map[link[0]]:
|
||||
for dst in old_id2new_id_map[link[1]]:
|
||||
new_links.append([src, dst])
|
||||
res[idx]["linking"] = deepcopy(new_links)
|
||||
|
||||
fn_info_map[anno_fn] = res
|
||||
|
||||
return fn_info_map
|
||||
|
||||
|
||||
def main():
|
||||
test_image_dir = "train_data/FUNSD/testing_data/images/"
|
||||
test_anno_dir = "train_data/FUNSD/testing_data/annotations/"
|
||||
test_output_dir = "train_data/FUNSD/test.json"
|
||||
|
||||
fn_info_map = load_funsd_label(test_image_dir, test_anno_dir)
|
||||
with open(test_output_dir, "w") as fout:
|
||||
for fn in fn_info_map:
|
||||
fout.write(
|
||||
fn
|
||||
+ ".png"
|
||||
+ "\t"
|
||||
+ json.dumps(fn_info_map[fn], ensure_ascii=False)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
train_image_dir = "train_data/FUNSD/training_data/images/"
|
||||
train_anno_dir = "train_data/FUNSD/training_data/annotations/"
|
||||
train_output_dir = "train_data/FUNSD/train.json"
|
||||
|
||||
fn_info_map = load_funsd_label(train_image_dir, train_anno_dir)
|
||||
with open(train_output_dir, "w") as fout:
|
||||
for fn in fn_info_map:
|
||||
fout.write(
|
||||
fn
|
||||
+ ".png"
|
||||
+ "\t"
|
||||
+ json.dumps(fn_info_map[fn], ensure_ascii=False)
|
||||
+ "\n"
|
||||
)
|
||||
print("====ok====")
|
||||
return
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
64
ppstructure/kie/tools/trans_xfun_data.py
Normal file
64
ppstructure/kie/tools/trans_xfun_data.py
Normal file
@@ -0,0 +1,64 @@
|
||||
# Copyright (c) 2021 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 json
|
||||
|
||||
|
||||
def transfer_xfun_data(json_path=None, output_file=None):
|
||||
with open(json_path, "r", encoding="utf-8") as fin:
|
||||
lines = fin.readlines()
|
||||
|
||||
json_info = json.loads(lines[0])
|
||||
documents = json_info["documents"]
|
||||
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"]
|
||||
|
||||
for doc in document:
|
||||
x1, y1, x2, y2 = doc["box"]
|
||||
points = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
|
||||
label_info.append(
|
||||
{
|
||||
"transcription": doc["text"],
|
||||
"label": doc["label"],
|
||||
"points": points,
|
||||
"id": doc["id"],
|
||||
"linking": doc["linking"],
|
||||
}
|
||||
)
|
||||
|
||||
fout.write(
|
||||
image_path + "\t" + json.dumps(label_info, ensure_ascii=False) + "\n"
|
||||
)
|
||||
|
||||
print("===ok====")
|
||||
|
||||
|
||||
def parser_args():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="args for paddleserving")
|
||||
parser.add_argument(
|
||||
"--ori_gt_path", type=str, required=True, help="origin xfun gt path"
|
||||
)
|
||||
parser.add_argument("--output_path", type=str, required=True, help="path to save")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
args = parser_args()
|
||||
transfer_xfun_data(args.ori_gt_path, args.output_path)
|
||||
470
ppstructure/layout/README.md
Normal file
470
ppstructure/layout/README.md
Normal file
@@ -0,0 +1,470 @@
|
||||
English | [简体中文](README_ch.md)
|
||||
|
||||
# Layout analysis
|
||||
|
||||
- [1. Introduction](#1-Introduction)
|
||||
- [2. Quick start](#2-Quick-start)
|
||||
- [3. Install](#3-Install)
|
||||
- [3.1 Install PaddlePaddle](#31-Install-paddlepaddle)
|
||||
- [3.2 Install PaddleDetection](#32-Install-paddledetection)
|
||||
- [4. Data preparation](#4-Data-preparation)
|
||||
- [4.1 English data set](#41-English-data-set)
|
||||
- [4.2 More datasets](#42-More-datasets)
|
||||
- [5. Start training](#5-Start-training)
|
||||
- [5.1 Train](#51-Train)
|
||||
- [5.2 FGD Distillation training](#52-Fgd-distillation-training)
|
||||
- [6. Model evaluation and prediction](#6-Model-evaluation-and-prediction)
|
||||
- [6.1 Indicator evaluation](#61-Indicator-evaluation)
|
||||
- [6.2 Test layout analysis results](#62-Test-layout-analysis-results)
|
||||
- [7. Model export and inference](#7-Model-export-and-inference)
|
||||
- [7.1 Model export](#71-Model-export)
|
||||
- [7.2 Model inference](#72-Model-inference)
|
||||
|
||||
|
||||
## 1. Introduction
|
||||
|
||||
Layout analysis refers to the regional division of documents in the form of pictures and the positioning of key areas, such as text, title, table, picture, etc. The layout analysis algorithm is based on the lightweight model PP-picodet of [PaddleDetection]( https://github.com/PaddlePaddle/PaddleDetection ), including English layout analysis, Chinese layout analysis and table layout analysis models. English layout analysis models can detect document layout elements such as text, title, table, figure, list. Chinese layout analysis models can detect document layout elements such as text, figure, figure caption, table, table caption, header, footer, reference, and equation. Table layout analysis models can detect table regions.
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/layout/layout.png" width="800">
|
||||
</div>
|
||||
|
||||
## 2. Quick start
|
||||
PP-Structure currently provides layout analysis models in Chinese, English and table documents. For the model link, see [models_list](../docs/models_list_en.md). The whl package is also provided for quick use, see [quickstart](../docs/quickstart_en.md) for details.
|
||||
|
||||
## 3. Install
|
||||
|
||||
### 3.1. Install PaddlePaddle
|
||||
|
||||
- **（1) Install PaddlePaddle**
|
||||
|
||||
```bash
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
# GPU Install
|
||||
python3 -m pip install "paddlepaddle-gpu>=2.3" -i https://mirror.baidu.com/pypi/simple
|
||||
|
||||
# CPU Install
|
||||
python3 -m pip install "paddlepaddle>=2.3" -i https://mirror.baidu.com/pypi/simple
|
||||
```
|
||||
For more requirements, please refer to the instructions in the [Install file](https://www.paddlepaddle.org.cn/install/quick)。
|
||||
|
||||
### 3.2. Install PaddleDetection
|
||||
|
||||
- **（1）Download PaddleDetection Source code**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/PaddlePaddle/PaddleDetection.git
|
||||
```
|
||||
|
||||
- **（2）Install third-party libraries**
|
||||
|
||||
```bash
|
||||
cd PaddleDetection
|
||||
python3 -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 4. Data preparation
|
||||
|
||||
If you want to experience the prediction process directly, you can skip data preparation and download the pre-training model.
|
||||
|
||||
### 4.1. English data set
|
||||
|
||||
Download document analysis data set [PubLayNet](https://developer.ibm.com/exchanges/data/all/publaynet/)（Dataset 96G），contains 5 classes：`{0: "Text", 1: "Title", 2: "List", 3:"Table", 4:"Figure"}`
|
||||
|
||||
```
|
||||
# Download data
|
||||
wget https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz
|
||||
# Decompress data
|
||||
tar -xvf publaynet.tar.gz
|
||||
```
|
||||
|
||||
Uncompressed **directory structure：**
|
||||
|
||||
```
|
||||
|-publaynet
|
||||
|- test
|
||||
|- PMC1277013_00004.jpg
|
||||
|- PMC1291385_00002.jpg
|
||||
| ...
|
||||
|- train.json
|
||||
|- train
|
||||
|- PMC1291385_00002.jpg
|
||||
|- PMC1277013_00004.jpg
|
||||
| ...
|
||||
|- val.json
|
||||
|- val
|
||||
|- PMC538274_00004.jpg
|
||||
|- PMC539300_00004.jpg
|
||||
| ...
|
||||
```
|
||||
|
||||
**data distribution：**
|
||||
|
||||
| File or Folder | Description | num |
|
||||
| :------------- | :------------- | ------- |
|
||||
| `train/` | Training set pictures | 335,703 |
|
||||
| `val/` | Verification set pictures | 11,245 |
|
||||
| `test/` | Test set pictures | 11,405 |
|
||||
| `train.json` | Training set annotation files | - |
|
||||
| `val.json` | Validation set dimension files | - |
|
||||
|
||||
**Data Annotation**
|
||||
|
||||
The JSON file contains the annotations of all images, and the data is stored in a dictionary nested manner.Contains the following keys：
|
||||
|
||||
- info，represents the dimension file info。
|
||||
|
||||
- licenses，represents the dimension file licenses。
|
||||
|
||||
- images，represents the list of image information in the annotation file，each element is the information of an image。The information of one of the images is as follows:
|
||||
|
||||
```
|
||||
{
|
||||
'file_name': 'PMC4055390_00006.jpg', # file_name
|
||||
'height': 601, # image height
|
||||
'width': 792, # image width
|
||||
'id': 341427 # image id
|
||||
}
|
||||
```
|
||||
|
||||
- annotations， represents the list of annotation information of the target object in the annotation file，each element is the annotation information of a target object。The following is the annotation information of one of the target objects:
|
||||
|
||||
```
|
||||
{
|
||||
|
||||
'segmentation': # Segmentation annotation of objects
|
||||
'area': 60518.099043117836, # Area of object
|
||||
'iscrowd': 0, # iscrowd
|
||||
'image_id': 341427, # image id
|
||||
'bbox': [50.58, 490.86, 240.15, 252.16], # bbox [x1,y1,w,h]
|
||||
'category_id': 1, # category_id
|
||||
'id': 3322348 # image id
|
||||
}
|
||||
```
|
||||
|
||||
### 4.2. More datasets
|
||||
|
||||
We provide CDLA(Chinese layout analysis), TableBank(Table layout analysis)etc. data set download links，process to the JSON format of the above annotation file，that is, the training can be conducted in the same way。
|
||||
|
||||
| dataset | 简介 |
|
||||
| ------------------------------------------------------------ | ------------------------------------------------------------ |
|
||||
| [cTDaR2019_cTDaR](https://cndplab-founder.github.io/cTDaR2019/) | For form detection (TRACKA) and form identification (TRACKB).Image types include historical data sets (beginning with cTDaR_t0, such as CTDAR_T00872.jpg) and modern data sets (beginning with cTDaR_t1, CTDAR_T10482.jpg). |
|
||||
| [IIIT-AR-13K](http://cvit.iiit.ac.in/usodi/iiitar13k.php) | Data sets constructed by manually annotating figures or pages from publicly available annual reports, containing 5 categories:table, figure, natural image, logo, and signature. |
|
||||
| [TableBank](https://github.com/doc-analysis/TableBank) | For table detection and recognition of large datasets, including Word and Latex document formats |
|
||||
| [CDLA](https://github.com/buptlihang/CDLA) | Chinese document layout analysis data set, for Chinese literature (paper) scenarios, including 10 categories:Text, Title, Figure, Figure caption, Table, Table caption, Header, Footer, Reference, Equation |
|
||||
| [DocBank](https://github.com/doc-analysis/DocBank) | Large-scale dataset (500K document pages) constructed using weakly supervised methods for document layout analysis, containing 12 categories:Author, Caption, Date, Equation, Figure, Footer, List, Paragraph, Reference, Section, Table, Title |
|
||||
|
||||
|
||||
## 5. Start training
|
||||
|
||||
Training scripts, evaluation scripts, and prediction scripts are provided, and the PubLayNet pre-training model is used as an example in this section.
|
||||
|
||||
If you do not want training and directly experience the following process of model evaluation, prediction, motion to static, and inference, you can download the provided pre-trained model (PubLayNet dataset) and skip this part.
|
||||
|
||||
```
|
||||
mkdir pretrained_model
|
||||
cd pretrained_model
|
||||
# Download PubLayNet pre-training model（Direct experience model evaluates, predicts, and turns static）
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout.pdparams
|
||||
# Download the PubLaynet inference model（Direct experience model reasoning）
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
```
|
||||
|
||||
If the test image is Chinese, the pre-trained model of Chinese CDLA dataset can be downloaded to identify 10 types of document regions：Table, Figure, Figure caption, Table, Table caption, Header, Footer, Reference, Equation，Download the training model and inference model of Model 'picodet_lcnet_x1_0_fgd_layout_cdla' in [layout analysis model](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md)。If only the table area in the image is detected, you can download the pre-trained model of the table dataset, and download the training model and inference model of the 'picodet_LCnet_x1_0_FGd_layout_table' model in [Layout Analysis model](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md)
|
||||
|
||||
### 5.1. Train
|
||||
|
||||
Start training with the PaddleDetection [layout analysis profile](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.5/configs/picodet/legacy_model/application/layout_analysis)
|
||||
|
||||
* Modify Profile
|
||||
|
||||
If you want to train your own data set, you need to modify the data configuration and the number of categories in the configuration file.
|
||||
|
||||
|
||||
Using 'configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml' as an example, the change is as follows:
|
||||
|
||||
```yaml
|
||||
metric: COCO
|
||||
# Number of categories
|
||||
num_classes: 5
|
||||
|
||||
TrainDataset:
|
||||
!COCODataSet
|
||||
# Modify to your own training data directory
|
||||
image_dir: train
|
||||
# Modify to your own training data label file
|
||||
anno_path: train.json
|
||||
# Modify to your own training data root directory
|
||||
dataset_dir: /root/publaynet/
|
||||
data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd']
|
||||
|
||||
EvalDataset:
|
||||
!COCODataSet
|
||||
# Modify to your own validation data directory
|
||||
image_dir: val
|
||||
# Modify to your own validation data label file
|
||||
anno_path: val.json
|
||||
# Modify to your own validation data root
|
||||
dataset_dir: /root/publaynet/
|
||||
|
||||
TestDataset:
|
||||
!ImageFolder
|
||||
# Modify to your own test data label file
|
||||
anno_path: /root/publaynet/val.json
|
||||
```
|
||||
|
||||
* Start training. During training, PP picodet pre training model will be downloaded by default. There is no need to download in advance.
|
||||
|
||||
```bash
|
||||
# GPU training supports single-card and multi-card training
|
||||
# The training log is automatically saved to the log directory
|
||||
|
||||
# Single card training
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
python3 tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--eval
|
||||
|
||||
# Multi-card training, with the -- GPUS parameter specifying the card number
|
||||
export CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||
python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--eval
|
||||
```
|
||||
|
||||
**Attention：**If the video memory is out during training, adjust Batch_size in TrainReader and base_LR in LearningRate. The published config is obtained by 8-card training. If the number of GPU cards is changed to 1, then the base_LR needs to be reduced by 8 times.
|
||||
|
||||
After starting training normally, you will see the following log output:
|
||||
|
||||
```
|
||||
[08/15 04:02:30] ppdet.utils.checkpoint INFO: Finish loading model weights: /root/.cache/paddle/weights/LCNet_x1_0_pretrained.pdparams
|
||||
[08/15 04:02:46] ppdet.engine INFO: Epoch: [0] [ 0/1929] learning_rate: 0.040000 loss_vfl: 1.216707 loss_bbox: 1.142163 loss_dfl: 0.544196 loss: 2.903065 eta: 17 days, 13:50:26 batch_cost: 15.7452 data_cost: 2.9112 ips: 1.5243 images/s
|
||||
[08/15 04:03:19] ppdet.engine INFO: Epoch: [0] [ 20/1929] learning_rate: 0.064000 loss_vfl: 1.180627 loss_bbox: 0.939552 loss_dfl: 0.442436 loss: 2.628206 eta: 2 days, 12:18:53 batch_cost: 1.5770 data_cost: 0.0008 ips: 15.2184 images/s
|
||||
[08/15 04:03:47] ppdet.engine INFO: Epoch: [0] [ 40/1929] learning_rate: 0.088000 loss_vfl: 0.543321 loss_bbox: 1.071401 loss_dfl: 0.457817 loss: 2.057003 eta: 2 days, 0:07:03 batch_cost: 1.3190 data_cost: 0.0007 ips: 18.1954 images/s
|
||||
[08/15 04:04:12] ppdet.engine INFO: Epoch: [0] [ 60/1929] learning_rate: 0.112000 loss_vfl: 0.630989 loss_bbox: 0.859183 loss_dfl: 0.384702 loss: 1.883143 eta: 1 day, 19:01:29 batch_cost: 1.2177 data_cost: 0.0006 ips: 19.7087 images/s
|
||||
```
|
||||
|
||||
- `--eval` indicates that the best model is saved as `output/picodet_lcnet_x1_0_layout/best_accuracy` by default during the evaluation process 。
|
||||
|
||||
**Note that the configuration file for prediction / evaluation must be consistent with the training.**
|
||||
|
||||
### 5.2. FGD Distillation Training
|
||||
|
||||
PaddleDetection supports FGD-based [Focal and Global Knowledge Distillation for Detectors]( https://arxiv.org/abs/2111.11837v1) The training process of the target detection model of distillation, FGD distillation is divided into two parts `Focal` and `Global`. `Focal` Distillation separates the foreground and background of the image, allowing the student model to focus on the key pixels of the foreground and background features of the teacher model respectively;` Global`Distillation section reconstructs the relationships between different pixels and transfers them from the teacher to the student to compensate for the global information lost in `Focal`Distillation.
|
||||
|
||||
Change the dataset and modify the data configuration and number of categories in the [TODO] configuration, referring to 4.1. Start training:
|
||||
|
||||
```bash
|
||||
# Single Card Training
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
python3 tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
--eval
|
||||
```
|
||||
|
||||
- `-c`: Specify the model configuration file.
|
||||
- `--slim_config`: Specify the compression policy profile.
|
||||
|
||||
## 6. Model evaluation and prediction
|
||||
|
||||
### 6.1. Indicator evaluation
|
||||
|
||||
Model parameters in training are saved by default in `output/picodet_ Lcnet_ X1_ 0_ Under the layout` directory. When evaluating indicators, you need to set `weights` to point to the saved parameter file.Assessment datasets can be accessed via `configs/picodet/legacy_ Model/application/layout_ Analysis/picodet_ Lcnet_ X1_ 0_ Layout. Yml` . Modify `EvalDataset` : `img_dir`,`anno_ Path`and`dataset_dir` setting.
|
||||
|
||||
```bash
|
||||
# GPU evaluation, weights as weights to be measured
|
||||
python3 tools/eval.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights=./output/picodet_lcnet_x1_0_layout/best_model
|
||||
```
|
||||
|
||||
The following information will be printed out, such as mAP, AP0.5, etc.
|
||||
|
||||
```py
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.935
|
||||
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.979
|
||||
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.956
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.404
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.782
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.969
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.539
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.938
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.949
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.495
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.818
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.978
|
||||
[08/15 07:07:09] ppdet.engine INFO: Total sample number: 11245, averge FPS: 24.405059207157436
|
||||
[08/15 07:07:09] ppdet.engine INFO: Best test bbox ap is 0.935.
|
||||
```
|
||||
|
||||
If you use the provided pre-training model for evaluation or the FGD distillation training model, replace the `weights` model path and execute the following command for evaluation:
|
||||
|
||||
```
|
||||
python3 tools/eval.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights=output/picodet_lcnet_x2_5_layout/best_model
|
||||
```
|
||||
|
||||
- `-c`: Specify the model configuration file.
|
||||
- `--slim_config`: Specify the distillation policy profile.
|
||||
- `-o weights`: Specify the model path trained by the distillation algorithm.
|
||||
|
||||
### 6.2. Test Layout Analysis Results
|
||||
|
||||
|
||||
The profile predicted to be used must be consistent with the training, for example, if you pass `python3 tools/train'. Py-c configs/picodet/legacy_ Model/application/layout_ Analysis/picodet_ Lcnet_ X1_ 0_ Layout. Yml` completed the training process for the model.
|
||||
|
||||
With trained PaddleDetection model, you can use the following commands to make model predictions.
|
||||
|
||||
```bash
|
||||
python3 tools/infer.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights='output/picodet_lcnet_x1_0_layout/best_model.pdparams' \
|
||||
--infer_img='docs/images/layout.jpg' \
|
||||
--output_dir=output_dir/ \
|
||||
--draw_threshold=0.5
|
||||
```
|
||||
|
||||
- `--infer_img`: Reasoning for a single picture can also be done via `--infer_ Dir`Inform all pictures in the file.
|
||||
- `--output_dir`: Specify the path to save the visualization results.
|
||||
- `--draw_threshold`:Specify the NMS threshold for drawing the result box.
|
||||
|
||||
If you use the provided pre-training model for prediction or the FGD distillation training model, change the `weights` model path and execute the following command to make the prediction:
|
||||
|
||||
```
|
||||
python3 tools/infer.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights='output/picodet_lcnet_x2_5_layout/best_model.pdparams' \
|
||||
--infer_img='docs/images/layout.jpg' \
|
||||
--output_dir=output_dir/ \
|
||||
--draw_threshold=0.5
|
||||
```
|
||||
|
||||
|
||||
## 7. Model Export and Inference
|
||||
|
||||
|
||||
### 7.1 Model Export
|
||||
|
||||
The inference model (the model saved by `paddle.jit.save`) is generally a solidified model saved after the model training is completed, and is mostly used to give prediction in deployment.
|
||||
|
||||
The model saved during the training process is the checkpoints model, which saves the parameters of the model and is mostly used to resume training.
|
||||
|
||||
Compared with the checkpoints model, the inference model will additionally save the structural information of the model. Therefore, it is easier to deploy because the model structure and model parameters are already solidified in the inference model file, and is suitable for integration with actual systems.
|
||||
|
||||
Layout analysis model to inference model steps are as follows：
|
||||
|
||||
```bash
|
||||
python3 tools/export_model.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights=output/picodet_lcnet_x1_0_layout/best_model \
|
||||
--output_dir=output_inference/
|
||||
```
|
||||
|
||||
* If no post-export processing is required, specify：`-o export.benchmark=True`（If -o already exists, delete -o here）
|
||||
* If you do not need to export NMS, specify：`-o export.nms=False`
|
||||
|
||||
After successful conversion, there are three files in the directory:
|
||||
|
||||
```
|
||||
output_inference/picodet_lcnet_x1_0_layout/
|
||||
├── model.pdiparams # inference Parameter file for model
|
||||
├── model.pdiparams.info # inference Model parameter information, ignorable
|
||||
└── model.pdmodel # inference Model Structure File for Model
|
||||
```
|
||||
|
||||
If you change the `weights` model path using the provided pre-training model to the Inference model, or using the FGD distillation training model, the model to inference model steps are as follows:
|
||||
|
||||
```bash
|
||||
python3 tools/export_model.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights=./output/picodet_lcnet_x2_5_layout/best_model \
|
||||
--output_dir=output_inference/
|
||||
```
|
||||
|
||||
### 7.2 Model inference
|
||||
|
||||
Replace model_with the provided inference training model for inference or the FGD distillation training `model_dir`Inference model path, execute the following commands for inference:
|
||||
|
||||
```bash
|
||||
python3 deploy/python/infer.py \
|
||||
--model_dir=output_inference/picodet_lcnet_x1_0_layout/ \
|
||||
--image_file=docs/images/layout.jpg \
|
||||
--device=CPU
|
||||
```
|
||||
|
||||
- --device：Specify the GPU or CPU device
|
||||
|
||||
When model inference is complete, you will see the following log output:
|
||||
|
||||
```
|
||||
------------------------------------------
|
||||
----------- Model Configuration -----------
|
||||
Model Arch: PicoDet
|
||||
Transform Order:
|
||||
--transform op: Resize
|
||||
--transform op: NormalizeImage
|
||||
--transform op: Permute
|
||||
--transform op: PadStride
|
||||
--------------------------------------------
|
||||
class_id:0, confidence:0.9921, left_top:[20.18,35.66],right_bottom:[341.58,600.99]
|
||||
class_id:0, confidence:0.9914, left_top:[19.77,611.42],right_bottom:[341.48,901.82]
|
||||
class_id:0, confidence:0.9904, left_top:[369.36,375.10],right_bottom:[691.29,600.59]
|
||||
class_id:0, confidence:0.9835, left_top:[369.60,608.60],right_bottom:[691.38,736.72]
|
||||
class_id:0, confidence:0.9830, left_top:[369.58,805.38],right_bottom:[690.97,901.80]
|
||||
class_id:0, confidence:0.9716, left_top:[383.68,271.44],right_bottom:[688.93,335.39]
|
||||
class_id:0, confidence:0.9452, left_top:[370.82,34.48],right_bottom:[688.10,63.54]
|
||||
class_id:1, confidence:0.8712, left_top:[370.84,771.03],right_bottom:[519.30,789.13]
|
||||
class_id:3, confidence:0.9856, left_top:[371.28,67.85],right_bottom:[685.73,267.72]
|
||||
save result to: output/layout.jpg
|
||||
Test iter 0
|
||||
------------------ Inference Time Info ----------------------
|
||||
total_time(ms): 2196.0, img_num: 1
|
||||
average latency time(ms): 2196.00, QPS: 0.455373
|
||||
preprocess_time(ms): 2172.50, inference_time(ms): 11.90, postprocess_time(ms): 11.60
|
||||
```
|
||||
|
||||
- Model：model structure
|
||||
- Transform Order：Preprocessing operation
|
||||
- class_id, confidence, left_top, right_bottom：Indicates category id, confidence level, upper left coordinate, lower right coordinate, respectively
|
||||
- save result to：Save path of visual layout analysis results, default save to ./output folder
|
||||
- inference time info：Inference time, where preprocess_time represents the preprocessing time, Inference_time represents the model prediction time, and postprocess_time represents the post-processing time
|
||||
|
||||
The result of visualization layout is shown in the following figure
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/layout/layout_res.jpg" width="800">
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
## Citations
|
||||
|
||||
```
|
||||
@inproceedings{zhong2019publaynet,
|
||||
title={PubLayNet: largest dataset ever for document layout analysis},
|
||||
author={Zhong, Xu and Tang, Jianbin and Yepes, Antonio Jimeno},
|
||||
booktitle={2019 International Conference on Document Analysis and Recognition (ICDAR)},
|
||||
year={2019},
|
||||
volume={},
|
||||
number={},
|
||||
pages={1015-1022},
|
||||
doi={10.1109/ICDAR.2019.00166},
|
||||
ISSN={1520-5363},
|
||||
month={Sep.},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
@inproceedings{yang2022focal,
|
||||
title={Focal and global knowledge distillation for detectors},
|
||||
author={Yang, Zhendong and Li, Zhe and Jiang, Xiaohu and Gong, Yuan and Yuan, Zehuan and Zhao, Danpei and Yuan, Chun},
|
||||
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
|
||||
pages={4643--4652},
|
||||
year={2022}
|
||||
}
|
||||
```
|
||||
469
ppstructure/layout/README_ch.md
Normal file
469
ppstructure/layout/README_ch.md
Normal file
@@ -0,0 +1,469 @@
|
||||
简体中文 | [English](README.md)
|
||||
|
||||
# 版面分析
|
||||
|
||||
- [1. 简介](#1-简介)
|
||||
- [2. 快速开始](#2-快速开始)
|
||||
- [3. 安装](#3-安装)
|
||||
- [3.1 安装PaddlePaddle](#31-安装paddlepaddle)
|
||||
- [3.2 安装PaddleDetection](#32-安装paddledetection)
|
||||
- [4. 数据准备](#4-数据准备)
|
||||
- [4.1 英文数据集](#41-英文数据集)
|
||||
- [4.2 更多数据集](#42-更多数据集)
|
||||
- [5. 开始训练](#5-开始训练)
|
||||
- [5.1 启动训练](#51-启动训练)
|
||||
- [5.2 FGD蒸馏训练](#52-fgd蒸馏训练)
|
||||
- [6. 模型评估与预测](#6-模型评估与预测)
|
||||
- [6.1 指标评估](#61-指标评估)
|
||||
- [6.2 测试版面分析结果](#62-测试版面分析结果)
|
||||
- [7 模型导出与预测](#7-模型导出与预测)
|
||||
- [7.1 模型导出](#71-模型导出)
|
||||
- [7.2 模型推理](#72-模型推理)
|
||||
|
||||
## 1. 简介
|
||||
|
||||
版面分析指的是对图片形式的文档进行区域划分，定位其中的关键区域，如文字、标题、表格、图片等。版面分析算法基于[PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection)的轻量模型PP-PicoDet进行开发，包含英文、中文、表格版面分析3类模型。其中，英文模型支持Text、Title、Tale、Figure、List5类区域的检测，中文模型支持Text、Title、Figure、Figure caption、Table、Table caption、Header、Footer、Reference、Equation10类区域的检测，表格版面分析支持Table区域的检测，版面分析效果如下图所示：
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/layout/layout.png" width="800">
|
||||
</div>
|
||||
|
||||
## 2. 快速开始
|
||||
|
||||
PP-Structure目前提供了中文、英文、表格三类文档版面分析模型，模型链接见 [models_list](../docs/models_list.md#1-版面分析模型)。也提供了whl包的形式方便快速使用，详见 [quickstart](../docs/quickstart.md)。
|
||||
|
||||
|
||||
## 3. 安装
|
||||
|
||||
### 3.1. 安装PaddlePaddle
|
||||
|
||||
- **（1) 安装PaddlePaddle**
|
||||
|
||||
```bash
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
# GPU安装
|
||||
python3 -m pip install "paddlepaddle-gpu>=2.3" -i https://mirror.baidu.com/pypi/simple
|
||||
|
||||
# CPU安装
|
||||
python3 -m pip install "paddlepaddle>=2.3" -i https://mirror.baidu.com/pypi/simple
|
||||
```
|
||||
更多需求，请参照[安装文档](https://www.paddlepaddle.org.cn/install/quick)中的说明进行操作。
|
||||
|
||||
### 3.2. 安装PaddleDetection
|
||||
|
||||
- **（1）下载PaddleDetection源码**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/PaddlePaddle/PaddleDetection.git
|
||||
```
|
||||
|
||||
- **（2）安装其他依赖**
|
||||
|
||||
```bash
|
||||
cd PaddleDetection
|
||||
python3 -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 4. 数据准备
|
||||
|
||||
如果希望直接体验预测过程，可以跳过数据准备，下载我们提供的预训练模型。
|
||||
|
||||
### 4.1. 英文数据集
|
||||
|
||||
下载文档分析数据集[PubLayNet](https://developer.ibm.com/exchanges/data/all/publaynet/)（数据集96G），包含5个类：`{0: "Text", 1: "Title", 2: "List", 3:"Table", 4:"Figure"}`
|
||||
|
||||
```
|
||||
# 下载数据
|
||||
wget https://dax-cdn.cdn.appdomain.cloud/dax-publaynet/1.0.0/publaynet.tar.gz
|
||||
# 解压数据
|
||||
tar -xvf publaynet.tar.gz
|
||||
```
|
||||
|
||||
解压之后的**目录结构：**
|
||||
|
||||
```
|
||||
|-publaynet
|
||||
|- test
|
||||
|- PMC1277013_00004.jpg
|
||||
|- PMC1291385_00002.jpg
|
||||
| ...
|
||||
|- train.json
|
||||
|- train
|
||||
|- PMC1291385_00002.jpg
|
||||
|- PMC1277013_00004.jpg
|
||||
| ...
|
||||
|- val.json
|
||||
|- val
|
||||
|- PMC538274_00004.jpg
|
||||
|- PMC539300_00004.jpg
|
||||
| ...
|
||||
```
|
||||
|
||||
**数据分布：**
|
||||
|
||||
| File or Folder | Description | num |
|
||||
| :------------- | :------------- | ------- |
|
||||
| `train/` | 训练集图片 | 335,703 |
|
||||
| `val/` | 验证集图片 | 11,245 |
|
||||
| `test/` | 测试集图片 | 11,405 |
|
||||
| `train.json` | 训练集标注文件 | - |
|
||||
| `val.json` | 验证集标注文件 | - |
|
||||
|
||||
**标注格式：**
|
||||
|
||||
json文件包含所有图像的标注，数据以字典嵌套的方式存放，包含以下key：
|
||||
|
||||
- info，表示标注文件info。
|
||||
|
||||
- licenses，表示标注文件licenses。
|
||||
|
||||
- images，表示标注文件中图像信息列表，每个元素是一张图像的信息。如下为其中一张图像的信息：
|
||||
|
||||
```
|
||||
{
|
||||
'file_name': 'PMC4055390_00006.jpg', # file_name
|
||||
'height': 601, # image height
|
||||
'width': 792, # image width
|
||||
'id': 341427 # image id
|
||||
}
|
||||
```
|
||||
|
||||
- annotations，表示标注文件中目标物体的标注信息列表，每个元素是一个目标物体的标注信息。如下为其中一个目标物体的标注信息：
|
||||
|
||||
```
|
||||
{
|
||||
|
||||
'segmentation': # 物体的分割标注
|
||||
'area': 60518.099043117836, # 物体的区域面积
|
||||
'iscrowd': 0, # iscrowd
|
||||
'image_id': 341427, # image id
|
||||
'bbox': [50.58, 490.86, 240.15, 252.16], # bbox [x1,y1,w,h]
|
||||
'category_id': 1, # category_id
|
||||
'id': 3322348 # image id
|
||||
}
|
||||
```
|
||||
|
||||
### 4.2. 更多数据集
|
||||
|
||||
我们提供了CDLA(中文版面分析)、TableBank(表格版面分析)等数据集的下连接，处理为上述标注文件json格式，即可以按相同方式进行训练。
|
||||
|
||||
| dataset | 简介 |
|
||||
| ------------------------------------------------------------ | ------------------------------------------------------------ |
|
||||
| [cTDaR2019_cTDaR](https://cndplab-founder.github.io/cTDaR2019/) | 用于表格检测(TRACKA)和表格识别(TRACKB)。图片类型包含历史数据集(以cTDaR_t0开头，如cTDaR_t00872.jpg)和现代数据集(以cTDaR_t1开头，cTDaR_t10482.jpg)。 |
|
||||
| [IIIT-AR-13K](http://cvit.iiit.ac.in/usodi/iiitar13k.php) | 手动注释公开的年度报告中的图形或页面而构建的数据集，包含5类：table, figure, natural image, logo, and signature |
|
||||
| [CDLA](https://github.com/buptlihang/CDLA) | 中文文档版面分析数据集，面向中文文献类（论文）场景，包含10类：Text、Title、Figure、Figure caption、Table、Table caption、Header、Footer、Reference、Equation |
|
||||
| [TableBank](https://github.com/doc-analysis/TableBank) | 用于表格检测和识别大型数据集，包含Word和Latex2种文档格式 |
|
||||
| [DocBank](https://github.com/doc-analysis/DocBank) | 使用弱监督方法构建的大规模数据集(500K文档页面)，用于文档布局分析，包含12类：Author、Caption、Date、Equation、Figure、Footer、List、Paragraph、Reference、Section、Table、Title |
|
||||
|
||||
|
||||
## 5. 开始训练
|
||||
|
||||
提供了训练脚本、评估脚本和预测脚本，本节将以PubLayNet预训练模型为例进行讲解。
|
||||
|
||||
如果不希望训练，直接体验后面的模型评估、预测、动转静、推理的流程，可以下载提供的预训练模型(PubLayNet数据集)，并跳过5.1和5.2。
|
||||
|
||||
```
|
||||
mkdir pretrained_model
|
||||
cd pretrained_model
|
||||
# 下载PubLayNet预训练模型（直接体验模型评估、预测、动转静）
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout.pdparams
|
||||
# 下载PubLaynet推理模型（直接体验模型推理）
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
```
|
||||
|
||||
如果测试图片为中文，可以下载中文CDLA数据集的预训练模型，识别10类文档区域：Table、Figure、Figure caption、Table、Table caption、Header、Footer、Reference、Equation，在[版面分析模型](../docs/models_list.md)中下载`picodet_lcnet_x1_0_fgd_layout_cdla`模型的训练模型和推理模型。如果只检测图片中的表格区域，可以下载表格数据集的预训练模型，在[版面分析模型](../docs/models_list.md)中下载`picodet_lcnet_x1_0_fgd_layout_table`模型的训练模型和推理模型。
|
||||
|
||||
### 5.1. 启动训练
|
||||
|
||||
使用PaddleDetection[版面分析配置文件](https://github.com/PaddlePaddle/PaddleDetection/tree/release/2.5/configs/picodet/legacy_model/application/layout_analysis)启动训练
|
||||
|
||||
* 修改配置文件
|
||||
|
||||
如果你希望训练自己的数据集，需要修改配置文件中的数据配置、类别数。
|
||||
|
||||
|
||||
以`configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml` 为例，修改的内容如下所示。
|
||||
|
||||
```yaml
|
||||
metric: COCO
|
||||
# 类别数
|
||||
num_classes: 5
|
||||
|
||||
TrainDataset:
|
||||
!COCODataSet
|
||||
# 修改为你自己的训练数据目录
|
||||
image_dir: train
|
||||
# 修改为你自己的训练数据标签文件
|
||||
anno_path: train.json
|
||||
# 修改为你自己的训练数据根目录
|
||||
dataset_dir: /root/publaynet/
|
||||
data_fields: ['image', 'gt_bbox', 'gt_class', 'is_crowd']
|
||||
|
||||
EvalDataset:
|
||||
!COCODataSet
|
||||
# 修改为你自己的验证数据目录
|
||||
image_dir: val
|
||||
# 修改为你自己的验证数据标签文件
|
||||
anno_path: val.json
|
||||
# 修改为你自己的验证数据根目录
|
||||
dataset_dir: /root/publaynet/
|
||||
|
||||
TestDataset:
|
||||
!ImageFolder
|
||||
# 修改为你自己的测试数据标签文件
|
||||
anno_path: /root/publaynet/val.json
|
||||
```
|
||||
|
||||
* 开始训练，在训练时，会默认下载PP-PicoDet预训练模型，这里无需预先下载。
|
||||
|
||||
```bash
|
||||
# GPU训练 支持单卡，多卡训练
|
||||
# 训练日志会自动保存到 log 目录中
|
||||
|
||||
# 单卡训练
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
python3 tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--eval
|
||||
|
||||
# 多卡训练，通过--gpus参数指定卡号
|
||||
export CUDA_VISIBLE_DEVICES=0,1,2,3
|
||||
python3 -m paddle.distributed.launch --gpus '0,1,2,3' tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--eval
|
||||
```
|
||||
|
||||
**注意：**如果训练时显存out memory，将TrainReader中batch_size调小，同时LearningRate中base_lr等比例减小。发布的config均由8卡训练得到，如果改变GPU卡数为1，那么base_lr需要减小8倍。
|
||||
|
||||
正常启动训练后，会看到以下log输出：
|
||||
|
||||
```
|
||||
[08/15 04:02:30] ppdet.utils.checkpoint INFO: Finish loading model weights: /root/.cache/paddle/weights/LCNet_x1_0_pretrained.pdparams
|
||||
[08/15 04:02:46] ppdet.engine INFO: Epoch: [0] [ 0/1929] learning_rate: 0.040000 loss_vfl: 1.216707 loss_bbox: 1.142163 loss_dfl: 0.544196 loss: 2.903065 eta: 17 days, 13:50:26 batch_cost: 15.7452 data_cost: 2.9112 ips: 1.5243 images/s
|
||||
[08/15 04:03:19] ppdet.engine INFO: Epoch: [0] [ 20/1929] learning_rate: 0.064000 loss_vfl: 1.180627 loss_bbox: 0.939552 loss_dfl: 0.442436 loss: 2.628206 eta: 2 days, 12:18:53 batch_cost: 1.5770 data_cost: 0.0008 ips: 15.2184 images/s
|
||||
[08/15 04:03:47] ppdet.engine INFO: Epoch: [0] [ 40/1929] learning_rate: 0.088000 loss_vfl: 0.543321 loss_bbox: 1.071401 loss_dfl: 0.457817 loss: 2.057003 eta: 2 days, 0:07:03 batch_cost: 1.3190 data_cost: 0.0007 ips: 18.1954 images/s
|
||||
[08/15 04:04:12] ppdet.engine INFO: Epoch: [0] [ 60/1929] learning_rate: 0.112000 loss_vfl: 0.630989 loss_bbox: 0.859183 loss_dfl: 0.384702 loss: 1.883143 eta: 1 day, 19:01:29 batch_cost: 1.2177 data_cost: 0.0006 ips: 19.7087 images/s
|
||||
```
|
||||
|
||||
- `--eval`表示训练的同时，进行评估， 评估过程中默认将最佳模型，保存为 `output/picodet_lcnet_x1_0_layout/best_accuracy` 。
|
||||
|
||||
**注意，预测/评估时的配置文件请务必与训练一致。**
|
||||
|
||||
### 5.2. FGD蒸馏训练
|
||||
|
||||
PaddleDetection支持了基于FGD([Focal and Global Knowledge Distillation for Detectors](https://arxiv.org/abs/2111.11837v1))蒸馏的目标检测模型训练过程，FGD蒸馏分为两个部分`Focal`和`Global`。`Focal`蒸馏分离图像的前景和背景，让学生模型分别关注教师模型的前景和背景部分特征的关键像素；`Global`蒸馏部分重建不同像素之间的关系并将其从教师转移到学生，以补偿`Focal`蒸馏中丢失的全局信息。
|
||||
|
||||
更换数据集，修改【TODO】配置中的数据配置、类别数，具体可以参考4.1。启动训练：
|
||||
|
||||
```bash
|
||||
# 单卡训练
|
||||
export CUDA_VISIBLE_DEVICES=0
|
||||
python3 tools/train.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
--eval
|
||||
```
|
||||
|
||||
- `-c`: 指定模型配置文件。
|
||||
- `--slim_config`: 指定压缩策略配置文件。
|
||||
|
||||
## 6. 模型评估与预测
|
||||
|
||||
### 6.1. 指标评估
|
||||
|
||||
训练中模型参数默认保存在`output/picodet_lcnet_x1_0_layout`目录下。在评估指标时，需要设置`weights`指向保存的参数文件。评估数据集可以通过 `configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml` 修改`EvalDataset`中的 `image_dir`、`anno_path`和`dataset_dir` 设置。
|
||||
|
||||
```bash
|
||||
# GPU 评估， weights 为待测权重
|
||||
python3 tools/eval.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights=./output/picodet_lcnet_x1_0_layout/best_model
|
||||
```
|
||||
|
||||
会输出以下信息，打印出mAP、AP0.5等信息。
|
||||
|
||||
```py
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.935
|
||||
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.979
|
||||
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.956
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.404
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.782
|
||||
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.969
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.539
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.938
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.949
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.495
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.818
|
||||
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.978
|
||||
[08/15 07:07:09] ppdet.engine INFO: Total sample number: 11245, averge FPS: 24.405059207157436
|
||||
[08/15 07:07:09] ppdet.engine INFO: Best test bbox ap is 0.935.
|
||||
```
|
||||
|
||||
若使用**提供的预训练模型进行评估**，或使用**FGD蒸馏训练的模型**，更换`weights`模型路径，执行如下命令进行评估：
|
||||
|
||||
```
|
||||
python3 tools/eval.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights=output/picodet_lcnet_x2_5_layout/best_model
|
||||
```
|
||||
|
||||
- `-c`: 指定模型配置文件。
|
||||
- `--slim_config`: 指定蒸馏策略配置文件。
|
||||
- `-o weights`: 指定蒸馏算法训好的模型路径。
|
||||
|
||||
### 6.2 测试版面分析结果
|
||||
|
||||
|
||||
预测使用的配置文件必须与训练一致，如您通过 `python3 tools/train.py -c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml` 完成了模型的训练过程。
|
||||
|
||||
使用 PaddleDetection 训练好的模型，您可以使用如下命令进行模型预测。
|
||||
|
||||
```bash
|
||||
python3 tools/infer.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights='output/picodet_lcnet_x1_0_layout/best_model.pdparams' \
|
||||
--infer_img='docs/images/layout.jpg' \
|
||||
--output_dir=output_dir/ \
|
||||
--draw_threshold=0.5
|
||||
```
|
||||
|
||||
- `--infer_img`: 推理单张图片，也可以通过`--infer_dir`推理文件中的所有图片。
|
||||
- `--output_dir`: 指定可视化结果保存路径。
|
||||
- `--draw_threshold`:指定绘制结果框的NMS阈值。
|
||||
|
||||
若使用**提供的预训练模型进行预测**，或使用**FGD蒸馏训练的模型**，更换`weights`模型路径，执行如下命令进行预测：
|
||||
|
||||
```
|
||||
python3 tools/infer.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights='output/picodet_lcnet_x2_5_layout/best_model.pdparams' \
|
||||
--infer_img='docs/images/layout.jpg' \
|
||||
--output_dir=output_dir/ \
|
||||
--draw_threshold=0.5
|
||||
```
|
||||
|
||||
|
||||
## 7. 模型导出与预测
|
||||
|
||||
|
||||
### 7.1 模型导出
|
||||
|
||||
inference 模型（`paddle.jit.save`保存的模型） 一般是模型训练，把模型结构和模型参数保存在文件中的固化模型，多用于预测部署场景。 训练过程中保存的模型是checkpoints模型，保存的只有模型的参数，多用于恢复训练等。 与checkpoints模型相比，inference 模型会额外保存模型的结构信息，在预测部署、加速推理上性能优越，灵活方便，适合于实际系统集成。
|
||||
|
||||
版面分析模型转inference模型步骤如下：
|
||||
|
||||
```bash
|
||||
python3 tools/export_model.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
-o weights=output/picodet_lcnet_x1_0_layout/best_model \
|
||||
--output_dir=output_inference/
|
||||
```
|
||||
|
||||
* 如无需导出后处理，请指定：`-o export.benchmark=True`（如果-o已出现过，此处删掉-o）
|
||||
* 如无需导出NMS，请指定：`-o export.nms=False`
|
||||
|
||||
转换成功后，在目录下有三个文件：
|
||||
|
||||
```
|
||||
output_inference/picodet_lcnet_x1_0_layout/
|
||||
├── model.pdiparams # inference模型的参数文件
|
||||
├── model.pdiparams.info # inference模型的参数信息，可忽略
|
||||
└── model.pdmodel # inference模型的模型结构文件
|
||||
```
|
||||
|
||||
若使用**提供的预训练模型转Inference模型**，或使用**FGD蒸馏训练的模型**，更换`weights`模型路径，模型转inference模型步骤如下：
|
||||
|
||||
```bash
|
||||
python3 tools/export_model.py \
|
||||
-c configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x1_0_layout.yml \
|
||||
--slim_config configs/picodet/legacy_model/application/layout_analysis/picodet_lcnet_x2_5_layout.yml \
|
||||
-o weights=./output/picodet_lcnet_x2_5_layout/best_model \
|
||||
--output_dir=output_inference/
|
||||
```
|
||||
|
||||
|
||||
|
||||
### 7.2 模型推理
|
||||
|
||||
若使用**提供的推理训练模型推理**，或使用**FGD蒸馏训练的模型**，更换`model_dir`推理模型路径，执行如下命令进行推理：
|
||||
|
||||
```bash
|
||||
python3 deploy/python/infer.py \
|
||||
--model_dir=output_inference/picodet_lcnet_x1_0_layout/ \
|
||||
--image_file=docs/images/layout.jpg \
|
||||
--device=CPU
|
||||
```
|
||||
|
||||
- --device：指定GPU、CPU设备
|
||||
|
||||
模型推理完成，会看到以下log输出
|
||||
|
||||
```
|
||||
------------------------------------------
|
||||
----------- Model Configuration -----------
|
||||
Model Arch: PicoDet
|
||||
Transform Order:
|
||||
--transform op: Resize
|
||||
--transform op: NormalizeImage
|
||||
--transform op: Permute
|
||||
--transform op: PadStride
|
||||
--------------------------------------------
|
||||
class_id:0, confidence:0.9921, left_top:[20.18,35.66],right_bottom:[341.58,600.99]
|
||||
class_id:0, confidence:0.9914, left_top:[19.77,611.42],right_bottom:[341.48,901.82]
|
||||
class_id:0, confidence:0.9904, left_top:[369.36,375.10],right_bottom:[691.29,600.59]
|
||||
class_id:0, confidence:0.9835, left_top:[369.60,608.60],right_bottom:[691.38,736.72]
|
||||
class_id:0, confidence:0.9830, left_top:[369.58,805.38],right_bottom:[690.97,901.80]
|
||||
class_id:0, confidence:0.9716, left_top:[383.68,271.44],right_bottom:[688.93,335.39]
|
||||
class_id:0, confidence:0.9452, left_top:[370.82,34.48],right_bottom:[688.10,63.54]
|
||||
class_id:1, confidence:0.8712, left_top:[370.84,771.03],right_bottom:[519.30,789.13]
|
||||
class_id:3, confidence:0.9856, left_top:[371.28,67.85],right_bottom:[685.73,267.72]
|
||||
save result to: output/layout.jpg
|
||||
Test iter 0
|
||||
------------------ Inference Time Info ----------------------
|
||||
total_time(ms): 2196.0, img_num: 1
|
||||
average latency time(ms): 2196.00, QPS: 0.455373
|
||||
preprocess_time(ms): 2172.50, inference_time(ms): 11.90, postprocess_time(ms): 11.60
|
||||
```
|
||||
|
||||
- Model：模型结构
|
||||
- Transform Order：预处理操作
|
||||
- class_id、confidence、left_top、right_bottom：分别表示类别id、置信度、左上角坐标、右下角坐标
|
||||
- save result to：可视化版面分析结果保存路径，默认保存到`./output`文件夹
|
||||
- Inference Time Info：推理时间，其中preprocess_time表示预处理耗时，inference_time表示模型预测耗时，postprocess_time表示后处理耗时
|
||||
|
||||
可视化版面结果如下图所示
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/layout/layout_res.jpg" width="800">
|
||||
</div>
|
||||
|
||||
|
||||
|
||||
## Citations
|
||||
|
||||
```
|
||||
@inproceedings{zhong2019publaynet,
|
||||
title={PubLayNet: largest dataset ever for document layout analysis},
|
||||
author={Zhong, Xu and Tang, Jianbin and Yepes, Antonio Jimeno},
|
||||
booktitle={2019 International Conference on Document Analysis and Recognition (ICDAR)},
|
||||
year={2019},
|
||||
volume={},
|
||||
number={},
|
||||
pages={1015-1022},
|
||||
doi={10.1109/ICDAR.2019.00166},
|
||||
ISSN={1520-5363},
|
||||
month={Sep.},
|
||||
organization={IEEE}
|
||||
}
|
||||
|
||||
@inproceedings{yang2022focal,
|
||||
title={Focal and global knowledge distillation for detectors},
|
||||
author={Yang, Zhendong and Li, Zhe and Jiang, Xiaohu and Gong, Yuan and Yuan, Zehuan and Zhao, Danpei and Yuan, Chun},
|
||||
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
|
||||
pages={4643--4652},
|
||||
year={2022}
|
||||
}
|
||||
```
|
||||
13
ppstructure/layout/__init__.py
Normal file
13
ppstructure/layout/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
143
ppstructure/layout/predict_layout.py
Executable file
143
ppstructure/layout/predict_layout.py
Executable file
@@ -0,0 +1,143 @@
|
||||
# Copyright (c) 2020 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.insert(0, os.path.abspath(os.path.join(__dir__, "../..")))
|
||||
|
||||
os.environ["FLAGS_allocator_strategy"] = "auto_growth"
|
||||
|
||||
import cv2
|
||||
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.utility import get_image_file_list, check_and_read
|
||||
from ppstructure.utility import parse_args
|
||||
from picodet_postprocess import PicoDetPostProcess
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class LayoutPredictor(object):
|
||||
def __init__(self, args):
|
||||
pre_process_list = [
|
||||
{"Resize": {"size": [800, 608]}},
|
||||
{
|
||||
"NormalizeImage": {
|
||||
"std": [0.229, 0.224, 0.225],
|
||||
"mean": [0.485, 0.456, 0.406],
|
||||
"scale": "1./255.",
|
||||
"order": "hwc",
|
||||
}
|
||||
},
|
||||
{"ToCHWImage": None},
|
||||
{"KeepKeys": {"keep_keys": ["image"]}},
|
||||
]
|
||||
postprocess_params = {
|
||||
"name": "PicoDetPostProcess",
|
||||
"layout_dict_path": args.layout_dict_path,
|
||||
"score_threshold": args.layout_score_threshold,
|
||||
"nms_threshold": args.layout_nms_threshold,
|
||||
}
|
||||
|
||||
self.preprocess_op = create_operators(pre_process_list)
|
||||
self.postprocess_op = build_post_process(postprocess_params)
|
||||
(
|
||||
self.predictor,
|
||||
self.input_tensor,
|
||||
self.output_tensors,
|
||||
self.config,
|
||||
) = utility.create_predictor(args, "layout", logger)
|
||||
self.use_onnx = args.use_onnx
|
||||
|
||||
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()
|
||||
|
||||
preds, elapse = 0, 1
|
||||
starttime = time.time()
|
||||
|
||||
np_score_list, np_boxes_list = [], []
|
||||
if self.use_onnx:
|
||||
input_dict = {}
|
||||
input_dict[self.input_tensor.name] = img
|
||||
outputs = self.predictor.run(self.output_tensors, input_dict)
|
||||
num_outs = int(len(outputs) / 2)
|
||||
for out_idx in range(num_outs):
|
||||
np_score_list.append(outputs[out_idx])
|
||||
np_boxes_list.append(outputs[out_idx + num_outs])
|
||||
else:
|
||||
self.input_tensor.copy_from_cpu(img)
|
||||
self.predictor.run()
|
||||
output_names = self.predictor.get_output_names()
|
||||
num_outs = int(len(output_names) / 2)
|
||||
for out_idx in range(num_outs):
|
||||
np_score_list.append(
|
||||
self.predictor.get_output_handle(
|
||||
output_names[out_idx]
|
||||
).copy_to_cpu()
|
||||
)
|
||||
np_boxes_list.append(
|
||||
self.predictor.get_output_handle(
|
||||
output_names[out_idx + num_outs]
|
||||
).copy_to_cpu()
|
||||
)
|
||||
preds = dict(boxes=np_score_list, boxes_num=np_boxes_list)
|
||||
|
||||
post_preds = self.postprocess_op(ori_im, img, preds)
|
||||
elapse = time.time() - starttime
|
||||
return post_preds, elapse
|
||||
|
||||
|
||||
def main(args):
|
||||
image_file_list = get_image_file_list(args.image_dir)
|
||||
layout_predictor = LayoutPredictor(args)
|
||||
count = 0
|
||||
total_time = 0
|
||||
|
||||
repeats = 50
|
||||
for image_file in image_file_list:
|
||||
img, flag, _ = check_and_read(image_file)
|
||||
if not flag:
|
||||
img = cv2.imread(image_file)
|
||||
if img is None:
|
||||
logger.info("error in loading image:{}".format(image_file))
|
||||
continue
|
||||
|
||||
layout_res, elapse = layout_predictor(img)
|
||||
|
||||
logger.info("result: {}".format(layout_res))
|
||||
|
||||
if count > 0:
|
||||
total_time += elapse
|
||||
count += 1
|
||||
logger.info("Predict time of {}: {}".format(image_file, elapse))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
49
ppstructure/pdf2word/README.md
Normal file
49
ppstructure/pdf2word/README.md
Normal file
@@ -0,0 +1,49 @@
|
||||
# PDF2WORD
|
||||
|
||||
PDF2Word是PaddleOCR社区开发者 [whjdark](https://github.com/whjdark) 基于PP-StructureV2版面分析与恢复模型实现的PDF转换Word应用程序，提供可直接安装的exe应用程序，**方便Windows用户免环境配置运行**
|
||||
|
||||
## 1.使用
|
||||
|
||||
### 应用程序
|
||||
|
||||
1. 下载与安装：针对Windows用户，根据[软件下载]()一节下载软件后，运行 `pdf2word.exe` 。若您下载的是lite版本，安装过程中会在线下载环境依赖、模型等必要资源，安装时间较长，请确保网络畅通。serve版本打包了相关依赖，安装时间较短，可按需下载。
|
||||
|
||||
2. 转换：由于PP-Structure根据中英文数据分别进行适配，在转换相应文件时可**根据文档语言进行相应选择**。
|
||||
|
||||
### 脚本运行
|
||||
|
||||
3. 打开结果：点击`显示结果`，即可打开转换完成后的文件夹
|
||||
|
||||
> 注意：
|
||||
>
|
||||
> - 初次安装程序根据不同设备需要等待1-2分钟不等
|
||||
> - 使用Office与WPS打开的Word结果会出现不同，推荐以Office为准
|
||||
> - 本程序使用 [QPT](https://github.com/QPT-Family/QPT) 进行应用程序打包，感谢 [GT-ZhangAcer](https://github.com/GT-ZhangAcer) 对打包过程的支持
|
||||
> - 应用程序仅支持正版win10，11系统，不支持盗版Windows系统，若在安装过程中出现报错或缺少依赖，推荐直接使用 `paddleocr` whl包应用PDF2Word功能，详情可查看[链接](https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.6/ppstructure/docs/quickstart.md)
|
||||
|
||||
### 脚本启动界面
|
||||
|
||||
首次运行需要将切换路径到PaddleOCR文件目录 ，然后运行代码
|
||||
|
||||
```
|
||||
cd ./ppstructure/pdf2word
|
||||
python pdf2word.py
|
||||
```
|
||||
|
||||
### PaddleOCR whl包
|
||||
|
||||
针对Linux、Mac用户或已经拥有Python环境的用户，**推荐安装 `paddleocr` whl包直接应用PDF2Word功能**，详情可查看[链接](https://github.com/PaddlePaddle/PaddleOCR/blob/release/2.6/ppstructure/docs/quickstart.md)
|
||||
|
||||
<a name="download"></a>
|
||||
|
||||
## 2.软件下载
|
||||
|
||||
如需获取已打包程序，可以扫描下方二维码，关注公众号填写问卷后，加入PaddleOCR官方交流群免费获取20G OCR学习大礼包，内含OCR场景应用集合（包含数码管、液晶屏、车牌、高精度SVTR模型等7个垂类模型）、《动手学OCR》电子书、课程回放视频、前沿论文等重磅资料
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/50011306/186369636-35f2008b-df5a-4784-b1f5-cebebcb2b7a5.jpg" width = "150" height = "150" />
|
||||
</div>
|
||||
|
||||
## 3.版本说明
|
||||
|
||||
v0.2版：新加入PDF解析功能，仅提供full版本，打包了所有依赖包与模型文件，尽可能避免安装失败问题。若仍然安装失败，推荐使用 `paddleocr` whl包
|
||||
BIN
ppstructure/pdf2word/icons/chinese.png
Normal file
BIN
ppstructure/pdf2word/icons/chinese.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 2.6 KiB |
BIN
ppstructure/pdf2word/icons/english.png
Normal file
BIN
ppstructure/pdf2word/icons/english.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 2.8 KiB |
BIN
ppstructure/pdf2word/icons/folder-open.png
Normal file
BIN
ppstructure/pdf2word/icons/folder-open.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 2.4 KiB |
BIN
ppstructure/pdf2word/icons/folder-plus.png
Normal file
BIN
ppstructure/pdf2word/icons/folder-plus.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 2.5 KiB |
537
ppstructure/pdf2word/pdf2word.py
Normal file
537
ppstructure/pdf2word/pdf2word.py
Normal file
@@ -0,0 +1,537 @@
|
||||
# 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 sys
|
||||
import tarfile
|
||||
import os
|
||||
import time
|
||||
import datetime
|
||||
import functools
|
||||
import cv2
|
||||
import platform
|
||||
import numpy as np
|
||||
from paddle.utils import try_import
|
||||
|
||||
fitz = try_import("fitz")
|
||||
from PIL import Image
|
||||
from qtpy.QtWidgets import (
|
||||
QApplication,
|
||||
QWidget,
|
||||
QPushButton,
|
||||
QProgressBar,
|
||||
QGridLayout,
|
||||
QMessageBox,
|
||||
QLabel,
|
||||
QFileDialog,
|
||||
QCheckBox,
|
||||
)
|
||||
from qtpy.QtCore import Signal, QThread, QObject
|
||||
from qtpy.QtGui import QImage, QPixmap, QIcon
|
||||
|
||||
file = os.path.dirname(os.path.abspath(__file__))
|
||||
root = os.path.abspath(os.path.join(file, "../../"))
|
||||
sys.path.append(file)
|
||||
sys.path.insert(0, root)
|
||||
|
||||
from ppstructure.predict_system import StructureSystem, save_structure_res
|
||||
from ppstructure.utility import parse_args, draw_structure_result
|
||||
from ppocr.utils.network import download_with_progressbar
|
||||
from ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
|
||||
|
||||
# from ScreenShotWidget import ScreenShotWidget
|
||||
|
||||
__APPNAME__ = "pdf2word"
|
||||
__VERSION__ = "0.2.2"
|
||||
|
||||
URLs_EN = {
|
||||
# 下载超英文轻量级PP-OCRv3模型的检测模型并解压
|
||||
"en_PP-OCRv3_det_infer": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar",
|
||||
# 下载英文轻量级PP-OCRv3模型的识别模型并解压
|
||||
"en_PP-OCRv3_rec_infer": "https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar",
|
||||
# 下载超轻量级英文表格英文模型并解压
|
||||
"en_ppstructure_mobile_v2.0_SLANet_infer": "https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar",
|
||||
# 英文版面分析模型
|
||||
"picodet_lcnet_x1_0_fgd_layout_infer": "https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar",
|
||||
}
|
||||
DICT_EN = {
|
||||
"rec_char_dict_path": "en_dict.txt",
|
||||
"layout_dict_path": "layout_publaynet_dict.txt",
|
||||
}
|
||||
|
||||
URLs_CN = {
|
||||
# 下载超中文轻量级PP-OCRv3模型的检测模型并解压
|
||||
"cn_PP-OCRv3_det_infer": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_det_infer.tar",
|
||||
# 下载中文轻量级PP-OCRv3模型的识别模型并解压
|
||||
"cn_PP-OCRv3_rec_infer": "https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar",
|
||||
# 下载超轻量级英文表格英文模型并解压
|
||||
"cn_ppstructure_mobile_v2.0_SLANet_infer": "https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar",
|
||||
# 中文版面分析模型
|
||||
"picodet_lcnet_x1_0_fgd_layout_cdla_infer": "https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_cdla_infer.tar",
|
||||
}
|
||||
DICT_CN = {
|
||||
"rec_char_dict_path": "ppocr_keys_v1.txt",
|
||||
"layout_dict_path": "layout_cdla_dict.txt",
|
||||
}
|
||||
|
||||
|
||||
def QImageToCvMat(incomingImage) -> np.array:
|
||||
"""
|
||||
Converts a QImage into an opencv MAT format
|
||||
"""
|
||||
|
||||
incomingImage = incomingImage.convertToFormat(QImage.Format.Format_RGBA8888)
|
||||
|
||||
width = incomingImage.width()
|
||||
height = incomingImage.height()
|
||||
|
||||
ptr = incomingImage.bits()
|
||||
ptr.setsize(height * width * 4)
|
||||
arr = np.frombuffer(ptr, np.uint8).reshape((height, width, 4))
|
||||
return arr
|
||||
|
||||
|
||||
def readImage(image_file) -> list:
|
||||
if os.path.basename(image_file)[-3:] == "pdf":
|
||||
imgs = []
|
||||
with fitz.open(image_file) as pdf:
|
||||
for pg in range(0, pdf.pageCount):
|
||||
page = pdf[pg]
|
||||
mat = fitz.Matrix(2, 2)
|
||||
pm = page.getPixmap(matrix=mat, alpha=False)
|
||||
|
||||
# if width or height > 2000 pixels, don't enlarge the image
|
||||
if pm.width > 2000 or pm.height > 2000:
|
||||
pm = page.getPixmap(matrix=fitz.Matrix(1, 1), alpha=False)
|
||||
|
||||
img = Image.frombytes("RGB", [pm.width, pm.height], pm.samples)
|
||||
img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
|
||||
imgs.append(img)
|
||||
else:
|
||||
img = cv2.imread(image_file, cv2.IMREAD_COLOR)
|
||||
if img is not None:
|
||||
imgs = [img]
|
||||
|
||||
return imgs
|
||||
|
||||
|
||||
class Worker(QThread):
|
||||
progressBarValue = Signal(int)
|
||||
progressBarRange = Signal(int)
|
||||
endsignal = Signal()
|
||||
exceptedsignal = Signal(str) # 发送一个异常信号
|
||||
loopFlag = True
|
||||
|
||||
def __init__(self, predictors, save_pdf, vis_font_path, use_pdf2docx_api):
|
||||
super(Worker, self).__init__()
|
||||
self.predictors = predictors
|
||||
self.save_pdf = save_pdf
|
||||
self.vis_font_path = vis_font_path
|
||||
self.lang = "EN"
|
||||
self.imagePaths = []
|
||||
self.use_pdf2docx_api = use_pdf2docx_api
|
||||
self.outputDir = None
|
||||
self.totalPageCnt = 0
|
||||
self.pageCnt = 0
|
||||
self.setStackSize(1024 * 1024)
|
||||
|
||||
def setImagePath(self, imagePaths):
|
||||
self.imagePaths = imagePaths
|
||||
|
||||
def setLang(self, lang):
|
||||
self.lang = lang
|
||||
|
||||
def setOutputDir(self, outputDir):
|
||||
self.outputDir = outputDir
|
||||
|
||||
def setPDFParser(self, enabled):
|
||||
self.use_pdf2docx_api = enabled
|
||||
|
||||
def resetPageCnt(self):
|
||||
self.pageCnt = 0
|
||||
|
||||
def resetTotalPageCnt(self):
|
||||
self.totalPageCnt = 0
|
||||
|
||||
def ppocrPrecitor(self, imgs, img_name):
|
||||
all_res = []
|
||||
# update progress bar ranges
|
||||
self.totalPageCnt += len(imgs)
|
||||
self.progressBarRange.emit(self.totalPageCnt)
|
||||
# processing pages
|
||||
for index, img in enumerate(imgs):
|
||||
res, time_dict = self.predictors[self.lang](img)
|
||||
|
||||
# save output
|
||||
save_structure_res(res, self.outputDir, img_name)
|
||||
# draw_img = draw_structure_result(img, res, self.vis_font_path)
|
||||
# img_save_path = os.path.join(self.outputDir, img_name, 'show_{}.jpg'.format(index))
|
||||
# if res != []:
|
||||
# cv2.imwrite(img_save_path, draw_img)
|
||||
|
||||
# recovery
|
||||
h, w, _ = img.shape
|
||||
res = sorted_layout_boxes(res, w)
|
||||
all_res += res
|
||||
self.pageCnt += 1
|
||||
self.progressBarValue.emit(self.pageCnt)
|
||||
|
||||
if all_res != []:
|
||||
try:
|
||||
convert_info_docx(imgs, all_res, self.outputDir, img_name)
|
||||
except Exception as ex:
|
||||
print(
|
||||
"error in layout recovery image:{}, err msg: {}".format(
|
||||
img_name, ex
|
||||
)
|
||||
)
|
||||
print("Predict time : {:.3f}s".format(time_dict["all"]))
|
||||
print("result save to {}".format(self.outputDir))
|
||||
|
||||
def run(self):
|
||||
self.resetPageCnt()
|
||||
self.resetTotalPageCnt()
|
||||
try:
|
||||
os.makedirs(self.outputDir, exist_ok=True)
|
||||
for i, image_file in enumerate(self.imagePaths):
|
||||
if not self.loopFlag:
|
||||
break
|
||||
# using use_pdf2docx_api for PDF parsing
|
||||
if self.use_pdf2docx_api and os.path.basename(image_file)[-3:] == "pdf":
|
||||
try_import("pdf2docx")
|
||||
from pdf2docx.converter import Converter
|
||||
|
||||
self.totalPageCnt += 1
|
||||
self.progressBarRange.emit(self.totalPageCnt)
|
||||
print("===============using use_pdf2docx_api===============")
|
||||
img_name = os.path.basename(image_file).split(".")[0]
|
||||
docx_file = os.path.join(self.outputDir, "{}.docx".format(img_name))
|
||||
cv = Converter(image_file)
|
||||
cv.convert(docx_file)
|
||||
cv.close()
|
||||
print("docx save to {}".format(docx_file))
|
||||
self.pageCnt += 1
|
||||
self.progressBarValue.emit(self.pageCnt)
|
||||
else:
|
||||
# using PPOCR for PDF/Image parsing
|
||||
imgs = readImage(image_file)
|
||||
if len(imgs) == 0:
|
||||
continue
|
||||
img_name = os.path.basename(image_file).split(".")[0]
|
||||
os.makedirs(os.path.join(self.outputDir, img_name), exist_ok=True)
|
||||
self.ppocrPrecitor(imgs, img_name)
|
||||
# file processed
|
||||
self.endsignal.emit()
|
||||
# self.exec()
|
||||
except Exception as e:
|
||||
self.exceptedsignal.emit(str(e)) # 将异常发送给UI进程
|
||||
|
||||
|
||||
class APP_Image2Doc(QWidget):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# self.setFixedHeight(100)
|
||||
# self.setFixedWidth(520)
|
||||
|
||||
# settings
|
||||
self.imagePaths = []
|
||||
# self.screenShotWg = ScreenShotWidget()
|
||||
self.screenShot = None
|
||||
self.save_pdf = False
|
||||
self.output_dir = None
|
||||
self.vis_font_path = os.path.join(root, "doc", "fonts", "simfang.ttf")
|
||||
self.use_pdf2docx_api = False
|
||||
|
||||
# ProgressBar
|
||||
self.pb = QProgressBar()
|
||||
self.pb.setRange(0, 100)
|
||||
self.pb.setValue(0)
|
||||
|
||||
# 初始化界面
|
||||
self.setupUi()
|
||||
|
||||
# 下载模型
|
||||
self.downloadModels(URLs_EN)
|
||||
self.downloadModels(URLs_CN)
|
||||
|
||||
# 初始化模型
|
||||
predictors = {
|
||||
"EN": self.initPredictor("EN"),
|
||||
"CN": self.initPredictor("CN"),
|
||||
}
|
||||
|
||||
# 设置工作进程
|
||||
self._thread = Worker(
|
||||
predictors, self.save_pdf, self.vis_font_path, self.use_pdf2docx_api
|
||||
)
|
||||
self._thread.progressBarValue.connect(self.handleProgressBarUpdateSingal)
|
||||
self._thread.endsignal.connect(self.handleEndsignalSignal)
|
||||
# self._thread.finished.connect(QObject.deleteLater)
|
||||
self._thread.progressBarRange.connect(self.handleProgressBarRangeSingal)
|
||||
self._thread.exceptedsignal.connect(self.handleThreadException)
|
||||
self.time_start = 0 # save start time
|
||||
|
||||
def setupUi(self):
|
||||
self.setObjectName("MainWindow")
|
||||
self.setWindowTitle(__APPNAME__ + " " + __VERSION__)
|
||||
|
||||
layout = QGridLayout()
|
||||
|
||||
self.openFileButton = QPushButton("打开文件")
|
||||
self.openFileButton.setIcon(QIcon(QPixmap("./icons/folder-plus.png")))
|
||||
layout.addWidget(self.openFileButton, 0, 0, 1, 1)
|
||||
self.openFileButton.clicked.connect(self.handleOpenFileSignal)
|
||||
|
||||
# screenShotButton = QPushButton("截图识别")
|
||||
# layout.addWidget(screenShotButton, 0, 1, 1, 1)
|
||||
# screenShotButton.clicked.connect(self.screenShotSlot)
|
||||
# screenShotButton.setEnabled(False) # temporarily disenble
|
||||
|
||||
self.startCNButton = QPushButton("中文转换")
|
||||
self.startCNButton.setIcon(QIcon(QPixmap("./icons/chinese.png")))
|
||||
layout.addWidget(self.startCNButton, 0, 1, 1, 1)
|
||||
self.startCNButton.clicked.connect(
|
||||
functools.partial(self.handleStartSignal, "CN", False)
|
||||
)
|
||||
|
||||
self.startENButton = QPushButton("英文转换")
|
||||
self.startENButton.setIcon(QIcon(QPixmap("./icons/english.png")))
|
||||
layout.addWidget(self.startENButton, 0, 2, 1, 1)
|
||||
self.startENButton.clicked.connect(
|
||||
functools.partial(self.handleStartSignal, "EN", False)
|
||||
)
|
||||
|
||||
self.PDFParserButton = QPushButton("PDF解析", self)
|
||||
layout.addWidget(self.PDFParserButton, 0, 3, 1, 1)
|
||||
self.PDFParserButton.clicked.connect(
|
||||
functools.partial(self.handleStartSignal, "CN", True)
|
||||
)
|
||||
|
||||
self.showResultButton = QPushButton("显示结果")
|
||||
self.showResultButton.setIcon(QIcon(QPixmap("./icons/folder-open.png")))
|
||||
layout.addWidget(self.showResultButton, 0, 4, 1, 1)
|
||||
self.showResultButton.clicked.connect(self.handleShowResultSignal)
|
||||
|
||||
# ProgressBar
|
||||
layout.addWidget(self.pb, 2, 0, 1, 5)
|
||||
# time estimate label
|
||||
self.timeEstLabel = QLabel(("Time Left: --"))
|
||||
layout.addWidget(self.timeEstLabel, 3, 0, 1, 5)
|
||||
|
||||
self.setLayout(layout)
|
||||
|
||||
def downloadModels(self, URLs):
|
||||
# using custom model
|
||||
tar_file_name_list = [
|
||||
"inference.pdiparams",
|
||||
"inference.pdiparams.info",
|
||||
"inference.pdmodel",
|
||||
"model.pdiparams",
|
||||
"model.pdiparams.info",
|
||||
"model.pdmodel",
|
||||
]
|
||||
model_path = os.path.join(root, "inference")
|
||||
os.makedirs(model_path, exist_ok=True)
|
||||
|
||||
# download and unzip models
|
||||
for name in URLs.keys():
|
||||
url = URLs[name]
|
||||
print("Try downloading file: {}".format(url))
|
||||
tarname = url.split("/")[-1]
|
||||
tarpath = os.path.join(model_path, tarname)
|
||||
if os.path.exists(tarpath):
|
||||
print("File have already exist. skip")
|
||||
else:
|
||||
try:
|
||||
download_with_progressbar(url, tarpath)
|
||||
except Exception as e:
|
||||
print("Error occurred when downloading file, error message:")
|
||||
print(e)
|
||||
|
||||
# unzip model tar
|
||||
try:
|
||||
with tarfile.open(tarpath, "r") as tarObj:
|
||||
storage_dir = os.path.join(model_path, name)
|
||||
os.makedirs(storage_dir, exist_ok=True)
|
||||
for member in tarObj.getmembers():
|
||||
filename = None
|
||||
for tar_file_name in tar_file_name_list:
|
||||
if tar_file_name in member.name:
|
||||
filename = tar_file_name
|
||||
if filename is None:
|
||||
continue
|
||||
file = tarObj.extractfile(member)
|
||||
with open(os.path.join(storage_dir, filename), "wb") as f:
|
||||
f.write(file.read())
|
||||
except Exception as e:
|
||||
print("Error occurred when unziping file, error message:")
|
||||
print(e)
|
||||
|
||||
def initPredictor(self, lang="EN"):
|
||||
# init predictor args
|
||||
args = parse_args()
|
||||
args.table_max_len = 488
|
||||
args.ocr = True
|
||||
args.recovery = True
|
||||
args.save_pdf = self.save_pdf
|
||||
args.table_char_dict_path = os.path.join(
|
||||
root, "ppocr", "utils", "dict", "table_structure_dict.txt"
|
||||
)
|
||||
if lang == "EN":
|
||||
args.det_model_dir = os.path.join(
|
||||
root, "inference", "en_PP-OCRv3_det_infer" # 此处从这里找到模型存放位置
|
||||
)
|
||||
args.rec_model_dir = os.path.join(
|
||||
root, "inference", "en_PP-OCRv3_rec_infer"
|
||||
)
|
||||
args.table_model_dir = os.path.join(
|
||||
root, "inference", "en_ppstructure_mobile_v2.0_SLANet_infer"
|
||||
)
|
||||
args.output = os.path.join(root, "output") # 结果保存路径
|
||||
args.layout_model_dir = os.path.join(
|
||||
root, "inference", "picodet_lcnet_x1_0_fgd_layout_infer"
|
||||
)
|
||||
lang_dict = DICT_EN
|
||||
elif lang == "CN":
|
||||
args.det_model_dir = os.path.join(
|
||||
root, "inference", "cn_PP-OCRv3_det_infer" # 此处从这里找到模型存放位置
|
||||
)
|
||||
args.rec_model_dir = os.path.join(
|
||||
root, "inference", "cn_PP-OCRv3_rec_infer"
|
||||
)
|
||||
args.table_model_dir = os.path.join(
|
||||
root, "inference", "cn_ppstructure_mobile_v2.0_SLANet_infer"
|
||||
)
|
||||
args.output = os.path.join(root, "output") # 结果保存路径
|
||||
args.layout_model_dir = os.path.join(
|
||||
root, "inference", "picodet_lcnet_x1_0_fgd_layout_cdla_infer"
|
||||
)
|
||||
lang_dict = DICT_CN
|
||||
else:
|
||||
raise ValueError("Unsupported language")
|
||||
args.rec_char_dict_path = os.path.join(
|
||||
root, "ppocr", "utils", lang_dict["rec_char_dict_path"]
|
||||
)
|
||||
args.layout_dict_path = os.path.join(
|
||||
root, "ppocr", "utils", "dict", "layout_dict", lang_dict["layout_dict_path"]
|
||||
)
|
||||
# init predictor
|
||||
return StructureSystem(args)
|
||||
|
||||
def handleOpenFileSignal(self):
|
||||
"""
|
||||
可以多选图像文件
|
||||
"""
|
||||
selectedFiles = QFileDialog.getOpenFileNames(
|
||||
self, "多文件选择", "/", "图片文件 (*.png *.jpeg *.jpg *.bmp *.pdf)"
|
||||
)[0]
|
||||
if len(selectedFiles) > 0:
|
||||
self.imagePaths = selectedFiles
|
||||
self.screenShot = None # discard screenshot temp image
|
||||
self.pb.setValue(0)
|
||||
|
||||
# def screenShotSlot(self):
|
||||
# '''
|
||||
# 选定图像文件和截图的转换过程只能同时进行一个
|
||||
# 截图只能同时转换一个
|
||||
# '''
|
||||
# self.screenShotWg.start()
|
||||
# if self.screenShotWg.captureImage:
|
||||
# self.screenShot = self.screenShotWg.captureImage
|
||||
# self.imagePaths.clear() # discard openfile temp list
|
||||
# self.pb.setRange(0, 1)
|
||||
# self.pb.setValue(0)
|
||||
|
||||
def handleStartSignal(self, lang="EN", pdfParser=False):
|
||||
if self.screenShot: # for screenShot
|
||||
img_name = "screenshot_" + time.strftime("%Y%m%d%H%M%S", time.localtime())
|
||||
image = QImageToCvMat(self.screenShot)
|
||||
self.predictAndSave(image, img_name, lang)
|
||||
# update Progress Bar
|
||||
self.pb.setValue(1)
|
||||
QMessageBox.information(self, "Information", "文档提取完成")
|
||||
elif len(self.imagePaths) > 0: # for image file selection
|
||||
# Must set image path list and language before start
|
||||
self.output_dir = os.path.join(
|
||||
os.path.dirname(self.imagePaths[0]), "output"
|
||||
) # output_dir should be same as imagepath
|
||||
self._thread.setOutputDir(self.output_dir)
|
||||
self._thread.setImagePath(self.imagePaths)
|
||||
self._thread.setLang(lang)
|
||||
self._thread.setPDFParser(pdfParser)
|
||||
# disable buttons
|
||||
self.openFileButton.setEnabled(False)
|
||||
self.startCNButton.setEnabled(False)
|
||||
self.startENButton.setEnabled(False)
|
||||
self.PDFParserButton.setEnabled(False)
|
||||
# 启动工作进程
|
||||
self._thread.start()
|
||||
self.time_start = time.time() # log start time
|
||||
QMessageBox.information(self, "Information", "开始转换")
|
||||
else:
|
||||
QMessageBox.warning(self, "Information", "请选择要识别的文件或截图")
|
||||
|
||||
def handleShowResultSignal(self):
|
||||
if self.output_dir is None:
|
||||
return
|
||||
if os.path.exists(self.output_dir):
|
||||
if platform.system() == "Windows":
|
||||
os.startfile(self.output_dir)
|
||||
else:
|
||||
os.system("open " + os.path.normpath(self.output_dir))
|
||||
else:
|
||||
QMessageBox.information(self, "Information", "输出文件不存在")
|
||||
|
||||
def handleProgressBarUpdateSingal(self, i):
|
||||
self.pb.setValue(i)
|
||||
# calculate time left of recognition
|
||||
lenbar = self.pb.maximum()
|
||||
avg_time = (
|
||||
time.time() - self.time_start
|
||||
) / i # Use average time to prevent time fluctuations
|
||||
time_left = str(datetime.timedelta(seconds=avg_time * (lenbar - i))).split(".")[
|
||||
0
|
||||
] # Remove microseconds
|
||||
self.timeEstLabel.setText(f"Time Left: {time_left}") # show time left
|
||||
|
||||
def handleProgressBarRangeSingal(self, max):
|
||||
self.pb.setRange(0, max)
|
||||
|
||||
def handleEndsignalSignal(self):
|
||||
# enable buttons
|
||||
self.openFileButton.setEnabled(True)
|
||||
self.startCNButton.setEnabled(True)
|
||||
self.startENButton.setEnabled(True)
|
||||
self.PDFParserButton.setEnabled(True)
|
||||
QMessageBox.information(self, "Information", "转换结束")
|
||||
|
||||
def handleCBChangeSignal(self):
|
||||
self._thread.setPDFParser(self.checkBox.isChecked())
|
||||
|
||||
def handleThreadException(self, message):
|
||||
self._thread.quit()
|
||||
QMessageBox.information(self, "Error", message)
|
||||
|
||||
|
||||
def main():
|
||||
app = QApplication(sys.argv)
|
||||
|
||||
window = APP_Image2Doc() # 创建对象
|
||||
window.show() # 全屏显示窗口
|
||||
|
||||
QApplication.processEvents()
|
||||
sys.exit(app.exec())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
415
ppstructure/predict_system.py
Normal file
415
ppstructure/predict_system.py
Normal file
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) 2020 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
|
||||
import subprocess
|
||||
|
||||
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
||||
sys.path.append(__dir__)
|
||||
sys.path.insert(0, 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 logging
|
||||
from copy import deepcopy
|
||||
|
||||
from paddle.utils import try_import
|
||||
from ppocr.utils.utility import get_image_file_list, check_and_read
|
||||
from ppocr.utils.logging import get_logger
|
||||
from ppocr.utils.visual import draw_ser_results, draw_re_results
|
||||
from tools.infer.predict_system import TextSystem
|
||||
from tools.infer.predict_rec import TextRecognizer
|
||||
from ppstructure.layout.predict_layout import LayoutPredictor
|
||||
from ppstructure.table.predict_table import TableSystem, to_excel
|
||||
from ppstructure.utility import parse_args, draw_structure_result, cal_ocr_word_box
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class StructureSystem(object):
|
||||
def __init__(self, args):
|
||||
self.mode = args.mode
|
||||
self.recovery = args.recovery
|
||||
|
||||
self.image_orientation_predictor = None
|
||||
if args.image_orientation:
|
||||
import paddleclas
|
||||
|
||||
self.image_orientation_predictor = paddleclas.PaddleClas(
|
||||
model_name="text_image_orientation"
|
||||
)
|
||||
|
||||
if self.mode == "structure":
|
||||
if not args.show_log:
|
||||
logger.setLevel(logging.INFO)
|
||||
if args.layout == False and args.ocr == True:
|
||||
args.ocr = False
|
||||
logger.warning(
|
||||
"When args.layout is false, args.ocr is automatically set to false"
|
||||
)
|
||||
# init model
|
||||
self.layout_predictor = None
|
||||
self.text_system = None
|
||||
self.table_system = None
|
||||
self.formula_system = None
|
||||
if args.layout:
|
||||
self.layout_predictor = LayoutPredictor(args)
|
||||
if args.ocr:
|
||||
self.text_system = TextSystem(args)
|
||||
if args.table:
|
||||
if self.text_system is not None:
|
||||
self.table_system = TableSystem(
|
||||
args,
|
||||
self.text_system.text_detector,
|
||||
self.text_system.text_recognizer,
|
||||
)
|
||||
else:
|
||||
self.table_system = TableSystem(args)
|
||||
if args.formula:
|
||||
args_formula = deepcopy(args)
|
||||
args_formula.rec_algorithm = args.formula_algorithm
|
||||
args_formula.rec_model_dir = args.formula_model_dir
|
||||
args_formula.rec_char_dict_path = args.formula_char_dict_path
|
||||
args_formula.rec_batch_num = args.formula_batch_num
|
||||
self.formula_system = TextRecognizer(args_formula)
|
||||
|
||||
elif self.mode == "kie":
|
||||
from ppstructure.kie.predict_kie_token_ser_re import SerRePredictor
|
||||
|
||||
self.kie_predictor = SerRePredictor(args)
|
||||
|
||||
self.return_word_box = args.return_word_box
|
||||
|
||||
def __call__(self, img, return_ocr_result_in_table=False, img_idx=0):
|
||||
time_dict = {
|
||||
"image_orientation": 0,
|
||||
"layout": 0,
|
||||
"table": 0,
|
||||
"table_match": 0,
|
||||
"formula": 0,
|
||||
"det": 0,
|
||||
"rec": 0,
|
||||
"kie": 0,
|
||||
"all": 0,
|
||||
}
|
||||
start = time.time()
|
||||
|
||||
if self.image_orientation_predictor is not None:
|
||||
tic = time.time()
|
||||
cls_result = self.image_orientation_predictor.predict(input_data=img)
|
||||
cls_res = next(cls_result)
|
||||
angle = cls_res[0]["label_names"][0]
|
||||
cv_rotate_code = {
|
||||
"90": cv2.ROTATE_90_COUNTERCLOCKWISE,
|
||||
"180": cv2.ROTATE_180,
|
||||
"270": cv2.ROTATE_90_CLOCKWISE,
|
||||
}
|
||||
if angle in cv_rotate_code:
|
||||
img = cv2.rotate(img, cv_rotate_code[angle])
|
||||
toc = time.time()
|
||||
time_dict["image_orientation"] = toc - tic
|
||||
|
||||
if self.mode == "structure":
|
||||
ori_im = img.copy()
|
||||
if self.layout_predictor is not None:
|
||||
layout_res, elapse = self.layout_predictor(img)
|
||||
time_dict["layout"] += elapse
|
||||
else:
|
||||
h, w = ori_im.shape[:2]
|
||||
layout_res = [dict(bbox=None, label="table", score=0.0)]
|
||||
|
||||
# As reported in issues such as #10270 and #11665, the old
|
||||
# implementation, which recognizes texts from the layout regions,
|
||||
# has problems with OCR recognition accuracy.
|
||||
#
|
||||
# To enhance the OCR recognition accuracy, we implement a patch fix
|
||||
# that first use text_system to detect and recognize all text information
|
||||
# and then filter out relevant texts according to the layout regions.
|
||||
text_res = None
|
||||
if self.text_system is not None:
|
||||
text_res, ocr_time_dict = self._predict_text(img)
|
||||
time_dict["det"] += ocr_time_dict["det"]
|
||||
time_dict["rec"] += ocr_time_dict["rec"]
|
||||
|
||||
res_list = []
|
||||
for region in layout_res:
|
||||
res = ""
|
||||
if region["bbox"] is not None:
|
||||
x1, y1, x2, y2 = region["bbox"]
|
||||
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||||
roi_img = ori_im[y1:y2, x1:x2, :]
|
||||
else:
|
||||
x1, y1, x2, y2 = 0, 0, w, h
|
||||
roi_img = ori_im
|
||||
bbox = [x1, y1, x2, y2]
|
||||
|
||||
if region["label"] == "table":
|
||||
if self.table_system is not None:
|
||||
res, table_time_dict = self.table_system(
|
||||
roi_img, return_ocr_result_in_table
|
||||
)
|
||||
time_dict["table"] += table_time_dict["table"]
|
||||
time_dict["table_match"] += table_time_dict["match"]
|
||||
time_dict["det"] += table_time_dict["det"]
|
||||
time_dict["rec"] += table_time_dict["rec"]
|
||||
|
||||
elif region["label"] == "equation" and self.formula_system is not None:
|
||||
latex_res, formula_time = self.formula_system([roi_img])
|
||||
time_dict["formula"] += formula_time
|
||||
res = {"latex": latex_res[0]}
|
||||
|
||||
else:
|
||||
if text_res is not None:
|
||||
# Filter the text results whose regions intersect with the current layout bbox.
|
||||
res = self._filter_text_res(text_res, bbox)
|
||||
|
||||
res_list.append(
|
||||
{
|
||||
"type": region["label"].lower(),
|
||||
"bbox": bbox,
|
||||
"img": roi_img,
|
||||
"res": res,
|
||||
"img_idx": img_idx,
|
||||
"score": region["score"],
|
||||
}
|
||||
)
|
||||
|
||||
end = time.time()
|
||||
time_dict["all"] = end - start
|
||||
return res_list, time_dict
|
||||
|
||||
elif self.mode == "kie":
|
||||
re_res, elapse = self.kie_predictor(img)
|
||||
time_dict["kie"] = elapse
|
||||
time_dict["all"] = elapse
|
||||
return re_res[0], time_dict
|
||||
|
||||
return None, None
|
||||
|
||||
def _predict_text(self, img):
|
||||
filter_boxes, filter_rec_res, ocr_time_dict = self.text_system(img)
|
||||
|
||||
# remove style char,
|
||||
# when using the recognition model trained on the PubtabNet dataset,
|
||||
# it will recognize the text format in the table, such as <b>
|
||||
style_token = [
|
||||
"<strike>",
|
||||
"<strike>",
|
||||
"<sup>",
|
||||
"</sub>",
|
||||
"<b>",
|
||||
"</b>",
|
||||
"<sub>",
|
||||
"</sup>",
|
||||
"<overline>",
|
||||
"</overline>",
|
||||
"<underline>",
|
||||
"</underline>",
|
||||
"<i>",
|
||||
"</i>",
|
||||
]
|
||||
res = []
|
||||
for box, rec_res in zip(filter_boxes, filter_rec_res):
|
||||
rec_str, rec_conf = rec_res[0], rec_res[1]
|
||||
for token in style_token:
|
||||
if token in rec_str:
|
||||
rec_str = rec_str.replace(token, "")
|
||||
if self.return_word_box:
|
||||
word_box_content_list, word_box_list = cal_ocr_word_box(
|
||||
rec_str, box, rec_res[2]
|
||||
)
|
||||
res.append(
|
||||
{
|
||||
"text": rec_str,
|
||||
"confidence": float(rec_conf),
|
||||
"text_region": box.tolist(),
|
||||
"text_word": word_box_content_list,
|
||||
"text_word_region": word_box_list,
|
||||
}
|
||||
)
|
||||
else:
|
||||
res.append(
|
||||
{
|
||||
"text": rec_str,
|
||||
"confidence": float(rec_conf),
|
||||
"text_region": box.tolist(),
|
||||
}
|
||||
)
|
||||
return res, ocr_time_dict
|
||||
|
||||
def _filter_text_res(self, text_res, bbox):
|
||||
res = []
|
||||
for r in text_res:
|
||||
box = r["text_region"]
|
||||
rect = box[0][0], box[0][1], box[2][0], box[2][1]
|
||||
if self._has_intersection(bbox, rect):
|
||||
res.append(r)
|
||||
return res
|
||||
|
||||
def _has_intersection(self, rect1, rect2):
|
||||
x_min1, y_min1, x_max1, y_max1 = rect1
|
||||
x_min2, y_min2, x_max2, y_max2 = rect2
|
||||
if x_min1 > x_max2 or x_max1 < x_min2:
|
||||
return False
|
||||
if y_min1 > y_max2 or y_max1 < y_min2:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def save_structure_res(res, save_folder, img_name, img_idx=0):
|
||||
excel_save_folder = os.path.join(save_folder, img_name)
|
||||
os.makedirs(excel_save_folder, exist_ok=True)
|
||||
res_cp = deepcopy(res)
|
||||
# save res
|
||||
with open(
|
||||
os.path.join(excel_save_folder, "res_{}.txt".format(img_idx)),
|
||||
"w",
|
||||
encoding="utf8",
|
||||
) as f:
|
||||
for region in res_cp:
|
||||
roi_img = region.pop("img")
|
||||
f.write("{}\n".format(json.dumps(region)))
|
||||
|
||||
if (
|
||||
region["type"].lower() == "table"
|
||||
and len(region["res"]) > 0
|
||||
and "html" in region["res"]
|
||||
):
|
||||
excel_path = os.path.join(
|
||||
excel_save_folder, "{}_{}.xlsx".format(region["bbox"], img_idx)
|
||||
)
|
||||
to_excel(region["res"]["html"], excel_path)
|
||||
elif region["type"].lower() == "figure":
|
||||
img_path = os.path.join(
|
||||
excel_save_folder, "{}_{}.jpg".format(region["bbox"], img_idx)
|
||||
)
|
||||
cv2.imwrite(img_path, roi_img)
|
||||
|
||||
|
||||
def main(args):
|
||||
image_file_list = get_image_file_list(args.image_dir)
|
||||
image_file_list = image_file_list
|
||||
image_file_list = image_file_list[args.process_id :: args.total_process_num]
|
||||
|
||||
if not args.use_pdf2docx_api:
|
||||
structure_sys = StructureSystem(args)
|
||||
save_folder = os.path.join(args.output, structure_sys.mode)
|
||||
os.makedirs(save_folder, exist_ok=True)
|
||||
img_num = len(image_file_list)
|
||||
|
||||
for i, image_file in enumerate(image_file_list):
|
||||
logger.info("[{}/{}] {}".format(i, img_num, image_file))
|
||||
img, flag_gif, flag_pdf = check_and_read(image_file)
|
||||
img_name = os.path.basename(image_file).split(".")[0]
|
||||
|
||||
if args.recovery and args.use_pdf2docx_api and flag_pdf:
|
||||
try_import("pdf2docx")
|
||||
from pdf2docx.converter import Converter
|
||||
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
docx_file = os.path.join(args.output, "{}_api.docx".format(img_name))
|
||||
cv = Converter(image_file)
|
||||
cv.convert(docx_file)
|
||||
cv.close()
|
||||
logger.info("docx save to {}".format(docx_file))
|
||||
continue
|
||||
|
||||
if not flag_gif and not flag_pdf:
|
||||
img = cv2.imread(image_file)
|
||||
|
||||
if not flag_pdf:
|
||||
if img is None:
|
||||
logger.error("error in loading image:{}".format(image_file))
|
||||
continue
|
||||
imgs = [img]
|
||||
else:
|
||||
imgs = img
|
||||
|
||||
all_res = []
|
||||
for index, img in enumerate(imgs):
|
||||
res, time_dict = structure_sys(img, img_idx=index)
|
||||
img_save_path = os.path.join(
|
||||
save_folder, img_name, "show_{}.jpg".format(index)
|
||||
)
|
||||
os.makedirs(os.path.join(save_folder, img_name), exist_ok=True)
|
||||
if structure_sys.mode == "structure" and res != []:
|
||||
draw_img = draw_structure_result(img, res, args.vis_font_path)
|
||||
save_structure_res(res, save_folder, img_name, index)
|
||||
elif structure_sys.mode == "kie":
|
||||
if structure_sys.kie_predictor.predictor is not None:
|
||||
draw_img = draw_re_results(img, res, font_path=args.vis_font_path)
|
||||
else:
|
||||
draw_img = draw_ser_results(img, res, font_path=args.vis_font_path)
|
||||
|
||||
with open(
|
||||
os.path.join(save_folder, img_name, "res_{}_kie.txt".format(index)),
|
||||
"w",
|
||||
encoding="utf8",
|
||||
) as f:
|
||||
res_str = "{}\t{}\n".format(
|
||||
image_file, json.dumps({"ocr_info": res}, ensure_ascii=False)
|
||||
)
|
||||
f.write(res_str)
|
||||
if res != []:
|
||||
cv2.imwrite(img_save_path, draw_img)
|
||||
logger.info("result save to {}".format(img_save_path))
|
||||
if args.recovery and res != []:
|
||||
from ppstructure.recovery.recovery_to_doc import (
|
||||
sorted_layout_boxes,
|
||||
convert_info_docx,
|
||||
)
|
||||
from ppstructure.recovery.recovery_to_markdown import (
|
||||
convert_info_markdown,
|
||||
)
|
||||
|
||||
h, w, _ = img.shape
|
||||
res = sorted_layout_boxes(res, w)
|
||||
all_res += res
|
||||
|
||||
if args.recovery and all_res != []:
|
||||
try:
|
||||
convert_info_docx(img, all_res, save_folder, img_name)
|
||||
if args.recovery_to_markdown:
|
||||
convert_info_markdown(all_res, save_folder, img_name)
|
||||
except Exception as ex:
|
||||
logger.error(
|
||||
"error in layout recovery image:{}, err msg: {}".format(
|
||||
image_file, ex
|
||||
)
|
||||
)
|
||||
continue
|
||||
logger.info("Predict time : {:.3f}s".format(time_dict["all"]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
if args.use_mp:
|
||||
p_list = []
|
||||
total_process_num = args.total_process_num
|
||||
for process_id in range(total_process_num):
|
||||
cmd = (
|
||||
[sys.executable, "-u"]
|
||||
+ sys.argv
|
||||
+ ["--process_id={}".format(process_id), "--use_mp={}".format(False)]
|
||||
)
|
||||
p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
|
||||
p_list.append(p)
|
||||
for p in p_list:
|
||||
p.wait()
|
||||
else:
|
||||
main(args)
|
||||
213
ppstructure/recovery/README.md
Normal file
213
ppstructure/recovery/README.md
Normal file
@@ -0,0 +1,213 @@
|
||||
English | [简体中文](README_ch.md)
|
||||
|
||||
# Layout Recovery
|
||||
|
||||
- [1. Introduction](#1)
|
||||
- [2. Install](#2)
|
||||
- [2.1 Install PaddlePaddle](#2.1)
|
||||
- [2.2 Install PaddleOCR](#2.2)
|
||||
- [3. Quick Start using standard PDF parse](#3)
|
||||
- [4. Quick Start using image format PDF parse ](#4)
|
||||
- [4.1 Download models](#4.1)
|
||||
- [4.2 Layout recovery](#4.2)
|
||||
- [5. More](#5)
|
||||
|
||||
<a name="1"></a>
|
||||
|
||||
## 1. Introduction
|
||||
|
||||
The layout recovery module is used to restore the image or pdf to an
|
||||
editable Word file consistent with the original image layout.
|
||||
|
||||
Two layout recovery methods are provided, you can choose by PDF format:
|
||||
|
||||
- **Standard PDF parse(the input is standard PDF)**: Python based PDF to word library [pdf2docx] (https://github.com/dothinking/pdf2docx) is optimized, the method extracts data from PDF with PyMuPDF, then parse layout with rule, finally, generate docx with python-docx.
|
||||
|
||||
- **Image format PDF parse(the input can be standard PDF or image format PDF)**: Layout recovery combines [layout analysis](../layout/README.md)、[table recognition](../table/README.md) to better recover images, tables, titles, etc. supports input files in PDF and document image formats in Chinese and English.
|
||||
|
||||
The input formats and application scenarios of the two methods are as follows:
|
||||
|
||||
| method | input formats | application scenarios/problem |
|
||||
| :-----: | :----------: | :----------------------------------------------------------: |
|
||||
| Standard PDF parse | pdf | Advantages: Better recovery for non-paper documents, each page remains on the same page after restoration<br>Disadvantages: English characters in some Chinese documents are garbled, some contents are still beyond the current page, the whole page content is restored to the table format, and the recovery effect of some pictures is not good |
|
||||
| Image format PDF parse( | pdf、picture | Advantages: More suitable for paper document content recovery, OCR recognition effect is more good<br>Disadvantages: Currently, the recovery is based on rules, the effect of content typesetting (spacing, fonts, etc.) need to be further improved, and the effect of layout recovery depends on layout analysis |
|
||||
|
||||
The following figure shows the effect of restoring the layout of documents by using PDF parse:
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/19808900/195319853-045123c9-f542-4596-b4e4-6081708dfc56.png" width = "700" />
|
||||
</div>
|
||||
|
||||
The following figures show the effect of restoring the layout of English and Chinese documents by using OCR technique:
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/recovery/recovery.jpg" width = "700" />
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/recovery/recovery_ch.jpg" width = "800" />
|
||||
</div>
|
||||
|
||||
|
||||
<a name="2"></a>
|
||||
|
||||
## 2. Install
|
||||
|
||||
<a name="2.1"></a>
|
||||
|
||||
### 2.1 Install PaddlePaddle
|
||||
|
||||
```bash
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
# If you have cuda9 or cuda10 installed on your machine, please run the following command to install
|
||||
python3 -m pip install "paddlepaddle-gpu" -i https://mirror.baidu.com/pypi/simple
|
||||
|
||||
# CPU installation
|
||||
python3 -m pip install "paddlepaddle" -i https://mirror.baidu.com/pypi/simple
|
||||
````
|
||||
|
||||
For more requirements, please refer to the instructions in [Installation Documentation](https://www.paddlepaddle.org.cn/en/install/quick?docurl=/documentation/docs/en/install/pip/macos-pip_en.html).
|
||||
|
||||
<a name="2.2"></a>
|
||||
|
||||
### 2.2 Install PaddleOCR
|
||||
|
||||
- **(1) Download source code**
|
||||
|
||||
```bash
|
||||
[Recommended] git clone https://github.com/PaddlePaddle/PaddleOCR
|
||||
|
||||
# If the pull cannot be successful due to network problems, you can also choose to use the hosting on the code cloud:
|
||||
git clone https://gitee.com/paddlepaddle/PaddleOCR
|
||||
|
||||
# Note: Code cloud hosting code may not be able to synchronize the update of this github project in real time, there is a delay of 3 to 5 days, please use the recommended method first.
|
||||
````
|
||||
|
||||
- **(2) Install recovery `requirements`**
|
||||
|
||||
The layout restoration is exported as docx files, so python-docx API need to be installed, and PyMuPDF api([requires Python >= 3.7](https://pypi.org/project/PyMuPDF/)) need to be installed to process the input files in pdf format.
|
||||
|
||||
Install all the libraries by running the following command:
|
||||
|
||||
```bash
|
||||
python3 -m pip install -r ppstructure/recovery/requirements.txt
|
||||
````
|
||||
|
||||
And if using pdf parse method, we need to install pdf2docx api.
|
||||
|
||||
```bash
|
||||
wget https://paddleocr.bj.bcebos.com/whl/pdf2docx-0.0.0-py3-none-any.whl
|
||||
pip3 install pdf2docx-0.0.0-py3-none-any.whl
|
||||
```
|
||||
|
||||
<a name="3"></a>
|
||||
|
||||
## 3. Quick Start using standard PDF parse
|
||||
|
||||
`use_pdf2docx_api` use PDF parse for layout recovery, The whl package is also provided for quick use, follow the above code, for more information please refer to [quickstart](../docs/quickstart_en.md) for details.
|
||||
|
||||
```bash
|
||||
# install paddleocr
|
||||
pip3 install "paddleocr>=2.6"
|
||||
paddleocr --image_dir=ppstructure/docs/recovery/UnrealText.pdf --type=structure --recovery=true --use_pdf2docx_api=true
|
||||
```
|
||||
|
||||
Command line:
|
||||
|
||||
```bash
|
||||
python3 predict_system.py \
|
||||
--image_dir=ppstructure/docs/recovery/UnrealText.pdf \
|
||||
--recovery=True \
|
||||
--use_pdf2docx_api=True \
|
||||
--output=../output/
|
||||
```
|
||||
|
||||
<a name="4"></a>
|
||||
## 4. Quick Start using image format PDF parse
|
||||
|
||||
Through layout analysis, we divided the image/PDF documents into regions, located the key regions, such as text, table, picture, etc., and recorded the location, category, and regional pixel value information of each region. Different regions are processed separately, where:
|
||||
|
||||
- OCR detection and recognition is performed in the text area, and the coordinates of the OCR detection box and the text content information are added on the basis of the previous information
|
||||
|
||||
- The table area identifies tables and records html and text information of tables
|
||||
- Save the image directly
|
||||
|
||||
We can restore the test picture through the layout information, OCR detection and recognition structure, table information, and saved pictures.
|
||||
|
||||
The whl package is also provided for quick use, follow the above code, for more information please refer to [quickstart](../docs/quickstart_en.md) for details.
|
||||
|
||||
```bash
|
||||
paddleocr --image_dir=ppstructure/docs/table/1.png --type=structure --recovery=true --lang='en'
|
||||
```
|
||||
|
||||
<a name="4.1"></a>
|
||||
### 4.1 Download models
|
||||
|
||||
If input is English document, download English models:
|
||||
|
||||
```bash
|
||||
cd PaddleOCR/ppstructure
|
||||
|
||||
# download model
|
||||
mkdir inference && cd inference
|
||||
# Download the detection model of the ultra-lightweight English PP-OCRv3 model and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar && tar xf en_PP-OCRv3_det_infer.tar
|
||||
# Download the recognition model of the ultra-lightweight English PP-OCRv3 model and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar && tar xf en_PP-OCRv3_rec_infer.tar
|
||||
# Download the ultra-lightweight English table inch model and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
tar xf en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
# Download the layout model of publaynet dataset and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
tar xf picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
cd ..
|
||||
```
|
||||
If input is Chinese document，download Chinese models:
|
||||
[Chinese and English ultra-lightweight PP-OCRv3 model](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/README.md#pp-ocr-series-model-listupdate-on-september-8th)、[table recognition model](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md#22-表格识别模型)、[layout analysis model](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md#1-版面分析模型)
|
||||
|
||||
<a name="4.2"></a>
|
||||
### 4.2 Layout recovery
|
||||
|
||||
|
||||
```bash
|
||||
python3 predict_system.py \
|
||||
--image_dir=./docs/table/1.png \
|
||||
--det_model_dir=inference/en_PP-OCRv3_det_infer \
|
||||
--rec_model_dir=inference/en_PP-OCRv3_rec_infer \
|
||||
--rec_char_dict_path=../ppocr/utils/en_dict.txt \
|
||||
--table_model_dir=inference/en_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--layout_model_dir=inference/picodet_lcnet_x1_0_fgd_layout_infer \
|
||||
--layout_dict_path=../ppocr/utils/dict/layout_dict/layout_publaynet_dict.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--recovery=True \
|
||||
--output=../output/
|
||||
```
|
||||
|
||||
After running, the docx of each picture will be saved in the directory specified by the output field
|
||||
|
||||
Field：
|
||||
|
||||
- image_dir：test file， can be picture, picture directory, pdf file, pdf file directory
|
||||
- det_model_dir：OCR detection model path
|
||||
- rec_model_dir：OCR recognition model path
|
||||
- rec_char_dict_path：OCR recognition dict path. If the Chinese model is used, change to "../ppocr/utils/ppocr_keys_v1.txt". And if you trained the model on your own dataset, change to the trained dictionary
|
||||
- table_model_dir：table recognition model path
|
||||
- table_char_dict_path：table recognition dict path. If the Chinese model is used, no need to change
|
||||
- layout_model_dir：layout analysis model path
|
||||
- layout_dict_path：layout analysis dict path. If the Chinese model is used, change to "../ppocr/utils/dict/layout_dict/layout_cdla_dict.txt"
|
||||
- recovery：whether to enable layout of recovery, default False
|
||||
- output：save the recovery result path
|
||||
|
||||
<a name="5"></a>
|
||||
|
||||
## 5. More
|
||||
|
||||
For training, evaluation and inference tutorial for text detection models, please refer to [text detection doc](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/doc/doc_en/detection_en.md).
|
||||
|
||||
For training, evaluation and inference tutorial for text recognition models, please refer to [text recognition doc](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/doc/doc_en/recognition_en.md).
|
||||
|
||||
For training, evaluation and inference tutorial for layout analysis models, please refer to [layout analysis doc](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/layout/README.md)
|
||||
|
||||
For training, evaluation and inference tutorial for table recognition models, please refer to [table recognition doc](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/table/README.md)
|
||||
222
ppstructure/recovery/README_ch.md
Normal file
222
ppstructure/recovery/README_ch.md
Normal file
@@ -0,0 +1,222 @@
|
||||
[English](README.md) | 简体中文
|
||||
|
||||
# 版面恢复
|
||||
|
||||
- [1. 简介](#1)
|
||||
- [2. 安装](#2)
|
||||
- [2.1 安装PaddlePaddle](#2.1)
|
||||
- [2.2 安装PaddleOCR](#2.2)
|
||||
- [3.使用标准PDF解析进行版面恢复](#3)
|
||||
- [4. 使用图片格式PDF解析进行版面恢复](#4)
|
||||
- [4.1 下载模型](#4.1)
|
||||
- [4.2 版面恢复](#4.2)
|
||||
- [5. 更多](#5)
|
||||
|
||||
<a name="1"></a>
|
||||
|
||||
## 1. 简介
|
||||
|
||||
版面恢复就是将输入的图片、pdf内容仍然像原文档那样排列着，段落不变、顺序不变的输出到word文档中等。
|
||||
|
||||
提供了2种版面恢复方法，可根据输入PDF的格式进行选择：
|
||||
|
||||
- **标准PDF解析(输入须为标准PDF)**：基于Python的pdf转word库[pdf2docx](https://github.com/dothinking/pdf2docx)进行优化，该方法通过PyMuPDF获取页面元素，然后利用规则解析章节、段落、表格等布局及样式，最后通过python-docx将解析的内容元素重建到word文档中。
|
||||
- **图片格式PDF解析(输入可为标准PDF或图片格式PDF)**：结合[版面分析](../layout/README_ch.md)、[表格识别](../table/README_ch.md)技术，从而更好地恢复图片、表格、标题等内容，支持中、英文pdf文档、文档图片格式的输入文件。
|
||||
|
||||
2种方法输入格式、适用场景如下：
|
||||
|
||||
| 方法 | 支持输入文件 | 适用场景/存在问题 |
|
||||
| :-------------: | :----------: | :----------------------------------------------------------: |
|
||||
| 标准PDF解析 | pdf | 优点：非论文文档恢复效果更优、每一页内容恢复后仍在同一页<br>缺点：有些中文文档中的英文乱码、仍存在内容超出当前页面的情况、整页内容恢复为表格格式、部分图片恢复效果不佳 |
|
||||
| 图片格式PDF解析 | pdf、图片 | 优点：更适合论文文档正文内容的恢复、中英文文档OCR识别效果好<br>缺点：目前内容恢复基于规则，内容排版效果(间距、字体等)待进一步提升、版面恢复效果依赖于版面分析效果 |
|
||||
|
||||
下图展示了通过PDF解析版面恢复效果：
|
||||
|
||||
<div align="center">
|
||||
<img src="https://user-images.githubusercontent.com/19808900/195319840-68fc60ec-ea66-4095-b734-0ec115860341.png" width = "700" />
|
||||
</div>
|
||||
|
||||
下图分别展示了通过OCR技术，英文文档和中文文档版面恢复的效果：
|
||||
|
||||
<div align="center">
|
||||
<img src="../docs/recovery/recovery.jpg" width = "700" />
|
||||
</div>
|
||||
<div align="center">
|
||||
<img src="../docs/recovery/recovery_ch.jpg" width = "800" />
|
||||
</div>
|
||||
<a name="2"></a>
|
||||
|
||||
## 2. 安装
|
||||
|
||||
<a name="2.1"></a>
|
||||
|
||||
### 2.1 安装PaddlePaddle
|
||||
|
||||
```bash
|
||||
python3 -m pip install --upgrade pip
|
||||
|
||||
# 您的机器安装的是CUDA9或CUDA10，请运行以下命令安装
|
||||
python3 -m pip install "paddlepaddle-gpu" -i https://mirror.baidu.com/pypi/simple
|
||||
|
||||
# 您的机器是CPU，请运行以下命令安装
|
||||
python3 -m pip install "paddlepaddle" -i https://mirror.baidu.com/pypi/simple
|
||||
|
||||
```
|
||||
|
||||
更多需求，请参照[安装文档](https://www.paddlepaddle.org.cn/install/quick)中的说明进行操作。
|
||||
|
||||
<a name="2.2"></a>
|
||||
|
||||
### 2.2 安装PaddleOCR
|
||||
|
||||
- **（1）下载版面恢复源码**
|
||||
|
||||
```bash
|
||||
【推荐】git clone https://github.com/PaddlePaddle/PaddleOCR
|
||||
|
||||
# 如果因为网络问题无法pull成功，也可选择使用码云上的托管：
|
||||
git clone https://gitee.com/paddlepaddle/PaddleOCR
|
||||
|
||||
# 注：码云托管代码可能无法实时同步本github项目更新，存在3~5天延时，请优先使用推荐方式。
|
||||
```
|
||||
|
||||
- **（2）安装recovery的`requirements`**
|
||||
|
||||
版面恢复导出为docx文件，所以需要安装Python处理word文档的python-docx API，同时处理pdf格式的输入文件，需要安装PyMuPDF API([要求Python >= 3.7](https://pypi.org/project/PyMuPDF/))。
|
||||
|
||||
通过如下命令安装全部库：
|
||||
|
||||
```bash
|
||||
python3 -m pip install -r ppstructure/recovery/requirements.txt
|
||||
```
|
||||
|
||||
使用pdf2docx库解析的方式恢复文档需要安装优化的pdf2docx。
|
||||
|
||||
```bash
|
||||
wget https://paddleocr.bj.bcebos.com/whl/pdf2docx-0.0.0-py3-none-any.whl
|
||||
pip3 install pdf2docx-0.0.0-py3-none-any.whl
|
||||
```
|
||||
|
||||
<a name="3"></a>
|
||||
|
||||
## 3.使用标准PDF解析进行版面恢复
|
||||
|
||||
`use_pdf2docx_api`表示使用PDF解析的方式进行版面恢复，通过whl包的形式方便快速使用，代码如下，更多信息详见 [quickstart](../docs/quickstart.md)。
|
||||
|
||||
```bash
|
||||
# 安装 paddleocr，推荐使用2.6版本
|
||||
pip3 install "paddleocr>=2.6"
|
||||
paddleocr --image_dir=ppstructure/docs/recovery/UnrealText.pdf --type=structure --recovery=true --use_pdf2docx_api=true
|
||||
```
|
||||
|
||||
通过命令行的方式：
|
||||
|
||||
```bash
|
||||
python3 predict_system.py \
|
||||
--image_dir=ppstructure/docs/recovery/UnrealText.pdf \
|
||||
--recovery=True \
|
||||
--use_pdf2docx_api=True \
|
||||
--output=../output/
|
||||
```
|
||||
|
||||
<a name="4"></a>
|
||||
|
||||
## 4.使用图片格式PDF解析进行版面恢复
|
||||
|
||||
我们通过版面分析对图片/pdf形式的文档进行区域划分，定位其中的关键区域，如文字、表格、图片等，记录每个区域的位置、类别、区域像素值信息。对不同的区域分别处理，其中：
|
||||
|
||||
- 文字区域直接进行OCR检测和识别，在之前信息基础上增加OCR检测框坐标和文本内容信息
|
||||
|
||||
- 表格区域进行表格识别，记录表格html和文字信息
|
||||
- 图片直接保存
|
||||
|
||||
我们通过版面信息、OCR检测和识别结构、表格信息、保存的图片，对测试图片进行恢复即可。
|
||||
|
||||
提供如下代码实现版面恢复，也提供了whl包的形式方便快速使用，代码如下，更多信息详见 [quickstart](../docs/quickstart.md)。
|
||||
|
||||
```bash
|
||||
# 安装 paddleocr，推荐使用2.6版本
|
||||
pip3 install "paddleocr>=2.6"
|
||||
# 中文测试图
|
||||
paddleocr --image_dir=ppstructure/docs/table/1.png --type=structure --recovery=true
|
||||
# 英文测试图
|
||||
paddleocr --image_dir=ppstructure/docs/table/1.png --type=structure --recovery=true --lang='en'
|
||||
# pdf测试文件
|
||||
paddleocr --image_dir=ppstructure/docs/recovery/UnrealText.pdf --type=structure --recovery=true --lang='en'
|
||||
```
|
||||
|
||||
<a name="4.1"></a>
|
||||
|
||||
### 4.1 下载模型
|
||||
|
||||
如果输入为英文文档类型，下载OCR检测和识别、版面分析、表格识别的英文模型
|
||||
|
||||
```bash
|
||||
cd PaddleOCR/ppstructure
|
||||
|
||||
# 下载模型
|
||||
mkdir inference && cd inference
|
||||
# 下载英文超轻量PP-OCRv3检测模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_det_infer.tar && tar xf en_PP-OCRv3_det_infer.tar
|
||||
# 下载英文超轻量PP-OCRv3识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/english/en_PP-OCRv3_rec_infer.tar && tar xf en_PP-OCRv3_rec_infer.tar
|
||||
# 下载英文表格识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
tar xf en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
# 下载英文版面分析模型
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/layout/picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
tar xf picodet_lcnet_x1_0_fgd_layout_infer.tar
|
||||
cd ..
|
||||
```
|
||||
|
||||
如果输入为中文文档类型，在下述链接中下载中文模型即可：
|
||||
|
||||
[PP-OCRv3中英文超轻量文本检测和识别模型](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/README_ch.md#pp-ocr%E7%B3%BB%E5%88%97%E6%A8%A1%E5%9E%8B%E5%88%97%E8%A1%A8%E6%9B%B4%E6%96%B0%E4%B8%AD)、[表格识别模型](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md#22-表格识别模型)、[版面分析模型](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/docs/models_list.md#1-版面分析模型)
|
||||
|
||||
<a name="4.2"></a>
|
||||
|
||||
### 4.2 版面恢复
|
||||
|
||||
使用下载的模型恢复给定文档的版面，以英文模型为例，执行如下命令：
|
||||
|
||||
```bash
|
||||
python3 predict_system.py \
|
||||
--image_dir=./docs/table/1.png \
|
||||
--det_model_dir=inference/en_PP-OCRv3_det_infer \
|
||||
--rec_model_dir=inference/en_PP-OCRv3_rec_infer \
|
||||
--rec_char_dict_path=../ppocr/utils/en_dict.txt \
|
||||
--table_model_dir=inference/en_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--layout_model_dir=inference/picodet_lcnet_x1_0_fgd_layout_infer \
|
||||
--layout_dict_path=../ppocr/utils/dict/layout_dict/layout_publaynet_dict.txt \
|
||||
--vis_font_path=../doc/fonts/simfang.ttf \
|
||||
--recovery=True \
|
||||
--output=../output/
|
||||
```
|
||||
|
||||
运行完成后，恢复版面的docx文档会保存到`output`字段指定的目录下
|
||||
|
||||
字段含义：
|
||||
|
||||
- image_dir：测试文件，可以是图片、图片目录、pdf文件、pdf文件目录
|
||||
- det_model_dir：OCR检测模型路径
|
||||
- rec_model_dir：OCR识别模型路径
|
||||
- rec_char_dict_path：OCR识别字典，如果更换为中文模型，需要更改为"../ppocr/utils/ppocr_keys_v1.txt"，如果您在自己的数据集上训练的模型，则更改为训练的字典的文件
|
||||
- table_model_dir：表格识别模型路径
|
||||
- table_char_dict_path：表格识别字典，如果更换为中文模型，不需要更换字典
|
||||
- layout_model_dir：版面分析模型路径
|
||||
- layout_dict_path：版面分析字典，如果更换为中文模型，需要更改为"../ppocr/utils/dict/layout_dict/layout_cdla_dict.txt"
|
||||
- recovery：是否进行版面恢复，默认False
|
||||
- output：版面恢复结果保存路径
|
||||
|
||||
<a name="5"></a>
|
||||
|
||||
## 5. 更多
|
||||
|
||||
关于OCR检测模型的训练评估与推理，请参考：[文本检测教程](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/doc/doc_ch/detection.md)
|
||||
|
||||
关于OCR识别模型的训练评估与推理，请参考：[文本识别教程](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/doc/doc_ch/recognition.md)
|
||||
|
||||
关于版面分析模型的训练评估与推理，请参考：[版面分析教程](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/layout/README_ch.md)
|
||||
|
||||
关于表格识别模型的训练评估与推理，请参考：[表格识别教程](https://github.com/PaddlePaddle/PaddleOCR/blob/dygraph/ppstructure/table/README_ch.md)
|
||||
13
ppstructure/recovery/__init__.py
Normal file
13
ppstructure/recovery/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
155
ppstructure/recovery/recovery_to_doc.py
Normal file
155
ppstructure/recovery/recovery_to_doc.py
Normal file
@@ -0,0 +1,155 @@
|
||||
# Copyright (c) 2020 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
|
||||
from copy import deepcopy
|
||||
|
||||
from docx import Document
|
||||
from docx import shared
|
||||
from docx.enum.text import WD_ALIGN_PARAGRAPH
|
||||
from docx.enum.section import WD_SECTION
|
||||
from docx.oxml.ns import qn
|
||||
from docx.enum.table import WD_TABLE_ALIGNMENT
|
||||
|
||||
from ppstructure.recovery.table_process import HtmlToDocx
|
||||
|
||||
from ppocr.utils.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def convert_info_docx(img, res, save_folder, img_name):
|
||||
doc = Document()
|
||||
doc.styles["Normal"].font.name = "Times New Roman"
|
||||
doc.styles["Normal"]._element.rPr.rFonts.set(qn("w:eastAsia"), "宋体")
|
||||
doc.styles["Normal"].font.size = shared.Pt(6.5)
|
||||
|
||||
flag = 1
|
||||
for i, region in enumerate(res):
|
||||
if not region["res"] and region["type"].lower() != "figure":
|
||||
continue
|
||||
img_idx = region["img_idx"]
|
||||
if flag == 2 and region["layout"] == "single":
|
||||
section = doc.add_section(WD_SECTION.CONTINUOUS)
|
||||
section._sectPr.xpath("./w:cols")[0].set(qn("w:num"), "1")
|
||||
flag = 1
|
||||
elif flag == 1 and region["layout"] == "double":
|
||||
section = doc.add_section(WD_SECTION.CONTINUOUS)
|
||||
section._sectPr.xpath("./w:cols")[0].set(qn("w:num"), "2")
|
||||
flag = 2
|
||||
|
||||
if region["type"].lower() == "figure":
|
||||
excel_save_folder = os.path.join(save_folder, img_name)
|
||||
img_path = os.path.join(
|
||||
excel_save_folder, "{}_{}.jpg".format(region["bbox"], img_idx)
|
||||
)
|
||||
paragraph_pic = doc.add_paragraph()
|
||||
paragraph_pic.alignment = WD_ALIGN_PARAGRAPH.CENTER
|
||||
run = paragraph_pic.add_run("")
|
||||
if flag == 1:
|
||||
run.add_picture(img_path, width=shared.Inches(5))
|
||||
elif flag == 2:
|
||||
run.add_picture(img_path, width=shared.Inches(2))
|
||||
elif region["type"].lower() == "title":
|
||||
doc.add_heading(region["res"][0]["text"])
|
||||
elif region["type"].lower() == "table":
|
||||
parser = HtmlToDocx()
|
||||
parser.table_style = "TableGrid"
|
||||
parser.handle_table(region["res"]["html"], doc)
|
||||
elif region["type"] == "equation" and "latex" in region["res"]:
|
||||
pass
|
||||
else:
|
||||
paragraph = doc.add_paragraph()
|
||||
paragraph_format = paragraph.paragraph_format
|
||||
for i, line in enumerate(region["res"]):
|
||||
if i == 0:
|
||||
paragraph_format.first_line_indent = shared.Inches(0.25)
|
||||
text_run = paragraph.add_run(line["text"] + " ")
|
||||
text_run.font.size = shared.Pt(10)
|
||||
|
||||
# save to docx
|
||||
docx_path = os.path.join(save_folder, "{}_ocr.docx".format(img_name))
|
||||
doc.save(docx_path)
|
||||
logger.info("docx save to {}".format(docx_path))
|
||||
|
||||
|
||||
def sorted_layout_boxes(res, w):
|
||||
"""
|
||||
Sort text boxes in order from top to bottom, left to right
|
||||
args:
|
||||
res(list):ppstructure results
|
||||
return:
|
||||
sorted results(list)
|
||||
"""
|
||||
num_boxes = len(res)
|
||||
if num_boxes == 1:
|
||||
res[0]["layout"] = "single"
|
||||
return res
|
||||
|
||||
sorted_boxes = sorted(res, key=lambda x: (x["bbox"][1], x["bbox"][0]))
|
||||
_boxes = list(sorted_boxes)
|
||||
|
||||
new_res = []
|
||||
res_left = []
|
||||
res_right = []
|
||||
i = 0
|
||||
|
||||
while True:
|
||||
if i >= num_boxes:
|
||||
break
|
||||
if i == num_boxes - 1:
|
||||
if (
|
||||
_boxes[i]["bbox"][1] > _boxes[i - 1]["bbox"][3]
|
||||
and _boxes[i]["bbox"][0] < w / 2
|
||||
and _boxes[i]["bbox"][2] > w / 2
|
||||
):
|
||||
new_res += res_left
|
||||
new_res += res_right
|
||||
_boxes[i]["layout"] = "single"
|
||||
new_res.append(_boxes[i])
|
||||
else:
|
||||
if _boxes[i]["bbox"][2] > w / 2:
|
||||
_boxes[i]["layout"] = "double"
|
||||
res_right.append(_boxes[i])
|
||||
new_res += res_left
|
||||
new_res += res_right
|
||||
elif _boxes[i]["bbox"][0] < w / 2:
|
||||
_boxes[i]["layout"] = "double"
|
||||
res_left.append(_boxes[i])
|
||||
new_res += res_left
|
||||
new_res += res_right
|
||||
res_left = []
|
||||
res_right = []
|
||||
break
|
||||
elif _boxes[i]["bbox"][0] < w / 4 and _boxes[i]["bbox"][2] < 3 * w / 4:
|
||||
_boxes[i]["layout"] = "double"
|
||||
res_left.append(_boxes[i])
|
||||
i += 1
|
||||
elif _boxes[i]["bbox"][0] > w / 4 and _boxes[i]["bbox"][2] > w / 2:
|
||||
_boxes[i]["layout"] = "double"
|
||||
res_right.append(_boxes[i])
|
||||
i += 1
|
||||
else:
|
||||
new_res += res_left
|
||||
new_res += res_right
|
||||
_boxes[i]["layout"] = "single"
|
||||
new_res.append(_boxes[i])
|
||||
res_left = []
|
||||
res_right = []
|
||||
i += 1
|
||||
if res_left:
|
||||
new_res += res_left
|
||||
if res_right:
|
||||
new_res += res_right
|
||||
return new_res
|
||||
187
ppstructure/recovery/recovery_to_markdown.py
Normal file
187
ppstructure/recovery/recovery_to_markdown.py
Normal file
@@ -0,0 +1,187 @@
|
||||
# Copyright (c) 2024 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 re
|
||||
|
||||
from ppocr.utils.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def check_merge_method(in_region):
|
||||
"""Select the function to merge paragraph.
|
||||
|
||||
Determine the paragraph merging method based on the positional
|
||||
relationship between the text bbox and the first line of text in the text bbox.
|
||||
|
||||
Args:
|
||||
in_region: Elements with text type in the layout result.
|
||||
|
||||
Returns:
|
||||
Merge the functions of paragraph, convert_text_space_head or convert_text_space_tail.
|
||||
"""
|
||||
text_bbox = in_region["bbox"]
|
||||
text_x1 = text_bbox[0]
|
||||
frist_line_box = in_region["res"][0]["text_region"]
|
||||
point_1 = frist_line_box[0]
|
||||
point_2 = frist_line_box[2]
|
||||
frist_line_x1 = point_1[0]
|
||||
frist_line_height = abs(point_2[1] - point_1[1])
|
||||
x1_distance = frist_line_x1 - text_x1
|
||||
return (
|
||||
convert_text_space_head
|
||||
if x1_distance > frist_line_height
|
||||
else convert_text_space_tail
|
||||
)
|
||||
|
||||
|
||||
def convert_text_space_head(in_region):
|
||||
"""The function to merge paragraph.
|
||||
|
||||
The sign of dividing paragraph is that there are two spaces at the beginning.
|
||||
|
||||
Args:
|
||||
in_region: Elements with text type in the layout result.
|
||||
|
||||
Returns:
|
||||
The text content of the current text box.
|
||||
"""
|
||||
text = ""
|
||||
pre_x = None
|
||||
frist_line = True
|
||||
for i, res in enumerate(in_region["res"]):
|
||||
point1 = res["text_region"][0]
|
||||
point2 = res["text_region"][2]
|
||||
h = point2[1] - point1[1]
|
||||
|
||||
if i == 0:
|
||||
text += res["text"]
|
||||
pre_x = point1[0]
|
||||
continue
|
||||
|
||||
x1 = point1[0]
|
||||
if frist_line:
|
||||
if abs(pre_x - x1) < h:
|
||||
text += "\n\n"
|
||||
text += res["text"]
|
||||
frist_line = True
|
||||
else:
|
||||
text += res["text"]
|
||||
frist_line = False
|
||||
else:
|
||||
same_paragh = abs(pre_x - x1) < h
|
||||
if same_paragh:
|
||||
text += res["text"]
|
||||
frist_line = False
|
||||
else:
|
||||
text += "\n\n"
|
||||
text += res["text"]
|
||||
frist_line = True
|
||||
pre_x = x1
|
||||
return text
|
||||
|
||||
|
||||
def convert_text_space_tail(in_region):
|
||||
"""The function to merge paragraph.
|
||||
|
||||
The symbol for dividing paragraph is a space at the end.
|
||||
|
||||
Args:
|
||||
in_region: Elements with text type in the layout result.
|
||||
|
||||
Returns:
|
||||
The text content of the current text box.
|
||||
"""
|
||||
text = ""
|
||||
frist_line = True
|
||||
text_bbox = in_region["bbox"]
|
||||
width = text_bbox[2] - text_bbox[0]
|
||||
for i, res in enumerate(in_region["res"]):
|
||||
point1 = res["text_region"][0]
|
||||
point2 = res["text_region"][2]
|
||||
row_width = point2[0] - point1[0]
|
||||
row_height = point2[1] - point1[1]
|
||||
full_row_threshold = width - row_height
|
||||
is_full = row_width >= full_row_threshold
|
||||
|
||||
if frist_line:
|
||||
text += "\n\n"
|
||||
text += res["text"]
|
||||
else:
|
||||
text += res["text"]
|
||||
|
||||
frist_line = not is_full
|
||||
return text
|
||||
|
||||
|
||||
def convert_info_markdown(res, save_folder, img_name):
|
||||
"""Save the recognition result as a markdown file.
|
||||
|
||||
Args:
|
||||
res: Recognition result
|
||||
save_folder: Folder to save the markdown file
|
||||
img_name: PDF file or image file name
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
|
||||
def replace_special_char(content):
|
||||
special_chars = ["*", "`", "~", "$"]
|
||||
for char in special_chars:
|
||||
content = content.replace(char, "\\" + char)
|
||||
return content
|
||||
|
||||
markdown_string = []
|
||||
|
||||
for i, region in enumerate(res):
|
||||
if not region["res"] and region["type"].lower() != "figure":
|
||||
continue
|
||||
img_idx = region["img_idx"]
|
||||
|
||||
if region["type"].lower() == "figure":
|
||||
img_file_name = "{}_{}.jpg".format(region["bbox"], img_idx)
|
||||
markdown_string.append(
|
||||
f"""<div align="center">\n\t<img src="{img_name+"/"+img_file_name}">\n</div>"""
|
||||
)
|
||||
elif region["type"].lower() == "title":
|
||||
markdown_string.append(
|
||||
f"""# {region['res'][0]['text']}"""
|
||||
+ "".join(
|
||||
[" " + one_region["text"] for one_region in region["res"][1:]]
|
||||
)
|
||||
)
|
||||
elif region["type"].lower() == "table":
|
||||
markdown_string.append(region["res"]["html"])
|
||||
elif region["type"].lower() == "header" or region["type"].lower() == "footer":
|
||||
pass
|
||||
elif region["type"].lower() == "equation" and "latex" in region["res"]:
|
||||
markdown_string.append(f"""$${region["res"]["latex"]}$$""")
|
||||
elif region["type"].lower() == "text":
|
||||
merge_func = check_merge_method(region)
|
||||
# logger.warning(f"use merge method:{merge_func.__name__}")
|
||||
markdown_string.append(replace_special_char(merge_func(region)))
|
||||
else:
|
||||
string = ""
|
||||
for line in region["res"]:
|
||||
string += line["text"] + " "
|
||||
markdown_string.append(string)
|
||||
|
||||
md_path = os.path.join(save_folder, "{}_ocr.md".format(img_name))
|
||||
markdown_string = "\n\n".join(markdown_string)
|
||||
markdown_string = re.sub(r"\n{3,}", "\n\n", markdown_string)
|
||||
with open(md_path, "w", encoding="utf-8") as f:
|
||||
f.write(markdown_string)
|
||||
logger.info("markdown save to {}".format(md_path))
|
||||
4
ppstructure/recovery/requirements.txt
Normal file
4
ppstructure/recovery/requirements.txt
Normal file
@@ -0,0 +1,4 @@
|
||||
python-docx
|
||||
beautifulsoup4
|
||||
fonttools>=4.43.0
|
||||
fire>=0.3.0
|
||||
325
ppstructure/recovery/table_process.py
Normal file
325
ppstructure/recovery/table_process.py
Normal file
@@ -0,0 +1,325 @@
|
||||
# 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.
|
||||
"""
|
||||
This code is refer from: https://github.com/weizwx/html2docx/blob/master/htmldocx/h2d.py
|
||||
"""
|
||||
|
||||
import re
|
||||
import docx
|
||||
from docx import Document
|
||||
from bs4 import BeautifulSoup
|
||||
from html.parser import HTMLParser
|
||||
|
||||
|
||||
def get_table_rows(table_soup):
|
||||
table_row_selectors = [
|
||||
"table > tr",
|
||||
"table > thead > tr",
|
||||
"table > tbody > tr",
|
||||
"table > tfoot > tr",
|
||||
]
|
||||
# If there's a header, body, footer or direct child tr tags, add row dimensions from there
|
||||
return table_soup.select(", ".join(table_row_selectors), recursive=False)
|
||||
|
||||
|
||||
def get_table_columns(row):
|
||||
# Get all columns for the specified row tag.
|
||||
return row.find_all(["th", "td"], recursive=False) if row else []
|
||||
|
||||
|
||||
def get_table_dimensions(table_soup):
|
||||
# Get rows for the table
|
||||
rows = get_table_rows(table_soup)
|
||||
# Table is either empty or has non-direct children between table and tr tags
|
||||
# Thus the row dimensions and column dimensions are assumed to be 0
|
||||
|
||||
cols = get_table_columns(rows[0]) if rows else []
|
||||
# Add colspan calculation column number
|
||||
col_count = 0
|
||||
for col in cols:
|
||||
colspan = col.attrs.get("colspan", 1)
|
||||
col_count += int(colspan)
|
||||
|
||||
return rows, col_count
|
||||
|
||||
|
||||
def get_cell_html(soup):
|
||||
# Returns string of td element with opening and closing <td> tags removed
|
||||
# Cannot use find_all as it only finds element tags and does not find text which
|
||||
# is not inside an element
|
||||
return " ".join([str(i) for i in soup.contents])
|
||||
|
||||
|
||||
def delete_paragraph(paragraph):
|
||||
# https://github.com/python-openxml/python-docx/issues/33#issuecomment-77661907
|
||||
p = paragraph._element
|
||||
p.getparent().remove(p)
|
||||
p._p = p._element = None
|
||||
|
||||
|
||||
def remove_whitespace(string, leading=False, trailing=False):
|
||||
"""Remove white space from a string.
|
||||
Args:
|
||||
string(str): The string to remove white space from.
|
||||
leading(bool, optional): Remove leading new lines when True.
|
||||
trailing(bool, optional): Remove trailing new lines when False.
|
||||
Returns:
|
||||
str: The input string with new line characters removed and white space squashed.
|
||||
Examples:
|
||||
Single or multiple new line characters are replaced with space.
|
||||
>>> remove_whitespace("abc\\ndef")
|
||||
'abc def'
|
||||
>>> remove_whitespace("abc\\n\\n\\ndef")
|
||||
'abc def'
|
||||
New line characters surrounded by white space are replaced with a single space.
|
||||
>>> remove_whitespace("abc \\n \\n \\n def")
|
||||
'abc def'
|
||||
>>> remove_whitespace("abc \\n \\n \\n def")
|
||||
'abc def'
|
||||
Leading and trailing new lines are replaced with a single space.
|
||||
>>> remove_whitespace("\\nabc")
|
||||
' abc'
|
||||
>>> remove_whitespace(" \\n abc")
|
||||
' abc'
|
||||
>>> remove_whitespace("abc\\n")
|
||||
'abc '
|
||||
>>> remove_whitespace("abc \\n ")
|
||||
'abc '
|
||||
Use ``leading=True`` to remove leading new line characters, including any surrounding
|
||||
white space:
|
||||
>>> remove_whitespace("\\nabc", leading=True)
|
||||
'abc'
|
||||
>>> remove_whitespace(" \\n abc", leading=True)
|
||||
'abc'
|
||||
Use ``trailing=True`` to remove trailing new line characters, including any surrounding
|
||||
white space:
|
||||
>>> remove_whitespace("abc \\n ", trailing=True)
|
||||
'abc'
|
||||
"""
|
||||
# Remove any leading new line characters along with any surrounding white space
|
||||
if leading:
|
||||
string = re.sub(r"^\s*\n+\s*", "", string)
|
||||
|
||||
# Remove any trailing new line characters along with any surrounding white space
|
||||
if trailing:
|
||||
string = re.sub(r"\s*\n+\s*$", "", string)
|
||||
|
||||
# Replace new line characters and absorb any surrounding space.
|
||||
string = re.sub(r"\s*\n\s*", " ", string)
|
||||
# TODO need some way to get rid of extra spaces in e.g. text <span> </span> text
|
||||
return re.sub(r"\s+", " ", string)
|
||||
|
||||
|
||||
font_styles = {
|
||||
"b": "bold",
|
||||
"strong": "bold",
|
||||
"em": "italic",
|
||||
"i": "italic",
|
||||
"u": "underline",
|
||||
"s": "strike",
|
||||
"sup": "superscript",
|
||||
"sub": "subscript",
|
||||
"th": "bold",
|
||||
}
|
||||
|
||||
font_names = {
|
||||
"code": "Courier",
|
||||
"pre": "Courier",
|
||||
}
|
||||
|
||||
|
||||
class HtmlToDocx(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.options = {
|
||||
"fix-html": True,
|
||||
"images": True,
|
||||
"tables": True,
|
||||
"styles": True,
|
||||
}
|
||||
self.table_row_selectors = [
|
||||
"table > tr",
|
||||
"table > thead > tr",
|
||||
"table > tbody > tr",
|
||||
"table > tfoot > tr",
|
||||
]
|
||||
self.table_style = None
|
||||
self.paragraph_style = None
|
||||
|
||||
def set_initial_attrs(self, document=None):
|
||||
self.tags = {
|
||||
"span": [],
|
||||
"list": [],
|
||||
}
|
||||
if document:
|
||||
self.doc = document
|
||||
else:
|
||||
self.doc = Document()
|
||||
self.bs = self.options["fix-html"] # whether or not to clean with BeautifulSoup
|
||||
self.document = self.doc
|
||||
self.include_tables = True # TODO add this option back in?
|
||||
self.include_images = self.options["images"]
|
||||
self.include_styles = self.options["styles"]
|
||||
self.paragraph = None
|
||||
self.skip = False
|
||||
self.skip_tag = None
|
||||
self.instances_to_skip = 0
|
||||
|
||||
def copy_settings_from(self, other):
|
||||
"""Copy settings from another instance of HtmlToDocx"""
|
||||
self.table_style = other.table_style
|
||||
self.paragraph_style = other.paragraph_style
|
||||
|
||||
def ignore_nested_tables(self, tables_soup):
|
||||
"""
|
||||
Returns array containing only the highest level tables
|
||||
Operates on the assumption that bs4 returns child elements immediately after
|
||||
the parent element in `find_all`. If this changes in the future, this method will need to be updated
|
||||
:return:
|
||||
"""
|
||||
new_tables = []
|
||||
nest = 0
|
||||
for table in tables_soup:
|
||||
if nest:
|
||||
nest -= 1
|
||||
continue
|
||||
new_tables.append(table)
|
||||
nest = len(table.find_all("table"))
|
||||
return new_tables
|
||||
|
||||
def get_tables(self):
|
||||
if not hasattr(self, "soup"):
|
||||
self.include_tables = False
|
||||
return
|
||||
# find other way to do it, or require this dependency?
|
||||
self.tables = self.ignore_nested_tables(self.soup.find_all("table"))
|
||||
self.table_no = 0
|
||||
|
||||
def run_process(self, html):
|
||||
if self.bs and BeautifulSoup:
|
||||
self.soup = BeautifulSoup(html, "html.parser")
|
||||
html = str(self.soup)
|
||||
if self.include_tables:
|
||||
self.get_tables()
|
||||
self.feed(html)
|
||||
|
||||
def add_html_to_cell(self, html, cell):
|
||||
if not isinstance(cell, docx.table._Cell):
|
||||
raise ValueError("Second argument needs to be a %s" % docx.table._Cell)
|
||||
unwanted_paragraph = cell.paragraphs[0]
|
||||
if unwanted_paragraph.text == "":
|
||||
delete_paragraph(unwanted_paragraph)
|
||||
self.set_initial_attrs(cell)
|
||||
self.run_process(html)
|
||||
# cells must end with a paragraph or will get message about corrupt file
|
||||
# https://stackoverflow.com/a/29287121
|
||||
if not self.doc.paragraphs:
|
||||
self.doc.add_paragraph("")
|
||||
|
||||
def apply_paragraph_style(self, style=None):
|
||||
try:
|
||||
if style:
|
||||
self.paragraph.style = style
|
||||
elif self.paragraph_style:
|
||||
self.paragraph.style = self.paragraph_style
|
||||
except KeyError as e:
|
||||
raise ValueError(f"Unable to apply style {self.paragraph_style}.") from e
|
||||
|
||||
def handle_table(self, html, doc):
|
||||
"""
|
||||
To handle nested tables, we will parse tables manually as follows:
|
||||
Get table soup
|
||||
Create docx table
|
||||
Iterate over soup and fill docx table with new instances of this parser
|
||||
Tell HTMLParser to ignore any tags until the corresponding closing table tag
|
||||
"""
|
||||
table_soup = BeautifulSoup(html, "html.parser")
|
||||
rows, cols_len = get_table_dimensions(table_soup)
|
||||
table = doc.add_table(len(rows), cols_len)
|
||||
table.style = doc.styles["Table Grid"]
|
||||
|
||||
num_rows = len(table.rows)
|
||||
num_cols = len(table.columns)
|
||||
|
||||
cell_row = 0
|
||||
for index, row in enumerate(rows):
|
||||
cols = get_table_columns(row)
|
||||
cell_col = 0
|
||||
for col in cols:
|
||||
colspan = int(col.attrs.get("colspan", 1))
|
||||
rowspan = int(col.attrs.get("rowspan", 1))
|
||||
|
||||
cell_html = get_cell_html(col)
|
||||
if col.name == "th":
|
||||
cell_html = "<b>%s</b>" % cell_html
|
||||
|
||||
if cell_row >= num_rows or cell_col >= num_cols:
|
||||
continue
|
||||
|
||||
docx_cell = table.cell(cell_row, cell_col)
|
||||
|
||||
while docx_cell.text != "": # Skip the merged cell
|
||||
cell_col += 1
|
||||
docx_cell = table.cell(cell_row, cell_col)
|
||||
|
||||
cell_to_merge = table.cell(
|
||||
cell_row + rowspan - 1, cell_col + colspan - 1
|
||||
)
|
||||
if docx_cell != cell_to_merge:
|
||||
docx_cell.merge(cell_to_merge)
|
||||
|
||||
child_parser = HtmlToDocx()
|
||||
child_parser.copy_settings_from(self)
|
||||
child_parser.add_html_to_cell(cell_html or " ", docx_cell)
|
||||
|
||||
cell_col += colspan
|
||||
cell_row += 1
|
||||
|
||||
def handle_data(self, data):
|
||||
if self.skip:
|
||||
return
|
||||
|
||||
# Only remove white space if we're not in a pre block.
|
||||
if "pre" not in self.tags:
|
||||
# remove leading and trailing whitespace in all instances
|
||||
data = remove_whitespace(data, True, True)
|
||||
|
||||
if not self.paragraph:
|
||||
self.paragraph = self.doc.add_paragraph()
|
||||
self.apply_paragraph_style()
|
||||
|
||||
# There can only be one nested link in a valid html document
|
||||
# You cannot have interactive content in an A tag, this includes links
|
||||
# https://html.spec.whatwg.org/#interactive-content
|
||||
link = self.tags.get("a")
|
||||
if link:
|
||||
self.handle_link(link["href"], data)
|
||||
else:
|
||||
# If there's a link, dont put the data directly in the run
|
||||
self.run = self.paragraph.add_run(data)
|
||||
spans = self.tags["span"]
|
||||
for span in spans:
|
||||
if "style" in span:
|
||||
style = self.parse_dict_string(span["style"])
|
||||
self.add_styles_to_run(style)
|
||||
|
||||
# add font style and name
|
||||
for tag in self.tags:
|
||||
if tag in font_styles:
|
||||
font_style = font_styles[tag]
|
||||
setattr(self.run.font, font_style, True)
|
||||
|
||||
if tag in font_names:
|
||||
font_name = font_names[tag]
|
||||
self.run.font.name = font_name
|
||||
159
ppstructure/table/README.md
Normal file
159
ppstructure/table/README.md
Normal file
@@ -0,0 +1,159 @@
|
||||
English | [简体中文](README_ch.md)
|
||||
|
||||
# Table Recognition
|
||||
|
||||
- [1. pipeline](#1-pipeline)
|
||||
- [2. Performance](#2-performance)
|
||||
- [3. Result](#3-result)
|
||||
- [4. How to use](#4-how-to-use)
|
||||
- [4.1 Quick start](#41-quick-start)
|
||||
- [4.2 Training, Evaluation and Inference](#42-training-evaluation-and-inference)
|
||||
- [4.3 Calculate TEDS](#43-calculate-teds)
|
||||
- [5. Reference](#5-reference)
|
||||
|
||||
|
||||
## 1. pipeline
|
||||
The table recognition mainly contains three models
|
||||
1. Single line text detection-DB
|
||||
2. Single line text recognition-CRNN
|
||||
3. Table structure and cell coordinate prediction-SLANet
|
||||
|
||||
The table recognition flow chart is as follows
|
||||
|
||||
|
||||
|
||||
1. The coordinates of single-line text is detected by DB model, and then sends it to the recognition model to get the recognition result.
|
||||
2. The table structure and cell coordinates is predicted by SLANet model.
|
||||
3. The recognition result of the cell is combined by the coordinates, recognition result of the single line and the coordinates of the cell.
|
||||
4. The cell recognition result and the table structure together construct the html string of the table.
|
||||
|
||||
## 2. Performance
|
||||
We evaluated the algorithm on the PubTabNet<sup>[1]</sup> eval dataset, and the performance is as follows:
|
||||
|
||||
|Method|Acc|[TEDS(Tree-Edit-Distance-based Similarity)](https://github.com/ibm-aur-nlp/PubTabNet/tree/master/src)|Speed|
|
||||
| --- | --- | --- | ---|
|
||||
| EDD<sup>[2]</sup> |x| 88.30% |x|
|
||||
| TableRec-RARE(ours) | 71.73%| 93.88% |779ms|
|
||||
| SLANet(ours) | 76.31%| 95.89%|766ms|
|
||||
|
||||
The performance indicators are explained as follows:
|
||||
- Acc: The accuracy of the table structure in each image, a wrong token is considered an error.
|
||||
- TEDS: The accuracy of the model's restoration of table information. This indicator evaluates not only the table structure, but also the text content in the table.
|
||||
- Speed: The inference speed of a single image when the model runs on the CPU machine and MKL is enabled.
|
||||
|
||||
## 3. Result
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## 4. How to use
|
||||
|
||||
### 4.1 Quick start
|
||||
|
||||
PP-Structure currently provides table recognition models in both Chinese and English. For the model link, see [models_list](../docs/models_list.md). The whl package is also provided for quick use, see [quickstart](../docs/quickstart_en.md) for details.
|
||||
|
||||
The following takes the Chinese table recognition model as an example to introduce how to recognize a table.
|
||||
|
||||
Use the following commands to quickly complete the identification of a table.
|
||||
|
||||
```python
|
||||
cd PaddleOCR/ppstructure
|
||||
|
||||
# download model
|
||||
mkdir inference && cd inference
|
||||
# Download the PP-OCRv3 text detection model and unzip it
|
||||
wget https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/PP-OCRv3_mobile_det_infer.tar && tar xf PP-OCRv3_mobile_det_infer.tar
|
||||
# Download the PP-OCRv3 text recognition model and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar && tar xf ch_PP-OCRv3_rec_infer.tar
|
||||
# Download the PP-StructureV2 form recognition model and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/ch_ppstructure_mobile_v2.0_SLANet_infer.tar && tar xf ch_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
cd ..
|
||||
# run
|
||||
python3 table/predict_table.py \
|
||||
--det_model_dir=inference/PP-OCRv3_mobile_det_infer \
|
||||
--rec_model_dir=inference/ch_PP-OCRv3_rec_infer \
|
||||
--table_model_dir=inference/ch_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--rec_char_dict_path=../ppocr/utils/ppocr_keys_v1.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict_ch.txt \
|
||||
--image_dir=docs/table/table.jpg \
|
||||
--output=../output/table
|
||||
|
||||
```
|
||||
|
||||
After the operation is completed, the excel table of each image will be saved to the directory specified by the output field, and an html file will be produced in the directory to visually view the cell coordinates and the recognized table.
|
||||
|
||||
**NOTE**
|
||||
1. If you want to use the English table recognition model, you need to download the English text detection and recognition model and the English table recognition model in [models_list](../docs/models_list_en.md), and replace `table_structure_dict_ch.txt` with `table_structure_dict.txt`.
|
||||
2. To use the TableRec-RARE model, you need to replace `table_structure_dict_ch.txt` with `table_structure_dict.txt`, and add parameter `--merge_no_span_structure=False`
|
||||
|
||||
### 4.2 Training, Evaluation and Inference
|
||||
|
||||
The training, evaluation and inference process of the text detection model can be referred to [detection](../../doc/doc_en/detection_en.md)
|
||||
|
||||
The training, evaluation and inference process of the text recognition model can be referred to [recognition](../../doc/doc_en/recognition_en.md)
|
||||
|
||||
The training, evaluation and inference process of the table recognition model can be referred to [table_recognition](../../doc/doc_en/table_recognition_en.md)
|
||||
|
||||
### 4.3 Calculate TEDS
|
||||
|
||||
The table uses [TEDS(Tree-Edit-Distance-based Similarity)](https://github.com/ibm-aur-nlp/PubTabNet/tree/master/src) as the evaluation metric of the model. Before the model evaluation, the three models in the pipeline need to be exported as inference models (we have provided them), and the gt for evaluation needs to be prepared. Examples of gt are as follows:
|
||||
```txt
|
||||
PMC5755158_010_01.png <html><body><table><thead><tr><td></td><td><b>Weaning</b></td><td><b>Week 15</b></td><td><b>Off-test</b></td></tr></thead><tbody><tr><td>Weaning</td><td>–</td><td>–</td><td>–</td></tr><tr><td>Week 15</td><td>–</td><td>0.17 ± 0.08</td><td>0.16 ± 0.03</td></tr><tr><td>Off-test</td><td>–</td><td>0.80 ± 0.24</td><td>0.19 ± 0.09</td></tr></tbody></table></body></html>
|
||||
```
|
||||
Each line in gt consists of the file name and the html string of the table. The file name and the html string of the table are separated by `\t`.
|
||||
|
||||
You can also use the following command to generate an evaluation gt file from the annotation file:
|
||||
```python
|
||||
python3 ppstructure/table/convert_label2html.py --ori_gt_path /path/to/your_label_file --save_path /path/to/save_file
|
||||
```
|
||||
|
||||
Use the following command to evaluate. After the evaluation is completed, the teds indicator will be output.
|
||||
```python
|
||||
python3 table/eval_table.py \
|
||||
--det_model_dir=path/to/det_model_dir \
|
||||
--rec_model_dir=path/to/rec_model_dir \
|
||||
--table_model_dir=path/to/table_model_dir \
|
||||
--image_dir=docs/table/table.jpg \
|
||||
--rec_char_dict_path=../ppocr/utils/dict/table_dict.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--det_limit_side_len=736 \
|
||||
--det_limit_type=min \
|
||||
--gt_path=path/to/gt.txt
|
||||
```
|
||||
|
||||
Evaluate on the PubLatNet dataset using the English model
|
||||
|
||||
```bash
|
||||
cd PaddleOCR/ppstructure
|
||||
# Download the model
|
||||
mkdir inference && cd inference
|
||||
# Download the text detection model trained on the PubTabNet dataset and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar && tar xf en_ppocr_mobile_v2.0_table_det_infer.tar
|
||||
# Download the text recognition model trained on the PubTabNet dataset and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar && tar xf en_ppocr_mobile_v2.0_table_rec_infer.tar
|
||||
# Download the table recognition model trained on the PubTabNet dataset and unzip it
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar && tar xf en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
cd ..
|
||||
|
||||
python3 table/eval_table.py \
|
||||
--det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer \
|
||||
--rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer \
|
||||
--table_model_dir=inference/en_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--image_dir=train_data/table/pubtabnet/val/ \
|
||||
--rec_char_dict_path=../ppocr/utils/dict/table_dict.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--det_limit_side_len=736 \
|
||||
--det_limit_type=min \
|
||||
--rec_image_shape=3,32,320 \
|
||||
--gt_path=path/to/gt.txt
|
||||
```
|
||||
|
||||
output is
|
||||
```bash
|
||||
teds: 95.89
|
||||
```
|
||||
|
||||
## 5. Reference
|
||||
1. https://github.com/ibm-aur-nlp/PubTabNet
|
||||
2. https://arxiv.org/pdf/1911.10683
|
||||
163
ppstructure/table/README_ch.md
Normal file
163
ppstructure/table/README_ch.md
Normal file
@@ -0,0 +1,163 @@
|
||||
[English](README.md) | 简体中文
|
||||
|
||||
# 表格识别
|
||||
|
||||
- [1. 表格识别 pipeline](#1-表格识别-pipeline)
|
||||
- [2. 性能](#2-性能)
|
||||
- [3. 效果演示](#3-效果演示)
|
||||
- [4. 使用](#4-使用)
|
||||
- [4.1 快速开始](#41-快速开始)
|
||||
- [4.2 模型训练、评估与推理](#42-模型训练评估与推理)
|
||||
- [4.3 计算TEDS](#43-计算teds)
|
||||
- [5. Reference](#5-reference)
|
||||
|
||||
|
||||
## 1. 表格识别 pipeline
|
||||
|
||||
表格识别主要包含三个模型
|
||||
1. 单行文本检测-DB
|
||||
2. 单行文本识别-CRNN
|
||||
3. 表格结构和cell坐标预测-SLANet
|
||||
|
||||
具体流程图如下
|
||||
|
||||
|
||||
|
||||
流程说明:
|
||||
|
||||
1. 图片由单行文字检测模型检测到单行文字的坐标，然后送入识别模型拿到识别结果。
|
||||
2. 图片由SLANet模型拿到表格的结构信息和单元格的坐标信息。
|
||||
3. 由单行文字的坐标、识别结果和单元格的坐标一起组合出单元格的识别结果。
|
||||
4. 单元格的识别结果和表格结构一起构造表格的html字符串。
|
||||
|
||||
|
||||
## 2. 性能
|
||||
|
||||
我们在 PubTabNet<sup>[1]</sup> 评估数据集上对算法进行了评估，性能如下
|
||||
|
||||
|
||||
|算法|Acc|[TEDS(Tree-Edit-Distance-based Similarity)](https://github.com/ibm-aur-nlp/PubTabNet/tree/master/src)|Speed|
|
||||
| --- | --- | --- | ---|
|
||||
| EDD<sup>[2]</sup> |x| 88.30% |x|
|
||||
| TableRec-RARE(ours) | 71.73%| 93.88% |779ms|
|
||||
| SLANet(ours) |76.31%| 95.89%|766ms|
|
||||
|
||||
性能指标解释如下：
|
||||
- Acc: 模型对每张图像里表格结构的识别准确率，错一个token就算错误。
|
||||
- TEDS: 模型对表格信息还原的准确度，此指标评价内容不仅包含表格结构，还包含表格内的文字内容。
|
||||
- Speed: 模型在CPU机器上，开启MKL的情况下，单张图片的推理速度。
|
||||
|
||||
## 3. 效果演示
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
## 4. 使用
|
||||
|
||||
### 4.1 快速开始
|
||||
|
||||
PP-Structure目前提供了中英文两种语言的表格识别模型，模型链接见 [models_list](../docs/models_list.md)。也提供了whl包的形式方便快速使用，详见 [quickstart](../docs/quickstart.md)。
|
||||
|
||||
下面以中文表格识别模型为例，介绍如何识别一张表格。
|
||||
|
||||
使用如下命令即可快速完成一张表格的识别。
|
||||
```python
|
||||
cd PaddleOCR/ppstructure
|
||||
|
||||
# 下载模型
|
||||
mkdir inference && cd inference
|
||||
# 下载PP-OCRv3文本检测模型并解压
|
||||
wget https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/PP-OCRv3_mobile_det_infer.tar && tar xf PP-OCRv3_mobile_det_infer.tar
|
||||
# 下载PP-OCRv3文本识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/PP-OCRv3/chinese/ch_PP-OCRv3_rec_infer.tar && tar xf ch_PP-OCRv3_rec_infer.tar
|
||||
# 下载PP-StructureV2中文表格识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/ch_ppstructure_mobile_v2.0_SLANet_infer.tar && tar xf ch_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
cd ..
|
||||
# 执行表格识别
|
||||
python table/predict_table.py \
|
||||
--det_model_dir=inference/PP-OCRv3_mobile_det_infer \
|
||||
--rec_model_dir=inference/ch_PP-OCRv3_rec_infer \
|
||||
--table_model_dir=inference/ch_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--rec_char_dict_path=../ppocr/utils/ppocr_keys_v1.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict_ch.txt \
|
||||
--image_dir=docs/table/table.jpg \
|
||||
--output=../output/table
|
||||
```
|
||||
运行完成后，每张图片的excel表格会保存到output字段指定的目录下，同时在该目录下回生产一个html文件，用于可视化查看单元格坐标和识别的表格。
|
||||
|
||||
**NOTE**
|
||||
1. 如果想使用英文模型，需要在 [models_list](../docs/models_list.md) 中下载英文文字检测识别模型和英文表格识别模型，同时替换`table_structure_dict_ch.txt`为`table_structure_dict.txt`即可。
|
||||
2. 如需使用TableRec-RARE模型，需要替换`table_structure_dict_ch.txt`为`table_structure_dict.txt`，同时参数`--merge_no_span_structure=False`
|
||||
|
||||
### 4.2 模型训练、评估与推理
|
||||
|
||||
文本检测模型的训练、评估和推理流程可参考 [detection](../../doc/doc_ch/detection.md)
|
||||
|
||||
文本识别模型的训练、评估和推理流程可参考 [recognition](../../doc/doc_ch/recognition.md)
|
||||
|
||||
表格识别模型的训练、评估和推理流程可参考 [table_recognition](../../doc/doc_ch/table_recognition.md)
|
||||
|
||||
### 4.3 计算TEDS
|
||||
|
||||
表格使用 [TEDS(Tree-Edit-Distance-based Similarity)](https://github.com/ibm-aur-nlp/PubTabNet/tree/master/src) 作为模型的评估指标。在进行模型评估之前，需要将pipeline中的三个模型分别导出为inference模型(我们已经提供好)，还需要准备评估的gt， gt示例如下:
|
||||
```txt
|
||||
PMC5755158_010_01.png <html><body><table><thead><tr><td></td><td><b>Weaning</b></td><td><b>Week 15</b></td><td><b>Off-test</b></td></tr></thead><tbody><tr><td>Weaning</td><td>–</td><td>–</td><td>–</td></tr><tr><td>Week 15</td><td>–</td><td>0.17 ± 0.08</td><td>0.16 ± 0.03</td></tr><tr><td>Off-test</td><td>–</td><td>0.80 ± 0.24</td><td>0.19 ± 0.09</td></tr></tbody></table></body></html>
|
||||
```
|
||||
gt每一行都由文件名和表格的html字符串组成，文件名和表格的html字符串之间使用`\t`分隔。
|
||||
|
||||
也可使用如下命令，由标注文件生成评估的gt文件：
|
||||
```python
|
||||
python3 ppstructure/table/convert_label2html.py --ori_gt_path /path/to/your_label_file --save_path /path/to/save_file
|
||||
```
|
||||
|
||||
准备完成后使用如下命令进行评估，评估完成后会输出teds指标。
|
||||
```python
|
||||
cd PaddleOCR/ppstructure
|
||||
python3 table/eval_table.py \
|
||||
--det_model_dir=path/to/det_model_dir \
|
||||
--rec_model_dir=path/to/rec_model_dir \
|
||||
--table_model_dir=path/to/table_model_dir \
|
||||
--image_dir=docs/table/table.jpg \
|
||||
--rec_char_dict_path=../ppocr/utils/dict/table_dict.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--det_limit_side_len=736 \
|
||||
--det_limit_type=min \
|
||||
--gt_path=path/to/gt.txt
|
||||
```
|
||||
|
||||
如使用英文表格识别模型在PubLatNet数据集上进行评估
|
||||
|
||||
```bash
|
||||
cd PaddleOCR/ppstructure
|
||||
# 下载模型
|
||||
mkdir inference && cd inference
|
||||
# 下载基于PubTabNet数据集训练的文本检测模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_det_infer.tar && tar xf en_ppocr_mobile_v2.0_table_det_infer.tar
|
||||
# 下载基于PubTabNet数据集训练的文本识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/table/en_ppocr_mobile_v2.0_table_rec_infer.tar && tar xf en_ppocr_mobile_v2.0_table_rec_infer.tar
|
||||
# 下载基于PubTabNet数据集训练的表格识别模型并解压
|
||||
wget https://paddleocr.bj.bcebos.com/ppstructure/models/slanet/paddle3.0b2/en_ppstructure_mobile_v2.0_SLANet_infer.tar && tar xf en_ppstructure_mobile_v2.0_SLANet_infer.tar
|
||||
cd ..
|
||||
|
||||
python3 table/eval_table.py \
|
||||
--det_model_dir=inference/en_ppocr_mobile_v2.0_table_det_infer \
|
||||
--rec_model_dir=inference/en_ppocr_mobile_v2.0_table_rec_infer \
|
||||
--table_model_dir=inference/en_ppstructure_mobile_v2.0_SLANet_infer \
|
||||
--image_dir=train_data/table/pubtabnet/val/ \
|
||||
--rec_char_dict_path=../ppocr/utils/dict/table_dict.txt \
|
||||
--table_char_dict_path=../ppocr/utils/dict/table_structure_dict.txt \
|
||||
--det_limit_side_len=736 \
|
||||
--det_limit_type=min \
|
||||
--rec_image_shape=3,32,320 \
|
||||
--gt_path=path/to/gt.txt
|
||||
```
|
||||
|
||||
将会输出
|
||||
```bash
|
||||
teds: 95.89
|
||||
```
|
||||
|
||||
## 5. Reference
|
||||
1. https://github.com/ibm-aur-nlp/PubTabNet
|
||||
2. https://arxiv.org/pdf/1911.10683
|
||||
13
ppstructure/table/__init__.py
Normal file
13
ppstructure/table/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
102
ppstructure/table/convert_label2html.py
Normal file
102
ppstructure/table/convert_label2html.py
Normal file
@@ -0,0 +1,102 @@
|
||||
# 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.
|
||||
"""
|
||||
convert table label to html
|
||||
"""
|
||||
|
||||
import json
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def save_pred_txt(key, val, tmp_file_path):
|
||||
with open(tmp_file_path, "a+", encoding="utf-8") as f:
|
||||
f.write("{}\t{}\n".format(key, val))
|
||||
|
||||
|
||||
def skip_char(text, sp_char_list):
|
||||
"""
|
||||
skip empty cell
|
||||
@param text: text in cell
|
||||
@param sp_char_list: style char and special code
|
||||
@return:
|
||||
"""
|
||||
for sp_char in sp_char_list:
|
||||
text = text.replace(sp_char, "")
|
||||
return text
|
||||
|
||||
|
||||
def gen_html(img):
|
||||
"""
|
||||
Formats HTML code from tokenized annotation of img
|
||||
"""
|
||||
html_code = img["html"]["structure"]["tokens"].copy()
|
||||
to_insert = [i for i, tag in enumerate(html_code) if tag in ("<td>", ">")]
|
||||
for i, cell in zip(to_insert[::-1], img["html"]["cells"][::-1]):
|
||||
if cell["tokens"]:
|
||||
text = "".join(cell["tokens"])
|
||||
# skip empty text
|
||||
sp_char_list = ["<b>", "</b>", "\u2028", " ", "<i>", "</i>"]
|
||||
text_remove_style = skip_char(text, sp_char_list)
|
||||
if len(text_remove_style) == 0:
|
||||
continue
|
||||
html_code.insert(i + 1, text)
|
||||
html_code = "".join(html_code)
|
||||
html_code = "<html><body><table>{}</table></body></html>".format(html_code)
|
||||
return html_code
|
||||
|
||||
|
||||
def load_gt_data(gt_path):
|
||||
"""
|
||||
load gt
|
||||
@param gt_path:
|
||||
@return:
|
||||
"""
|
||||
data_list = {}
|
||||
with open(gt_path, "rb") as f:
|
||||
lines = f.readlines()
|
||||
for line in tqdm(lines):
|
||||
data_line = line.decode("utf-8").strip("\n")
|
||||
info = json.loads(data_line)
|
||||
data_list[info["filename"]] = info
|
||||
return data_list
|
||||
|
||||
|
||||
def convert(origin_gt_path, save_path):
|
||||
"""
|
||||
gen html from label file
|
||||
@param origin_gt_path:
|
||||
@param save_path:
|
||||
@return:
|
||||
"""
|
||||
data_dict = load_gt_data(origin_gt_path)
|
||||
for img_name, gt in tqdm(data_dict.items()):
|
||||
html = gen_html(gt)
|
||||
save_pred_txt(img_name, html, save_path)
|
||||
print("convert finish")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="args for paddleserving")
|
||||
parser.add_argument("--ori_gt_path", type=str, required=True, help="label gt path")
|
||||
parser.add_argument(
|
||||
"--save_path", type=str, required=True, help="path to save file"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
convert(args.ori_gt_path, args.save_path)
|
||||
107
ppstructure/table/eval_table.py
Executable file
107
ppstructure/table/eval_table.py
Executable file
@@ -0,0 +1,107 @@
|
||||
# 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.insert(0, os.path.abspath(os.path.join(__dir__, "../..")))
|
||||
|
||||
import cv2
|
||||
import pickle
|
||||
import paddle
|
||||
from tqdm import tqdm
|
||||
from ppstructure.table.table_metric import TEDS
|
||||
from ppstructure.table.predict_table import TableSystem
|
||||
from ppstructure.utility import init_args
|
||||
from ppocr.utils.logging import get_logger
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = init_args()
|
||||
parser.add_argument("--gt_path", type=str)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_txt(txt_path):
|
||||
pred_html_dict = {}
|
||||
if not os.path.exists(txt_path):
|
||||
return pred_html_dict
|
||||
with open(txt_path, encoding="utf-8") as f:
|
||||
lines = f.readlines()
|
||||
for line in lines:
|
||||
line = line.strip().split("\t")
|
||||
img_name, pred_html = line
|
||||
pred_html_dict[img_name] = pred_html
|
||||
return pred_html_dict
|
||||
|
||||
|
||||
def load_result(path):
|
||||
data = {}
|
||||
if os.path.exists(path):
|
||||
data = pickle.load(open(path, "rb"))
|
||||
return data
|
||||
|
||||
|
||||
def save_result(path, data):
|
||||
old_data = load_result(path)
|
||||
old_data.update(data)
|
||||
with open(path, "wb") as f:
|
||||
pickle.dump(old_data, f)
|
||||
|
||||
|
||||
def main(gt_path, img_root, args):
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
# init TableSystem
|
||||
text_sys = TableSystem(args)
|
||||
# load gt and preds html result
|
||||
gt_html_dict = load_txt(gt_path)
|
||||
|
||||
ocr_result = load_result(os.path.join(args.output, "ocr.pickle"))
|
||||
structure_result = load_result(os.path.join(args.output, "structure.pickle"))
|
||||
|
||||
pred_htmls = []
|
||||
gt_htmls = []
|
||||
for img_name, gt_html in tqdm(gt_html_dict.items()):
|
||||
img = cv2.imread(os.path.join(img_root, img_name))
|
||||
# run ocr and save result
|
||||
if img_name not in ocr_result:
|
||||
dt_boxes, rec_res, _, _ = text_sys._ocr(img)
|
||||
ocr_result[img_name] = [dt_boxes, rec_res]
|
||||
save_result(os.path.join(args.output, "ocr.pickle"), ocr_result)
|
||||
# run structure and save result
|
||||
if img_name not in structure_result:
|
||||
structure_res, _ = text_sys._structure(img)
|
||||
structure_result[img_name] = structure_res
|
||||
save_result(os.path.join(args.output, "structure.pickle"), structure_result)
|
||||
dt_boxes, rec_res = ocr_result[img_name]
|
||||
structure_res = structure_result[img_name]
|
||||
# match ocr and structure
|
||||
pred_html = text_sys.match(structure_res, dt_boxes, rec_res)
|
||||
|
||||
pred_htmls.append(pred_html)
|
||||
gt_htmls.append(gt_html)
|
||||
|
||||
# compute teds
|
||||
teds = TEDS(n_jobs=16)
|
||||
scores = teds.batch_evaluate_html(gt_htmls, pred_htmls)
|
||||
logger.info("teds: {}".format(sum(scores) / len(scores)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args.gt_path, args.image_dir, args)
|
||||
206
ppstructure/table/matcher.py
Executable file
206
ppstructure/table/matcher.py
Executable file
@@ -0,0 +1,206 @@
|
||||
# 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 numpy as np
|
||||
from ppstructure.table.table_master_match import deal_eb_token, deal_bb
|
||||
import html
|
||||
|
||||
|
||||
def distance(box_1, box_2):
|
||||
x1, y1, x2, y2 = box_1
|
||||
x3, y3, x4, y4 = box_2
|
||||
dis = abs(x3 - x1) + abs(y3 - y1) + abs(x4 - x2) + abs(y4 - y2)
|
||||
dis_2 = abs(x3 - x1) + abs(y3 - y1)
|
||||
dis_3 = abs(x4 - x2) + abs(y4 - y2)
|
||||
return dis + min(dis_2, dis_3)
|
||||
|
||||
|
||||
def compute_iou(rec1, rec2):
|
||||
"""
|
||||
computing IoU
|
||||
:param rec1: (y0, x0, y1, x1), which reflects
|
||||
(top, left, bottom, right)
|
||||
:param rec2: (y0, x0, y1, x1)
|
||||
:return: scala value of IoU
|
||||
"""
|
||||
# computing area of each rectangles
|
||||
S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1])
|
||||
S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1])
|
||||
|
||||
# computing the sum_area
|
||||
sum_area = S_rec1 + S_rec2
|
||||
|
||||
# find the each edge of intersect rectangle
|
||||
left_line = max(rec1[1], rec2[1])
|
||||
right_line = min(rec1[3], rec2[3])
|
||||
top_line = max(rec1[0], rec2[0])
|
||||
bottom_line = min(rec1[2], rec2[2])
|
||||
|
||||
# judge if there is an intersect
|
||||
if left_line >= right_line or top_line >= bottom_line:
|
||||
return 0.0
|
||||
else:
|
||||
intersect = (right_line - left_line) * (bottom_line - top_line)
|
||||
return (intersect / (sum_area - intersect)) * 1.0
|
||||
|
||||
|
||||
class TableMatch:
|
||||
def __init__(self, filter_ocr_result=False, use_master=False):
|
||||
self.filter_ocr_result = filter_ocr_result
|
||||
self.use_master = use_master
|
||||
|
||||
def __call__(self, structure_res, dt_boxes, rec_res):
|
||||
pred_structures, pred_bboxes = structure_res
|
||||
if self.filter_ocr_result:
|
||||
dt_boxes, rec_res = self._filter_ocr_result(pred_bboxes, dt_boxes, rec_res)
|
||||
matched_index = self.match_result(dt_boxes, pred_bboxes)
|
||||
if self.use_master:
|
||||
pred_html, pred = self.get_pred_html_master(
|
||||
pred_structures, matched_index, rec_res
|
||||
)
|
||||
else:
|
||||
pred_html, pred = self.get_pred_html(
|
||||
pred_structures, matched_index, rec_res
|
||||
)
|
||||
return pred_html
|
||||
|
||||
def match_result(self, dt_boxes, pred_bboxes):
|
||||
matched = {}
|
||||
for i, gt_box in enumerate(dt_boxes):
|
||||
distances = []
|
||||
for j, pred_box in enumerate(pred_bboxes):
|
||||
if len(pred_box) == 8:
|
||||
pred_box = [
|
||||
np.min(pred_box[0::2]),
|
||||
np.min(pred_box[1::2]),
|
||||
np.max(pred_box[0::2]),
|
||||
np.max(pred_box[1::2]),
|
||||
]
|
||||
distances.append(
|
||||
(distance(gt_box, pred_box), 1.0 - compute_iou(gt_box, pred_box))
|
||||
) # compute iou and l1 distance
|
||||
sorted_distances = distances.copy()
|
||||
# select det box by iou and l1 distance
|
||||
sorted_distances = sorted(
|
||||
sorted_distances, key=lambda item: (item[1], item[0])
|
||||
)
|
||||
if distances.index(sorted_distances[0]) not in matched.keys():
|
||||
matched[distances.index(sorted_distances[0])] = [i]
|
||||
else:
|
||||
matched[distances.index(sorted_distances[0])].append(i)
|
||||
return matched
|
||||
|
||||
def get_pred_html(self, pred_structures, matched_index, ocr_contents):
|
||||
end_html = []
|
||||
td_index = 0
|
||||
for tag in pred_structures:
|
||||
if "</td>" in tag:
|
||||
if "<td></td>" == tag:
|
||||
end_html.extend("<td>")
|
||||
if td_index in matched_index.keys():
|
||||
b_with = False
|
||||
if (
|
||||
"<b>" in ocr_contents[matched_index[td_index][0]]
|
||||
and len(matched_index[td_index]) > 1
|
||||
):
|
||||
b_with = True
|
||||
end_html.extend("<b>")
|
||||
for i, td_index_index in enumerate(matched_index[td_index]):
|
||||
content = ocr_contents[td_index_index][0]
|
||||
if len(matched_index[td_index]) > 1:
|
||||
if len(content) == 0:
|
||||
continue
|
||||
if content[0] == " ":
|
||||
content = content[1:]
|
||||
if "<b>" in content:
|
||||
content = content[3:]
|
||||
if "</b>" in content:
|
||||
content = content[:-4]
|
||||
if len(content) == 0:
|
||||
continue
|
||||
if (
|
||||
i != len(matched_index[td_index]) - 1
|
||||
and " " != content[-1]
|
||||
):
|
||||
content += " "
|
||||
# escape content
|
||||
content = html.escape(content)
|
||||
end_html.extend(content)
|
||||
if b_with:
|
||||
end_html.extend("</b>")
|
||||
if "<td></td>" == tag:
|
||||
end_html.append("</td>")
|
||||
else:
|
||||
end_html.append(tag)
|
||||
td_index += 1
|
||||
else:
|
||||
end_html.append(tag)
|
||||
return "".join(end_html), end_html
|
||||
|
||||
def get_pred_html_master(self, pred_structures, matched_index, ocr_contents):
|
||||
end_html = []
|
||||
td_index = 0
|
||||
for token in pred_structures:
|
||||
if "</td>" in token:
|
||||
txt = ""
|
||||
b_with = False
|
||||
if td_index in matched_index.keys():
|
||||
if (
|
||||
"<b>" in ocr_contents[matched_index[td_index][0]]
|
||||
and len(matched_index[td_index]) > 1
|
||||
):
|
||||
b_with = True
|
||||
for i, td_index_index in enumerate(matched_index[td_index]):
|
||||
content = ocr_contents[td_index_index][0]
|
||||
if len(matched_index[td_index]) > 1:
|
||||
if len(content) == 0:
|
||||
continue
|
||||
if content[0] == " ":
|
||||
content = content[1:]
|
||||
if "<b>" in content:
|
||||
content = content[3:]
|
||||
if "</b>" in content:
|
||||
content = content[:-4]
|
||||
if len(content) == 0:
|
||||
continue
|
||||
if (
|
||||
i != len(matched_index[td_index]) - 1
|
||||
and " " != content[-1]
|
||||
):
|
||||
content += " "
|
||||
txt += content
|
||||
if b_with:
|
||||
txt = "<b>{}</b>".format(txt)
|
||||
if "<td></td>" == token:
|
||||
token = "<td>{}</td>".format(txt)
|
||||
else:
|
||||
token = "{}</td>".format(txt)
|
||||
td_index += 1
|
||||
token = deal_eb_token(token)
|
||||
end_html.append(token)
|
||||
html = "".join(end_html)
|
||||
html = deal_bb(html)
|
||||
return html, end_html
|
||||
|
||||
def _filter_ocr_result(self, pred_bboxes, dt_boxes, rec_res):
|
||||
y1 = pred_bboxes[:, 1::2].min()
|
||||
new_dt_boxes = []
|
||||
new_rec_res = []
|
||||
|
||||
for box, rec in zip(dt_boxes, rec_res):
|
||||
if np.max(box[1::2]) < y1:
|
||||
continue
|
||||
new_dt_boxes.append(box)
|
||||
new_rec_res.append(rec)
|
||||
return new_dt_boxes, new_rec_res
|
||||
207
ppstructure/table/predict_structure.py
Executable file
207
ppstructure/table/predict_structure.py
Executable file
@@ -0,0 +1,207 @@
|
||||
# Copyright (c) 2020 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.insert(0, os.path.abspath(os.path.join(__dir__, "../..")))
|
||||
|
||||
os.environ["FLAGS_allocator_strategy"] = "auto_growth"
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import time
|
||||
import json
|
||||
|
||||
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.utility import get_image_file_list, check_and_read
|
||||
from ppocr.utils.visual import draw_rectangle
|
||||
from ppstructure.utility import parse_args
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def build_pre_process_list(args):
|
||||
resize_op = {
|
||||
"ResizeTableImage": {
|
||||
"max_len": args.table_max_len,
|
||||
}
|
||||
}
|
||||
pad_op = {"PaddingTableImage": {"size": [args.table_max_len, args.table_max_len]}}
|
||||
normalize_op = {
|
||||
"NormalizeImage": {
|
||||
"std": (
|
||||
[0.229, 0.224, 0.225]
|
||||
if args.table_algorithm not in ["TableMaster"]
|
||||
else [0.5, 0.5, 0.5]
|
||||
),
|
||||
"mean": (
|
||||
[0.485, 0.456, 0.406]
|
||||
if args.table_algorithm not in ["TableMaster"]
|
||||
else [0.5, 0.5, 0.5]
|
||||
),
|
||||
"scale": "1./255.",
|
||||
"order": "hwc",
|
||||
}
|
||||
}
|
||||
to_chw_op = {"ToCHWImage": None}
|
||||
keep_keys_op = {"KeepKeys": {"keep_keys": ["image", "shape"]}}
|
||||
if args.table_algorithm not in ["TableMaster"]:
|
||||
pre_process_list = [resize_op, normalize_op, pad_op, to_chw_op, keep_keys_op]
|
||||
else:
|
||||
pre_process_list = [resize_op, pad_op, normalize_op, to_chw_op, keep_keys_op]
|
||||
return pre_process_list
|
||||
|
||||
|
||||
class TableStructurer(object):
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.use_onnx = args.use_onnx
|
||||
pre_process_list = build_pre_process_list(args)
|
||||
if args.table_algorithm not in ["TableMaster"]:
|
||||
postprocess_params = {
|
||||
"name": "TableLabelDecode",
|
||||
"character_dict_path": args.table_char_dict_path,
|
||||
"merge_no_span_structure": args.merge_no_span_structure,
|
||||
}
|
||||
else:
|
||||
postprocess_params = {
|
||||
"name": "TableMasterLabelDecode",
|
||||
"character_dict_path": args.table_char_dict_path,
|
||||
"box_shape": "pad",
|
||||
"merge_no_span_structure": args.merge_no_span_structure,
|
||||
}
|
||||
|
||||
self.preprocess_op = create_operators(pre_process_list)
|
||||
self.postprocess_op = build_post_process(postprocess_params)
|
||||
(
|
||||
self.predictor,
|
||||
self.input_tensor,
|
||||
self.output_tensors,
|
||||
self.config,
|
||||
) = utility.create_predictor(args, "table", logger)
|
||||
|
||||
if args.benchmark:
|
||||
import auto_log
|
||||
|
||||
pid = os.getpid()
|
||||
gpu_id = utility.get_infer_gpuid()
|
||||
self.autolog = auto_log.AutoLogger(
|
||||
model_name="table",
|
||||
model_precision=args.precision,
|
||||
batch_size=1,
|
||||
data_shape="dynamic",
|
||||
save_path=None, # args.save_log_path,
|
||||
inference_config=self.config,
|
||||
pids=pid,
|
||||
process_name=None,
|
||||
gpu_ids=gpu_id if args.use_gpu else None,
|
||||
time_keys=["preprocess_time", "inference_time", "postprocess_time"],
|
||||
warmup=0,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
def __call__(self, img):
|
||||
starttime = time.time()
|
||||
if self.args.benchmark:
|
||||
self.autolog.times.start()
|
||||
|
||||
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()
|
||||
if self.args.benchmark:
|
||||
self.autolog.times.stamp()
|
||||
if self.use_onnx:
|
||||
input_dict = {}
|
||||
input_dict[self.input_tensor.name] = img
|
||||
outputs = self.predictor.run(self.output_tensors, input_dict)
|
||||
else:
|
||||
self.input_tensor.copy_from_cpu(img)
|
||||
self.predictor.run()
|
||||
outputs = []
|
||||
for output_tensor in self.output_tensors:
|
||||
output = output_tensor.copy_to_cpu()
|
||||
outputs.append(output)
|
||||
if self.args.benchmark:
|
||||
self.autolog.times.stamp()
|
||||
|
||||
preds = {}
|
||||
preds["structure_probs"] = outputs[1]
|
||||
preds["loc_preds"] = outputs[0]
|
||||
|
||||
shape_list = np.expand_dims(data[-1], axis=0)
|
||||
post_result = self.postprocess_op(preds, [shape_list])
|
||||
|
||||
structure_str_list = post_result["structure_batch_list"][0]
|
||||
bbox_list = post_result["bbox_batch_list"][0]
|
||||
structure_str_list = structure_str_list[0]
|
||||
structure_str_list = (
|
||||
["<html>", "<body>", "<table>"]
|
||||
+ structure_str_list
|
||||
+ ["</table>", "</body>", "</html>"]
|
||||
)
|
||||
elapse = time.time() - starttime
|
||||
if self.args.benchmark:
|
||||
self.autolog.times.end(stamp=True)
|
||||
return (structure_str_list, bbox_list), elapse
|
||||
|
||||
|
||||
def main(args):
|
||||
image_file_list = get_image_file_list(args.image_dir)
|
||||
table_structurer = TableStructurer(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(image_file)
|
||||
if not flag:
|
||||
img = cv2.imread(image_file)
|
||||
if img is None:
|
||||
logger.info("error in loading image:{}".format(image_file))
|
||||
continue
|
||||
structure_res, elapse = table_structurer(img)
|
||||
structure_str_list, bbox_list = structure_res
|
||||
bbox_list_str = json.dumps(bbox_list.tolist())
|
||||
logger.info("result: {}, {}".format(structure_str_list, bbox_list_str))
|
||||
f_w.write("result: {}, {}\n".format(structure_str_list, bbox_list_str))
|
||||
|
||||
if len(bbox_list) > 0 and len(bbox_list[0]) == 4:
|
||||
img = draw_rectangle(image_file, bbox_list)
|
||||
else:
|
||||
img = utility.draw_boxes(img, bbox_list)
|
||||
img_save_path = os.path.join(args.output, os.path.basename(image_file))
|
||||
cv2.imwrite(img_save_path, img)
|
||||
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 args.benchmark:
|
||||
table_structurer.autolog.report()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
241
ppstructure/table/predict_table.py
Normal file
241
ppstructure/table/predict_table.py
Normal file
@@ -0,0 +1,241 @@
|
||||
# Copyright (c) 2020 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.insert(0, 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"
|
||||
import cv2
|
||||
import copy
|
||||
import logging
|
||||
import numpy as np
|
||||
import time
|
||||
import tools.infer.predict_rec as predict_rec
|
||||
import tools.infer.predict_det as predict_det
|
||||
import tools.infer.utility as utility
|
||||
from tools.infer.predict_system import sorted_boxes
|
||||
from ppocr.utils.utility import get_image_file_list, check_and_read
|
||||
from ppocr.utils.logging import get_logger
|
||||
from ppstructure.table.matcher import TableMatch
|
||||
from ppstructure.table.table_master_match import TableMasterMatcher
|
||||
from ppstructure.utility import parse_args
|
||||
import ppstructure.table.predict_structure as predict_strture
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
def expand(pix, det_box, shape):
|
||||
x0, y0, x1, y1 = det_box
|
||||
# print(shape)
|
||||
h, w, c = shape
|
||||
tmp_x0 = x0 - pix
|
||||
tmp_x1 = x1 + pix
|
||||
tmp_y0 = y0 - pix
|
||||
tmp_y1 = y1 + pix
|
||||
x0_ = tmp_x0 if tmp_x0 >= 0 else 0
|
||||
x1_ = tmp_x1 if tmp_x1 <= w else w
|
||||
y0_ = tmp_y0 if tmp_y0 >= 0 else 0
|
||||
y1_ = tmp_y1 if tmp_y1 <= h else h
|
||||
return x0_, y0_, x1_, y1_
|
||||
|
||||
|
||||
class TableSystem(object):
|
||||
def __init__(self, args, text_detector=None, text_recognizer=None):
|
||||
self.args = args
|
||||
if not args.show_log:
|
||||
logger.setLevel(logging.INFO)
|
||||
benchmark_tmp = False
|
||||
if args.benchmark:
|
||||
benchmark_tmp = args.benchmark
|
||||
args.benchmark = False
|
||||
self.text_detector = (
|
||||
predict_det.TextDetector(copy.deepcopy(args))
|
||||
if text_detector is None
|
||||
else text_detector
|
||||
)
|
||||
self.text_recognizer = (
|
||||
predict_rec.TextRecognizer(copy.deepcopy(args))
|
||||
if text_recognizer is None
|
||||
else text_recognizer
|
||||
)
|
||||
if benchmark_tmp:
|
||||
args.benchmark = True
|
||||
self.table_structurer = predict_strture.TableStructurer(args)
|
||||
if args.table_algorithm in ["TableMaster"]:
|
||||
self.match = TableMasterMatcher()
|
||||
else:
|
||||
self.match = TableMatch(filter_ocr_result=True)
|
||||
|
||||
(
|
||||
self.predictor,
|
||||
self.input_tensor,
|
||||
self.output_tensors,
|
||||
self.config,
|
||||
) = utility.create_predictor(args, "table", logger)
|
||||
|
||||
def __call__(self, img, return_ocr_result_in_table=False):
|
||||
result = dict()
|
||||
time_dict = {"det": 0, "rec": 0, "table": 0, "all": 0, "match": 0}
|
||||
start = time.time()
|
||||
structure_res, elapse = self._structure(copy.deepcopy(img))
|
||||
result["cell_bbox"] = structure_res[1].tolist()
|
||||
time_dict["table"] = elapse
|
||||
|
||||
dt_boxes, rec_res, det_elapse, rec_elapse = self._ocr(copy.deepcopy(img))
|
||||
time_dict["det"] = det_elapse
|
||||
time_dict["rec"] = rec_elapse
|
||||
|
||||
if return_ocr_result_in_table:
|
||||
result["boxes"] = [x.tolist() for x in dt_boxes]
|
||||
result["rec_res"] = rec_res
|
||||
|
||||
tic = time.time()
|
||||
pred_html = self.match(structure_res, dt_boxes, rec_res)
|
||||
toc = time.time()
|
||||
time_dict["match"] = toc - tic
|
||||
result["html"] = pred_html
|
||||
end = time.time()
|
||||
time_dict["all"] = end - start
|
||||
return result, time_dict
|
||||
|
||||
def _structure(self, img):
|
||||
structure_res, elapse = self.table_structurer(copy.deepcopy(img))
|
||||
return structure_res, elapse
|
||||
|
||||
def _ocr(self, img):
|
||||
h, w = img.shape[:2]
|
||||
dt_boxes, det_elapse = self.text_detector(copy.deepcopy(img))
|
||||
dt_boxes = sorted_boxes(dt_boxes)
|
||||
|
||||
r_boxes = []
|
||||
for box in dt_boxes:
|
||||
x_min = max(0, box[:, 0].min() - 1)
|
||||
x_max = min(w, box[:, 0].max() + 1)
|
||||
y_min = max(0, box[:, 1].min() - 1)
|
||||
y_max = min(h, box[:, 1].max() + 1)
|
||||
box = [x_min, y_min, x_max, y_max]
|
||||
r_boxes.append(box)
|
||||
dt_boxes = np.array(r_boxes)
|
||||
logger.debug("dt_boxes num : {}, elapse : {}".format(len(dt_boxes), det_elapse))
|
||||
if dt_boxes is None:
|
||||
return None, None
|
||||
|
||||
img_crop_list = []
|
||||
for i in range(len(dt_boxes)):
|
||||
det_box = dt_boxes[i]
|
||||
x0, y0, x1, y1 = expand(2, det_box, img.shape)
|
||||
text_rect = img[int(y0) : int(y1), int(x0) : int(x1), :]
|
||||
img_crop_list.append(text_rect)
|
||||
rec_res, rec_elapse = self.text_recognizer(img_crop_list)
|
||||
logger.debug("rec_res num : {}, elapse : {}".format(len(rec_res), rec_elapse))
|
||||
return dt_boxes, rec_res, det_elapse, rec_elapse
|
||||
|
||||
|
||||
def to_excel(html_table, excel_path):
|
||||
from tablepyxl import tablepyxl
|
||||
|
||||
tablepyxl.document_to_xl(html_table, excel_path)
|
||||
|
||||
|
||||
def main(args):
|
||||
image_file_list = get_image_file_list(args.image_dir)
|
||||
image_file_list = image_file_list[args.process_id :: args.total_process_num]
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
table_sys = TableSystem(args)
|
||||
img_num = len(image_file_list)
|
||||
|
||||
f_html = open(os.path.join(args.output, "show.html"), mode="w", encoding="utf-8")
|
||||
f_html.write("<html>\n<body>\n")
|
||||
f_html.write('<table border="1">\n')
|
||||
f_html.write(
|
||||
'<meta http-equiv="Content-Type" content="text/html; charset=utf-8" />'
|
||||
)
|
||||
f_html.write("<tr>\n")
|
||||
f_html.write("<td>img name\n")
|
||||
f_html.write("<td>ori image</td>")
|
||||
f_html.write("<td>table html</td>")
|
||||
f_html.write("<td>cell box</td>")
|
||||
f_html.write("</tr>\n")
|
||||
|
||||
for i, image_file in enumerate(image_file_list):
|
||||
logger.info("[{}/{}] {}".format(i, img_num, image_file))
|
||||
img, flag, _ = check_and_read(image_file)
|
||||
excel_path = os.path.join(
|
||||
args.output, os.path.basename(image_file).split(".")[0] + ".xlsx"
|
||||
)
|
||||
if not flag:
|
||||
img = cv2.imread(image_file)
|
||||
if img is None:
|
||||
logger.error("error in loading image:{}".format(image_file))
|
||||
continue
|
||||
starttime = time.time()
|
||||
pred_res, _ = table_sys(img)
|
||||
pred_html = pred_res["html"]
|
||||
logger.info(pred_html)
|
||||
to_excel(pred_html, excel_path)
|
||||
logger.info("excel saved to {}".format(excel_path))
|
||||
elapse = time.time() - starttime
|
||||
logger.info("Predict time : {:.3f}s".format(elapse))
|
||||
|
||||
if len(pred_res["cell_bbox"]) > 0 and len(pred_res["cell_bbox"][0]) == 4:
|
||||
img = predict_strture.draw_rectangle(image_file, pred_res["cell_bbox"])
|
||||
else:
|
||||
img = utility.draw_boxes(img, pred_res["cell_bbox"])
|
||||
img_save_path = os.path.join(args.output, os.path.basename(image_file))
|
||||
cv2.imwrite(img_save_path, img)
|
||||
|
||||
f_html.write("<tr>\n")
|
||||
f_html.write(f"<td> {os.path.basename(image_file)} <br/>\n")
|
||||
f_html.write(f'<td><img src="{image_file}" width=640></td>\n')
|
||||
f_html.write(
|
||||
'<td><table border="1">'
|
||||
+ pred_html.replace("<html><body><table>", "").replace(
|
||||
"</table></body></html>", ""
|
||||
)
|
||||
+ "</table></td>\n"
|
||||
)
|
||||
f_html.write(f'<td><img src="{os.path.basename(image_file)}" width=640></td>\n')
|
||||
f_html.write("</tr>\n")
|
||||
f_html.write("</table>\n")
|
||||
f_html.close()
|
||||
|
||||
if args.benchmark:
|
||||
table_sys.table_structurer.autolog.report()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
if args.use_mp:
|
||||
import subprocess
|
||||
|
||||
p_list = []
|
||||
total_process_num = args.total_process_num
|
||||
for process_id in range(total_process_num):
|
||||
cmd = (
|
||||
[sys.executable, "-u"]
|
||||
+ sys.argv
|
||||
+ ["--process_id={}".format(process_id), "--use_mp={}".format(False)]
|
||||
)
|
||||
p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
|
||||
p_list.append(p)
|
||||
for p in p_list:
|
||||
p.wait()
|
||||
else:
|
||||
main(args)
|
||||
995
ppstructure/table/table_master_match.py
Normal file
995
ppstructure/table/table_master_match.py
Normal file
@@ -0,0 +1,995 @@
|
||||
# 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.
|
||||
"""
|
||||
This code is refer from:
|
||||
https://github.com/JiaquanYe/TableMASTER-mmocr/blob/master/table_recognition/match.py
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import cv2
|
||||
import glob
|
||||
import copy
|
||||
import math
|
||||
import pickle
|
||||
import numpy as np
|
||||
|
||||
from shapely.geometry import Polygon, MultiPoint
|
||||
|
||||
"""
|
||||
Useful function in matching.
|
||||
"""
|
||||
|
||||
|
||||
def remove_empty_bboxes(bboxes):
|
||||
"""
|
||||
remove [0., 0., 0., 0.] in structure master bboxes.
|
||||
len(bboxes.shape) must be 2.
|
||||
:param bboxes:
|
||||
:return:
|
||||
"""
|
||||
new_bboxes = []
|
||||
for bbox in bboxes:
|
||||
if sum(bbox) == 0.0:
|
||||
continue
|
||||
new_bboxes.append(bbox)
|
||||
return np.array(new_bboxes)
|
||||
|
||||
|
||||
def xywh2xyxy(bboxes):
|
||||
if len(bboxes.shape) == 1:
|
||||
new_bboxes = np.empty_like(bboxes)
|
||||
new_bboxes[0] = bboxes[0] - bboxes[2] / 2
|
||||
new_bboxes[1] = bboxes[1] - bboxes[3] / 2
|
||||
new_bboxes[2] = bboxes[0] + bboxes[2] / 2
|
||||
new_bboxes[3] = bboxes[1] + bboxes[3] / 2
|
||||
return new_bboxes
|
||||
elif len(bboxes.shape) == 2:
|
||||
new_bboxes = np.empty_like(bboxes)
|
||||
new_bboxes[:, 0] = bboxes[:, 0] - bboxes[:, 2] / 2
|
||||
new_bboxes[:, 1] = bboxes[:, 1] - bboxes[:, 3] / 2
|
||||
new_bboxes[:, 2] = bboxes[:, 0] + bboxes[:, 2] / 2
|
||||
new_bboxes[:, 3] = bboxes[:, 1] + bboxes[:, 3] / 2
|
||||
return new_bboxes
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
|
||||
def xyxy2xywh(bboxes):
|
||||
if len(bboxes.shape) == 1:
|
||||
new_bboxes = np.empty_like(bboxes)
|
||||
new_bboxes[0] = bboxes[0] + (bboxes[2] - bboxes[0]) / 2
|
||||
new_bboxes[1] = bboxes[1] + (bboxes[3] - bboxes[1]) / 2
|
||||
new_bboxes[2] = bboxes[2] - bboxes[0]
|
||||
new_bboxes[3] = bboxes[3] - bboxes[1]
|
||||
return new_bboxes
|
||||
elif len(bboxes.shape) == 2:
|
||||
new_bboxes = np.empty_like(bboxes)
|
||||
new_bboxes[:, 0] = bboxes[:, 0] + (bboxes[:, 2] - bboxes[:, 0]) / 2
|
||||
new_bboxes[:, 1] = bboxes[:, 1] + (bboxes[:, 3] - bboxes[:, 1]) / 2
|
||||
new_bboxes[:, 2] = bboxes[:, 2] - bboxes[:, 0]
|
||||
new_bboxes[:, 3] = bboxes[:, 3] - bboxes[:, 1]
|
||||
return new_bboxes
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
|
||||
def pickle_load(path, prefix="end2end"):
|
||||
if os.path.isfile(path):
|
||||
data = pickle.load(open(path, "rb"))
|
||||
elif os.path.isdir(path):
|
||||
data = dict()
|
||||
search_path = os.path.join(path, "{}_*.pkl".format(prefix))
|
||||
pkls = glob.glob(search_path)
|
||||
for pkl in pkls:
|
||||
this_data = pickle.load(open(pkl, "rb"))
|
||||
data.update(this_data)
|
||||
else:
|
||||
raise ValueError
|
||||
return data
|
||||
|
||||
|
||||
def convert_coord(xyxy):
|
||||
"""
|
||||
Convert two points format to four points format.
|
||||
:param xyxy:
|
||||
:return:
|
||||
"""
|
||||
new_bbox = np.zeros([4, 2], dtype=np.float32)
|
||||
new_bbox[0, 0], new_bbox[0, 1] = xyxy[0], xyxy[1]
|
||||
new_bbox[1, 0], new_bbox[1, 1] = xyxy[2], xyxy[1]
|
||||
new_bbox[2, 0], new_bbox[2, 1] = xyxy[2], xyxy[3]
|
||||
new_bbox[3, 0], new_bbox[3, 1] = xyxy[0], xyxy[3]
|
||||
return new_bbox
|
||||
|
||||
|
||||
def cal_iou(bbox1, bbox2):
|
||||
bbox1_poly = Polygon(bbox1).convex_hull
|
||||
bbox2_poly = Polygon(bbox2).convex_hull
|
||||
union_poly = np.concatenate((bbox1, bbox2))
|
||||
|
||||
if not bbox1_poly.intersects(bbox2_poly):
|
||||
iou = 0
|
||||
else:
|
||||
inter_area = bbox1_poly.intersection(bbox2_poly).area
|
||||
union_area = MultiPoint(union_poly).convex_hull.area
|
||||
if union_area == 0:
|
||||
iou = 0
|
||||
else:
|
||||
iou = float(inter_area) / union_area
|
||||
return iou
|
||||
|
||||
|
||||
def cal_distance(p1, p2):
|
||||
delta_x = p1[0] - p2[0]
|
||||
delta_y = p1[1] - p2[1]
|
||||
d = math.sqrt((delta_x**2) + (delta_y**2))
|
||||
return d
|
||||
|
||||
|
||||
def is_inside(center_point, corner_point):
|
||||
"""
|
||||
Find if center_point inside the bbox(corner_point) or not.
|
||||
:param center_point: center point (x, y)
|
||||
:param corner_point: corner point ((x1,y1),(x2,y2))
|
||||
:return:
|
||||
"""
|
||||
x_flag = False
|
||||
y_flag = False
|
||||
if (center_point[0] >= corner_point[0][0]) and (
|
||||
center_point[0] <= corner_point[1][0]
|
||||
):
|
||||
x_flag = True
|
||||
if (center_point[1] >= corner_point[0][1]) and (
|
||||
center_point[1] <= corner_point[1][1]
|
||||
):
|
||||
y_flag = True
|
||||
if x_flag and y_flag:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def find_no_match(match_list, all_end2end_nums, type="end2end"):
|
||||
"""
|
||||
Find out no match end2end bbox in previous match list.
|
||||
:param match_list: matching pairs.
|
||||
:param all_end2end_nums: numbers of end2end_xywh
|
||||
:param type: 'end2end' corresponding to idx 0, 'master' corresponding to idx 1.
|
||||
:return: no match pse bbox index list
|
||||
"""
|
||||
if type == "end2end":
|
||||
idx = 0
|
||||
elif type == "master":
|
||||
idx = 1
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
no_match_indexs = []
|
||||
# m[0] is end2end index m[1] is master index
|
||||
matched_bbox_indexs = [m[idx] for m in match_list]
|
||||
for n in range(all_end2end_nums):
|
||||
if n not in matched_bbox_indexs:
|
||||
no_match_indexs.append(n)
|
||||
return no_match_indexs
|
||||
|
||||
|
||||
def is_abs_lower_than_threshold(this_bbox, target_bbox, threshold=3):
|
||||
# only consider y axis, for grouping in row.
|
||||
delta = abs(this_bbox[1] - target_bbox[1])
|
||||
if delta < threshold:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def sort_line_bbox(g, bg):
|
||||
"""
|
||||
Sorted the bbox in the same line(group)
|
||||
compare coord 'x' value, where 'y' value is closed in the same group.
|
||||
:param g: index in the same group
|
||||
:param bg: bbox in the same group
|
||||
:return:
|
||||
"""
|
||||
|
||||
xs = [bg_item[0] for bg_item in bg]
|
||||
xs_sorted = sorted(xs)
|
||||
|
||||
g_sorted = [None] * len(xs_sorted)
|
||||
bg_sorted = [None] * len(xs_sorted)
|
||||
for g_item, bg_item in zip(g, bg):
|
||||
idx = xs_sorted.index(bg_item[0])
|
||||
bg_sorted[idx] = bg_item
|
||||
g_sorted[idx] = g_item
|
||||
|
||||
return g_sorted, bg_sorted
|
||||
|
||||
|
||||
def flatten(sorted_groups, sorted_bbox_groups):
|
||||
idxs = []
|
||||
bboxes = []
|
||||
for group, bbox_group in zip(sorted_groups, sorted_bbox_groups):
|
||||
for g, bg in zip(group, bbox_group):
|
||||
idxs.append(g)
|
||||
bboxes.append(bg)
|
||||
return idxs, bboxes
|
||||
|
||||
|
||||
def sort_bbox(end2end_xywh_bboxes, no_match_end2end_indexes):
|
||||
"""
|
||||
This function will group the render end2end bboxes in row.
|
||||
:param end2end_xywh_bboxes:
|
||||
:param no_match_end2end_indexes:
|
||||
:return:
|
||||
"""
|
||||
groups = []
|
||||
bbox_groups = []
|
||||
for index, end2end_xywh_bbox in zip(no_match_end2end_indexes, end2end_xywh_bboxes):
|
||||
this_bbox = end2end_xywh_bbox
|
||||
if len(groups) == 0:
|
||||
groups.append([index])
|
||||
bbox_groups.append([this_bbox])
|
||||
else:
|
||||
flag = False
|
||||
for g, bg in zip(groups, bbox_groups):
|
||||
# this_bbox is belong to bg's row or not
|
||||
if is_abs_lower_than_threshold(this_bbox, bg[0]):
|
||||
g.append(index)
|
||||
bg.append(this_bbox)
|
||||
flag = True
|
||||
break
|
||||
if not flag:
|
||||
# this_bbox is not belong to bg's row, create a row.
|
||||
groups.append([index])
|
||||
bbox_groups.append([this_bbox])
|
||||
|
||||
# sorted bboxes in a group
|
||||
tmp_groups, tmp_bbox_groups = [], []
|
||||
for g, bg in zip(groups, bbox_groups):
|
||||
g_sorted, bg_sorted = sort_line_bbox(g, bg)
|
||||
tmp_groups.append(g_sorted)
|
||||
tmp_bbox_groups.append(bg_sorted)
|
||||
|
||||
# sorted groups, sort by coord y's value.
|
||||
sorted_groups = [None] * len(tmp_groups)
|
||||
sorted_bbox_groups = [None] * len(tmp_bbox_groups)
|
||||
ys = [bg[0][1] for bg in tmp_bbox_groups]
|
||||
sorted_ys = sorted(ys)
|
||||
for g, bg in zip(tmp_groups, tmp_bbox_groups):
|
||||
idx = sorted_ys.index(bg[0][1])
|
||||
sorted_groups[idx] = g
|
||||
sorted_bbox_groups[idx] = bg
|
||||
|
||||
# flatten, get final result
|
||||
end2end_sorted_idx_list, end2end_sorted_bbox_list = flatten(
|
||||
sorted_groups, sorted_bbox_groups
|
||||
)
|
||||
|
||||
return (
|
||||
end2end_sorted_idx_list,
|
||||
end2end_sorted_bbox_list,
|
||||
sorted_groups,
|
||||
sorted_bbox_groups,
|
||||
)
|
||||
|
||||
|
||||
def get_bboxes_list(end2end_result, structure_master_result):
|
||||
"""
|
||||
This function is use to convert end2end results and structure master results to
|
||||
List of xyxy bbox format and List of xywh bbox format
|
||||
:param end2end_result: bbox's format is xyxy
|
||||
:param structure_master_result: bbox's format is xywh
|
||||
:return: 4 kind list of bbox ()
|
||||
"""
|
||||
# end2end
|
||||
end2end_xyxy_list = []
|
||||
end2end_xywh_list = []
|
||||
for end2end_item in end2end_result:
|
||||
src_bbox = end2end_item["bbox"]
|
||||
end2end_xyxy_list.append(src_bbox)
|
||||
xywh_bbox = xyxy2xywh(src_bbox)
|
||||
end2end_xywh_list.append(xywh_bbox)
|
||||
end2end_xyxy_bboxes = np.array(end2end_xyxy_list)
|
||||
end2end_xywh_bboxes = np.array(end2end_xywh_list)
|
||||
|
||||
# structure master
|
||||
src_bboxes = structure_master_result["bbox"]
|
||||
src_bboxes = remove_empty_bboxes(src_bboxes)
|
||||
structure_master_xyxy_bboxes = src_bboxes
|
||||
xywh_bbox = xyxy2xywh(src_bboxes)
|
||||
structure_master_xywh_bboxes = xywh_bbox
|
||||
|
||||
return (
|
||||
end2end_xyxy_bboxes,
|
||||
end2end_xywh_bboxes,
|
||||
structure_master_xywh_bboxes,
|
||||
structure_master_xyxy_bboxes,
|
||||
)
|
||||
|
||||
|
||||
def center_rule_match(end2end_xywh_bboxes, structure_master_xyxy_bboxes):
|
||||
"""
|
||||
Judge end2end Bbox's center point is inside structure master Bbox or not,
|
||||
if end2end Bbox's center is in structure master Bbox, get matching pair.
|
||||
:param end2end_xywh_bboxes:
|
||||
:param structure_master_xyxy_bboxes:
|
||||
:return: match pairs list, e.g. [[0,1], [1,2], ...]
|
||||
"""
|
||||
match_pairs_list = []
|
||||
for i, end2end_xywh in enumerate(end2end_xywh_bboxes):
|
||||
for j, master_xyxy in enumerate(structure_master_xyxy_bboxes):
|
||||
x_end2end, y_end2end = end2end_xywh[0], end2end_xywh[1]
|
||||
x_master1, y_master1, x_master2, y_master2 = (
|
||||
master_xyxy[0],
|
||||
master_xyxy[1],
|
||||
master_xyxy[2],
|
||||
master_xyxy[3],
|
||||
)
|
||||
center_point_end2end = (x_end2end, y_end2end)
|
||||
corner_point_master = ((x_master1, y_master1), (x_master2, y_master2))
|
||||
if is_inside(center_point_end2end, corner_point_master):
|
||||
match_pairs_list.append([i, j])
|
||||
return match_pairs_list
|
||||
|
||||
|
||||
def iou_rule_match(
|
||||
end2end_xyxy_bboxes, end2end_xyxy_indexes, structure_master_xyxy_bboxes
|
||||
):
|
||||
"""
|
||||
Use iou to find matching list.
|
||||
choose max iou value bbox as match pair.
|
||||
:param end2end_xyxy_bboxes:
|
||||
:param end2end_xyxy_indexes: original end2end indexes.
|
||||
:param structure_master_xyxy_bboxes:
|
||||
:return: match pairs list, e.g. [[0,1], [1,2], ...]
|
||||
"""
|
||||
match_pair_list = []
|
||||
for end2end_xyxy_index, end2end_xyxy in zip(
|
||||
end2end_xyxy_indexes, end2end_xyxy_bboxes
|
||||
):
|
||||
max_iou = 0
|
||||
max_match = [None, None]
|
||||
for j, master_xyxy in enumerate(structure_master_xyxy_bboxes):
|
||||
end2end_4xy = convert_coord(end2end_xyxy)
|
||||
master_4xy = convert_coord(master_xyxy)
|
||||
iou = cal_iou(end2end_4xy, master_4xy)
|
||||
if iou > max_iou:
|
||||
max_match[0], max_match[1] = end2end_xyxy_index, j
|
||||
max_iou = iou
|
||||
|
||||
if max_match[0] is None:
|
||||
# no match
|
||||
continue
|
||||
match_pair_list.append(max_match)
|
||||
return match_pair_list
|
||||
|
||||
|
||||
def distance_rule_match(end2end_indexes, end2end_bboxes, master_indexes, master_bboxes):
|
||||
"""
|
||||
Get matching between no-match end2end bboxes and no-match master bboxes.
|
||||
Use min distance to match.
|
||||
This rule will only run (no-match end2end nums > 0) and (no-match master nums > 0)
|
||||
It will Return master_bboxes_nums match-pairs.
|
||||
:param end2end_indexes:
|
||||
:param end2end_bboxes:
|
||||
:param master_indexes:
|
||||
:param master_bboxes:
|
||||
:return: match_pairs list, e.g. [[0,1], [1,2], ...]
|
||||
"""
|
||||
min_match_list = []
|
||||
for j, master_bbox in zip(master_indexes, master_bboxes):
|
||||
min_distance = np.inf
|
||||
min_match = [0, 0] # i, j
|
||||
for i, end2end_bbox in zip(end2end_indexes, end2end_bboxes):
|
||||
x_end2end, y_end2end = end2end_bbox[0], end2end_bbox[1]
|
||||
x_master, y_master = master_bbox[0], master_bbox[1]
|
||||
end2end_point = (x_end2end, y_end2end)
|
||||
master_point = (x_master, y_master)
|
||||
dist = cal_distance(master_point, end2end_point)
|
||||
if dist < min_distance:
|
||||
min_match[0], min_match[1] = i, j
|
||||
min_distance = dist
|
||||
min_match_list.append(min_match)
|
||||
return min_match_list
|
||||
|
||||
|
||||
def extra_match(no_match_end2end_indexes, master_bbox_nums):
|
||||
"""
|
||||
This function will create some virtual master bboxes,
|
||||
and get match with the no match end2end indexes.
|
||||
:param no_match_end2end_indexes:
|
||||
:param master_bbox_nums:
|
||||
:return:
|
||||
"""
|
||||
end_nums = len(no_match_end2end_indexes) + master_bbox_nums
|
||||
extra_match_list = []
|
||||
for i in range(master_bbox_nums, end_nums):
|
||||
end2end_index = no_match_end2end_indexes[i - master_bbox_nums]
|
||||
extra_match_list.append([end2end_index, i])
|
||||
return extra_match_list
|
||||
|
||||
|
||||
def get_match_dict(match_list):
|
||||
"""
|
||||
Convert match_list to a dict, where key is master bbox's index, value is end2end bbox index.
|
||||
:param match_list:
|
||||
:return:
|
||||
"""
|
||||
match_dict = dict()
|
||||
for match_pair in match_list:
|
||||
end2end_index, master_index = match_pair[0], match_pair[1]
|
||||
if master_index not in match_dict.keys():
|
||||
match_dict[master_index] = [end2end_index]
|
||||
else:
|
||||
match_dict[master_index].append(end2end_index)
|
||||
return match_dict
|
||||
|
||||
|
||||
def deal_successive_space(text):
|
||||
"""
|
||||
deal successive space character for text
|
||||
1. Replace ' '*3 with '<space>' which is real space is text
|
||||
2. Remove ' ', which is split token, not true space
|
||||
3. Replace '<space>' with ' ', to get real text
|
||||
:param text:
|
||||
:return:
|
||||
"""
|
||||
text = text.replace(" " * 3, "<space>")
|
||||
text = text.replace(" ", "")
|
||||
text = text.replace("<space>", " ")
|
||||
return text
|
||||
|
||||
|
||||
def reduce_repeat_bb(text_list, break_token):
|
||||
"""
|
||||
convert ['<b>Local</b>', '<b>government</b>', '<b>unit</b>'] to ['<b>Local government unit</b>']
|
||||
PS: maybe style <i>Local</i> is also exist, too. it can be processed like this.
|
||||
:param text_list:
|
||||
:param break_token:
|
||||
:return:
|
||||
"""
|
||||
count = 0
|
||||
for text in text_list:
|
||||
if text.startswith("<b>"):
|
||||
count += 1
|
||||
if count == len(text_list):
|
||||
new_text_list = []
|
||||
for text in text_list:
|
||||
text = text.replace("<b>", "").replace("</b>", "")
|
||||
new_text_list.append(text)
|
||||
return ["<b>" + break_token.join(new_text_list) + "</b>"]
|
||||
else:
|
||||
return text_list
|
||||
|
||||
|
||||
def get_match_text_dict(match_dict, end2end_info, break_token=" "):
|
||||
match_text_dict = dict()
|
||||
for master_index, end2end_index_list in match_dict.items():
|
||||
text_list = [
|
||||
end2end_info[end2end_index]["text"] for end2end_index in end2end_index_list
|
||||
]
|
||||
text_list = reduce_repeat_bb(text_list, break_token)
|
||||
text = break_token.join(text_list)
|
||||
match_text_dict[master_index] = text
|
||||
return match_text_dict
|
||||
|
||||
|
||||
def merge_span_token(master_token_list):
|
||||
"""
|
||||
Merge the span style token (row span or col span).
|
||||
:param master_token_list:
|
||||
:return:
|
||||
"""
|
||||
new_master_token_list = []
|
||||
pointer = 0
|
||||
if master_token_list[-1] != "</tbody>":
|
||||
master_token_list.append("</tbody>")
|
||||
while master_token_list[pointer] != "</tbody>":
|
||||
try:
|
||||
if master_token_list[pointer] == "<td":
|
||||
if master_token_list[pointer + 1].startswith(
|
||||
" colspan="
|
||||
) or master_token_list[pointer + 1].startswith(" rowspan="):
|
||||
"""
|
||||
example:
|
||||
pattern <td colspan="3">
|
||||
'<td' + 'colspan=" "' + '>' + '</td>'
|
||||
"""
|
||||
tmp = "".join(master_token_list[pointer : pointer + 3 + 1])
|
||||
pointer += 4
|
||||
new_master_token_list.append(tmp)
|
||||
|
||||
elif master_token_list[pointer + 2].startswith(
|
||||
" colspan="
|
||||
) or master_token_list[pointer + 2].startswith(" rowspan="):
|
||||
"""
|
||||
example:
|
||||
pattern <td rowspan="2" colspan="3">
|
||||
'<td' + 'rowspan=" "' + 'colspan=" "' + '>' + '</td>'
|
||||
"""
|
||||
tmp = "".join(master_token_list[pointer : pointer + 4 + 1])
|
||||
pointer += 5
|
||||
new_master_token_list.append(tmp)
|
||||
|
||||
else:
|
||||
new_master_token_list.append(master_token_list[pointer])
|
||||
pointer += 1
|
||||
else:
|
||||
new_master_token_list.append(master_token_list[pointer])
|
||||
pointer += 1
|
||||
except:
|
||||
print("Break in merge...")
|
||||
break
|
||||
new_master_token_list.append("</tbody>")
|
||||
|
||||
return new_master_token_list
|
||||
|
||||
|
||||
def deal_eb_token(master_token):
|
||||
"""
|
||||
post process with <eb></eb>, <eb1></eb1>, ...
|
||||
emptyBboxTokenDict = {
|
||||
"[]": '<eb></eb>',
|
||||
"[' ']": '<eb1></eb1>',
|
||||
"['<b>', ' ', '</b>']": '<eb2></eb2>',
|
||||
"['\\u2028', '\\u2028']": '<eb3></eb3>',
|
||||
"['<sup>', ' ', '</sup>']": '<eb4></eb4>',
|
||||
"['<b>', '</b>']": '<eb5></eb5>',
|
||||
"['<i>', ' ', '</i>']": '<eb6></eb6>',
|
||||
"['<b>', '<i>', '</i>', '</b>']": '<eb7></eb7>',
|
||||
"['<b>', '<i>', ' ', '</i>', '</b>']": '<eb8></eb8>',
|
||||
"['<i>', '</i>']": '<eb9></eb9>',
|
||||
"['<b>', ' ', '\\u2028', ' ', '\\u2028', ' ', '</b>']": '<eb10></eb10>',
|
||||
}
|
||||
:param master_token:
|
||||
:return:
|
||||
"""
|
||||
master_token = master_token.replace("<eb></eb>", "<td></td>")
|
||||
master_token = master_token.replace("<eb1></eb1>", "<td> </td>")
|
||||
master_token = master_token.replace("<eb2></eb2>", "<td><b> </b></td>")
|
||||
master_token = master_token.replace("<eb3></eb3>", "<td>\u2028\u2028</td>")
|
||||
master_token = master_token.replace("<eb4></eb4>", "<td><sup> </sup></td>")
|
||||
master_token = master_token.replace("<eb5></eb5>", "<td><b></b></td>")
|
||||
master_token = master_token.replace("<eb6></eb6>", "<td><i> </i></td>")
|
||||
master_token = master_token.replace("<eb7></eb7>", "<td><b><i></i></b></td>")
|
||||
master_token = master_token.replace("<eb8></eb8>", "<td><b><i> </i></b></td>")
|
||||
master_token = master_token.replace("<eb9></eb9>", "<td><i></i></td>")
|
||||
master_token = master_token.replace(
|
||||
"<eb10></eb10>", "<td><b> \u2028 \u2028 </b></td>"
|
||||
)
|
||||
return master_token
|
||||
|
||||
|
||||
def insert_text_to_token(master_token_list, match_text_dict):
|
||||
"""
|
||||
Insert OCR text result to structure token.
|
||||
:param master_token_list:
|
||||
:param match_text_dict:
|
||||
:return:
|
||||
"""
|
||||
master_token_list = merge_span_token(master_token_list)
|
||||
merged_result_list = []
|
||||
text_count = 0
|
||||
for master_token in master_token_list:
|
||||
if master_token.startswith("<td"):
|
||||
if text_count > len(match_text_dict) - 1:
|
||||
text_count += 1
|
||||
continue
|
||||
elif text_count not in match_text_dict.keys():
|
||||
text_count += 1
|
||||
continue
|
||||
else:
|
||||
master_token = master_token.replace(
|
||||
"><", ">{}<".format(match_text_dict[text_count])
|
||||
)
|
||||
text_count += 1
|
||||
master_token = deal_eb_token(master_token)
|
||||
merged_result_list.append(master_token)
|
||||
|
||||
return "".join(merged_result_list)
|
||||
|
||||
|
||||
def deal_isolate_span(thead_part):
|
||||
"""
|
||||
Deal with isolate span cases in this function.
|
||||
It causes by wrong prediction in structure recognition model.
|
||||
eg. predict <td rowspan="2"></td> to <td></td> rowspan="2"></b></td>.
|
||||
:param thead_part:
|
||||
:return:
|
||||
"""
|
||||
# 1. find out isolate span tokens.
|
||||
isolate_pattern = (
|
||||
'<td></td> rowspan="(\d)+" colspan="(\d)+"></b></td>|'
|
||||
'<td></td> colspan="(\d)+" rowspan="(\d)+"></b></td>|'
|
||||
'<td></td> rowspan="(\d)+"></b></td>|'
|
||||
'<td></td> colspan="(\d)+"></b></td>'
|
||||
)
|
||||
isolate_iter = re.finditer(isolate_pattern, thead_part)
|
||||
isolate_list = [i.group() for i in isolate_iter]
|
||||
|
||||
# 2. find out span number, by step 1 results.
|
||||
span_pattern = (
|
||||
' rowspan="(\d)+" colspan="(\d)+"|'
|
||||
' colspan="(\d)+" rowspan="(\d)+"|'
|
||||
' rowspan="(\d)+"|'
|
||||
' colspan="(\d)+"'
|
||||
)
|
||||
corrected_list = []
|
||||
for isolate_item in isolate_list:
|
||||
span_part = re.search(span_pattern, isolate_item)
|
||||
spanStr_in_isolateItem = span_part.group()
|
||||
# 3. merge the span number into the span token format string.
|
||||
if spanStr_in_isolateItem is not None:
|
||||
corrected_item = "<td{}></td>".format(spanStr_in_isolateItem)
|
||||
corrected_list.append(corrected_item)
|
||||
else:
|
||||
corrected_list.append(None)
|
||||
|
||||
# 4. replace original isolated token.
|
||||
for corrected_item, isolate_item in zip(corrected_list, isolate_list):
|
||||
if corrected_item is not None:
|
||||
thead_part = thead_part.replace(isolate_item, corrected_item)
|
||||
else:
|
||||
pass
|
||||
return thead_part
|
||||
|
||||
|
||||
def deal_duplicate_bb(thead_part):
|
||||
"""
|
||||
Deal duplicate <b> or </b> after replace.
|
||||
Keep one <b></b> in a <td></td> token.
|
||||
:param thead_part:
|
||||
:return:
|
||||
"""
|
||||
# 1. find out <td></td> in <thead></thead>.
|
||||
td_pattern = (
|
||||
'<td rowspan="(\d)+" colspan="(\d)+">(.+?)</td>|'
|
||||
'<td colspan="(\d)+" rowspan="(\d)+">(.+?)</td>|'
|
||||
'<td rowspan="(\d)+">(.+?)</td>|'
|
||||
'<td colspan="(\d)+">(.+?)</td>|'
|
||||
"<td>(.*?)</td>"
|
||||
)
|
||||
td_iter = re.finditer(td_pattern, thead_part)
|
||||
td_list = [t.group() for t in td_iter]
|
||||
|
||||
# 2. is multiply <b></b> in <td></td> or not?
|
||||
new_td_list = []
|
||||
for td_item in td_list:
|
||||
if td_item.count("<b>") > 1 or td_item.count("</b>") > 1:
|
||||
# multiply <b></b> in <td></td> case.
|
||||
# 1. remove all <b></b>
|
||||
td_item = td_item.replace("<b>", "").replace("</b>", "")
|
||||
# 2. replace <tb> -> <tb><b>, </tb> -> </b></tb>.
|
||||
td_item = td_item.replace("<td>", "<td><b>").replace("</td>", "</b></td>")
|
||||
new_td_list.append(td_item)
|
||||
else:
|
||||
new_td_list.append(td_item)
|
||||
|
||||
# 3. replace original thead part.
|
||||
for td_item, new_td_item in zip(td_list, new_td_list):
|
||||
thead_part = thead_part.replace(td_item, new_td_item)
|
||||
return thead_part
|
||||
|
||||
|
||||
def deal_bb(result_token):
|
||||
"""
|
||||
In our opinion, <b></b> always occurs in <thead></thead> text's context.
|
||||
This function will find out all tokens in <thead></thead> and insert <b></b> by manual.
|
||||
:param result_token:
|
||||
:return:
|
||||
"""
|
||||
# find out <thead></thead> parts.
|
||||
thead_pattern = "<thead>(.*?)</thead>"
|
||||
if re.search(thead_pattern, result_token) is None:
|
||||
return result_token
|
||||
thead_part = re.search(thead_pattern, result_token).group()
|
||||
origin_thead_part = copy.deepcopy(thead_part)
|
||||
|
||||
# check "rowspan" or "colspan" occur in <thead></thead> parts or not .
|
||||
span_pattern = '<td rowspan="(\d)+" colspan="(\d)+">|<td colspan="(\d)+" rowspan="(\d)+">|<td rowspan="(\d)+">|<td colspan="(\d)+">'
|
||||
span_iter = re.finditer(span_pattern, thead_part)
|
||||
span_list = [s.group() for s in span_iter]
|
||||
has_span_in_head = True if len(span_list) > 0 else False
|
||||
|
||||
if not has_span_in_head:
|
||||
# <thead></thead> not include "rowspan" or "colspan" branch 1.
|
||||
# 1. replace <td> to <td><b>, and </td> to </b></td>
|
||||
# 2. it is possible to predict text include <b> or </b> by Text-line recognition,
|
||||
# so we replace <b><b> to <b>, and </b></b> to </b>
|
||||
thead_part = (
|
||||
thead_part.replace("<td>", "<td><b>")
|
||||
.replace("</td>", "</b></td>")
|
||||
.replace("<b><b>", "<b>")
|
||||
.replace("</b></b>", "</b>")
|
||||
)
|
||||
else:
|
||||
# <thead></thead> include "rowspan" or "colspan" branch 2.
|
||||
# Firstly, we deal rowspan or colspan cases.
|
||||
# 1. replace > to ><b>
|
||||
# 2. replace </td> to </b></td>
|
||||
# 3. it is possible to predict text include <b> or </b> by Text-line recognition,
|
||||
# so we replace <b><b> to <b>, and </b><b> to </b>
|
||||
|
||||
# Secondly, deal ordinary cases like branch 1
|
||||
|
||||
# replace ">" to "<b>"
|
||||
replaced_span_list = []
|
||||
for sp in span_list:
|
||||
replaced_span_list.append(sp.replace(">", "><b>"))
|
||||
for sp, rsp in zip(span_list, replaced_span_list):
|
||||
thead_part = thead_part.replace(sp, rsp)
|
||||
|
||||
# replace "</td>" to "</b></td>"
|
||||
thead_part = thead_part.replace("</td>", "</b></td>")
|
||||
|
||||
# remove duplicated <b> by re.sub
|
||||
mb_pattern = "(<b>)+"
|
||||
single_b_string = "<b>"
|
||||
thead_part = re.sub(mb_pattern, single_b_string, thead_part)
|
||||
|
||||
mgb_pattern = "(</b>)+"
|
||||
single_gb_string = "</b>"
|
||||
thead_part = re.sub(mgb_pattern, single_gb_string, thead_part)
|
||||
|
||||
# ordinary cases like branch 1
|
||||
thead_part = thead_part.replace("<td>", "<td><b>").replace("<b><b>", "<b>")
|
||||
|
||||
# convert <tb><b></b></tb> back to <tb></tb>, empty cell has no <b></b>.
|
||||
# but space cell(<tb> </tb>) is suitable for <td><b> </b></td>
|
||||
thead_part = thead_part.replace("<td><b></b></td>", "<td></td>")
|
||||
# deal with duplicated <b></b>
|
||||
thead_part = deal_duplicate_bb(thead_part)
|
||||
# deal with isolate span tokens, which causes by wrong predict by structure prediction.
|
||||
# eg.PMC5994107_011_00.png
|
||||
thead_part = deal_isolate_span(thead_part)
|
||||
# replace original result with new thead part.
|
||||
result_token = result_token.replace(origin_thead_part, thead_part)
|
||||
return result_token
|
||||
|
||||
|
||||
class Matcher:
|
||||
def __init__(self, end2end_file, structure_master_file):
|
||||
"""
|
||||
This class process the end2end results and structure recognition results.
|
||||
:param end2end_file: end2end results predict by end2end inference.
|
||||
:param structure_master_file: structure recognition results predict by structure master inference.
|
||||
"""
|
||||
self.end2end_file = end2end_file
|
||||
self.structure_master_file = structure_master_file
|
||||
self.end2end_results = pickle_load(end2end_file, prefix="end2end")
|
||||
self.structure_master_results = pickle_load(
|
||||
structure_master_file, prefix="structure"
|
||||
)
|
||||
|
||||
def match(self):
|
||||
"""
|
||||
Match process:
|
||||
pre-process : convert end2end and structure master results to xyxy, xywh ndnarray format.
|
||||
1. Use pseBbox is inside masterBbox judge rule
|
||||
2. Use iou between pseBbox and masterBbox rule
|
||||
3. Use min distance of center point rule
|
||||
:return:
|
||||
"""
|
||||
match_results = dict()
|
||||
for idx, (file_name, end2end_result) in enumerate(self.end2end_results.items()):
|
||||
match_list = []
|
||||
if file_name not in self.structure_master_results:
|
||||
continue
|
||||
structure_master_result = self.structure_master_results[file_name]
|
||||
(
|
||||
end2end_xyxy_bboxes,
|
||||
end2end_xywh_bboxes,
|
||||
structure_master_xywh_bboxes,
|
||||
structure_master_xyxy_bboxes,
|
||||
) = get_bboxes_list(end2end_result, structure_master_result)
|
||||
|
||||
# rule 1: center rule
|
||||
center_rule_match_list = center_rule_match(
|
||||
end2end_xywh_bboxes, structure_master_xyxy_bboxes
|
||||
)
|
||||
match_list.extend(center_rule_match_list)
|
||||
|
||||
# rule 2: iou rule
|
||||
# firstly, find not match index in previous step.
|
||||
center_no_match_end2end_indexs = find_no_match(
|
||||
match_list, len(end2end_xywh_bboxes), type="end2end"
|
||||
)
|
||||
if len(center_no_match_end2end_indexs) > 0:
|
||||
center_no_match_end2end_xyxy = end2end_xyxy_bboxes[
|
||||
center_no_match_end2end_indexs
|
||||
]
|
||||
# secondly, iou rule match
|
||||
iou_rule_match_list = iou_rule_match(
|
||||
center_no_match_end2end_xyxy,
|
||||
center_no_match_end2end_indexs,
|
||||
structure_master_xyxy_bboxes,
|
||||
)
|
||||
match_list.extend(iou_rule_match_list)
|
||||
|
||||
# rule 3: distance rule
|
||||
# match between no-match end2end bboxes and no-match master bboxes.
|
||||
# it will return master_bboxes_nums match-pairs.
|
||||
# firstly, find not match index in previous step.
|
||||
centerIou_no_match_end2end_indexs = find_no_match(
|
||||
match_list, len(end2end_xywh_bboxes), type="end2end"
|
||||
)
|
||||
centerIou_no_match_master_indexs = find_no_match(
|
||||
match_list, len(structure_master_xywh_bboxes), type="master"
|
||||
)
|
||||
if (
|
||||
len(centerIou_no_match_master_indexs) > 0
|
||||
and len(centerIou_no_match_end2end_indexs) > 0
|
||||
):
|
||||
centerIou_no_match_end2end_xywh = end2end_xywh_bboxes[
|
||||
centerIou_no_match_end2end_indexs
|
||||
]
|
||||
centerIou_no_match_master_xywh = structure_master_xywh_bboxes[
|
||||
centerIou_no_match_master_indexs
|
||||
]
|
||||
distance_match_list = distance_rule_match(
|
||||
centerIou_no_match_end2end_indexs,
|
||||
centerIou_no_match_end2end_xywh,
|
||||
centerIou_no_match_master_indexs,
|
||||
centerIou_no_match_master_xywh,
|
||||
)
|
||||
match_list.extend(distance_match_list)
|
||||
|
||||
# TODO:
|
||||
# The render no-match pseBbox, insert the last
|
||||
# After step3 distance rule, a master bbox at least match one end2end bbox.
|
||||
# But end2end bbox maybe overmuch, because numbers of master bbox will cut by max length.
|
||||
# For these render end2end bboxes, we will make some virtual master bboxes, and get matching.
|
||||
# The above extra insert bboxes will be further processed in "formatOutput" function.
|
||||
# After this operation, it will increase TEDS score.
|
||||
no_match_end2end_indexes = find_no_match(
|
||||
match_list, len(end2end_xywh_bboxes), type="end2end"
|
||||
)
|
||||
if len(no_match_end2end_indexes) > 0:
|
||||
no_match_end2end_xywh = end2end_xywh_bboxes[no_match_end2end_indexes]
|
||||
# sort the render no-match end2end bbox in row
|
||||
(
|
||||
end2end_sorted_indexes_list,
|
||||
end2end_sorted_bboxes_list,
|
||||
sorted_groups,
|
||||
sorted_bboxes_groups,
|
||||
) = sort_bbox(no_match_end2end_xywh, no_match_end2end_indexes)
|
||||
# make virtual master bboxes, and get matching with the no-match end2end bboxes.
|
||||
extra_match_list = extra_match(
|
||||
end2end_sorted_indexes_list, len(structure_master_xywh_bboxes)
|
||||
)
|
||||
match_list_add_extra_match = copy.deepcopy(match_list)
|
||||
match_list_add_extra_match.extend(extra_match_list)
|
||||
else:
|
||||
# no no-match end2end bboxes
|
||||
match_list_add_extra_match = copy.deepcopy(match_list)
|
||||
sorted_groups = []
|
||||
sorted_bboxes_groups = []
|
||||
|
||||
match_result_dict = {
|
||||
"match_list": match_list,
|
||||
"match_list_add_extra_match": match_list_add_extra_match,
|
||||
"sorted_groups": sorted_groups,
|
||||
"sorted_bboxes_groups": sorted_bboxes_groups,
|
||||
}
|
||||
|
||||
# format output
|
||||
match_result_dict = self._format(match_result_dict, file_name)
|
||||
|
||||
match_results[file_name] = match_result_dict
|
||||
|
||||
return match_results
|
||||
|
||||
def _format(self, match_result, file_name):
|
||||
"""
|
||||
Extend the master token(insert virtual master token), and format matching result.
|
||||
:param match_result:
|
||||
:param file_name:
|
||||
:return:
|
||||
"""
|
||||
end2end_info = self.end2end_results[file_name]
|
||||
master_info = self.structure_master_results[file_name]
|
||||
master_token = master_info["text"]
|
||||
sorted_groups = match_result["sorted_groups"]
|
||||
|
||||
# creat virtual master token
|
||||
virtual_master_token_list = []
|
||||
for line_group in sorted_groups:
|
||||
tmp_list = ["<tr>"]
|
||||
item_nums = len(line_group)
|
||||
for _ in range(item_nums):
|
||||
tmp_list.append("<td></td>")
|
||||
tmp_list.append("</tr>")
|
||||
virtual_master_token_list.extend(tmp_list)
|
||||
|
||||
# insert virtual master token
|
||||
master_token_list = master_token.split(",")
|
||||
if master_token_list[-1] == "</tbody>":
|
||||
# complete predict(no cut by max length)
|
||||
# This situation insert virtual master token will drop TEDs score in val set.
|
||||
# So we will not extend virtual token in this situation.
|
||||
|
||||
# fake extend virtual
|
||||
master_token_list[:-1].extend(virtual_master_token_list)
|
||||
|
||||
# real extend virtual
|
||||
# master_token_list = master_token_list[:-1]
|
||||
# master_token_list.extend(virtual_master_token_list)
|
||||
# master_token_list.append('</tbody>')
|
||||
|
||||
elif master_token_list[-1] == "<td></td>":
|
||||
master_token_list.append("</tr>")
|
||||
master_token_list.extend(virtual_master_token_list)
|
||||
master_token_list.append("</tbody>")
|
||||
else:
|
||||
master_token_list.extend(virtual_master_token_list)
|
||||
master_token_list.append("</tbody>")
|
||||
|
||||
# format output
|
||||
match_result.setdefault("matched_master_token_list", master_token_list)
|
||||
return match_result
|
||||
|
||||
def get_merge_result(self, match_results):
|
||||
"""
|
||||
Merge the OCR result into structure token to get final results.
|
||||
:param match_results:
|
||||
:return:
|
||||
"""
|
||||
merged_results = dict()
|
||||
|
||||
# break_token is linefeed token, when one master bbox has multiply end2end bboxes.
|
||||
break_token = " "
|
||||
|
||||
for idx, (file_name, match_info) in enumerate(match_results.items()):
|
||||
end2end_info = self.end2end_results[file_name]
|
||||
master_token_list = match_info["matched_master_token_list"]
|
||||
match_list = match_info["match_list_add_extra_match"]
|
||||
|
||||
match_dict = get_match_dict(match_list)
|
||||
match_text_dict = get_match_text_dict(match_dict, end2end_info, break_token)
|
||||
merged_result = insert_text_to_token(master_token_list, match_text_dict)
|
||||
merged_result = deal_bb(merged_result)
|
||||
|
||||
merged_results[file_name] = merged_result
|
||||
|
||||
return merged_results
|
||||
|
||||
|
||||
class TableMasterMatcher(Matcher):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, structure_res, dt_boxes, rec_res, img_name=1):
|
||||
end2end_results = {img_name: []}
|
||||
for dt_box, res in zip(dt_boxes, rec_res):
|
||||
d = dict(
|
||||
bbox=np.array(dt_box),
|
||||
text=res[0],
|
||||
)
|
||||
end2end_results[img_name].append(d)
|
||||
|
||||
self.end2end_results = end2end_results
|
||||
|
||||
structure_master_result_dict = {img_name: {}}
|
||||
pred_structures, pred_bboxes = structure_res
|
||||
pred_structures = ",".join(pred_structures[3:-3])
|
||||
structure_master_result_dict[img_name]["text"] = pred_structures
|
||||
structure_master_result_dict[img_name]["bbox"] = pred_bboxes
|
||||
self.structure_master_results = structure_master_result_dict
|
||||
|
||||
# match
|
||||
match_results = self.match()
|
||||
merged_results = self.get_merge_result(match_results)
|
||||
pred_html = merged_results[img_name]
|
||||
pred_html = "<html><body><table>" + pred_html + "</table></body></html>"
|
||||
return pred_html
|
||||
16
ppstructure/table/table_metric/__init__.py
Executable file
16
ppstructure/table/table_metric/__init__.py
Executable file
@@ -0,0 +1,16 @@
|
||||
# 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.
|
||||
|
||||
__all__ = ["TEDS"]
|
||||
from .table_metric import TEDS
|
||||
55
ppstructure/table/table_metric/parallel.py
Executable file
55
ppstructure/table/table_metric/parallel.py
Executable file
@@ -0,0 +1,55 @@
|
||||
from tqdm import tqdm
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
|
||||
def parallel_process(array, function, n_jobs=16, use_kwargs=False, front_num=0):
|
||||
"""
|
||||
A parallel version of the map function with a progress bar.
|
||||
Args:
|
||||
array (array-like): An array to iterate over.
|
||||
function (function): A python function to apply to the elements of array
|
||||
n_jobs (int, default=16): The number of cores to use
|
||||
use_kwargs (boolean, default=False): Whether to consider the elements of array as dictionaries of
|
||||
keyword arguments to function
|
||||
front_num (int, default=3): The number of iterations to run serially before kicking off the parallel job.
|
||||
Useful for catching bugs
|
||||
Returns:
|
||||
[function(array[0]), function(array[1]), ...]
|
||||
"""
|
||||
# We run the first few iterations serially to catch bugs
|
||||
if front_num > 0:
|
||||
front = [
|
||||
function(**a) if use_kwargs else function(a) for a in array[:front_num]
|
||||
]
|
||||
else:
|
||||
front = []
|
||||
# If we set n_jobs to 1, just run a list comprehension. This is useful for benchmarking and debugging.
|
||||
if n_jobs == 1:
|
||||
return front + [
|
||||
function(**a) if use_kwargs else function(a)
|
||||
for a in tqdm(array[front_num:])
|
||||
]
|
||||
# Assemble the workers
|
||||
with ProcessPoolExecutor(max_workers=n_jobs) as pool:
|
||||
# Pass the elements of array into function
|
||||
if use_kwargs:
|
||||
futures = [pool.submit(function, **a) for a in array[front_num:]]
|
||||
else:
|
||||
futures = [pool.submit(function, a) for a in array[front_num:]]
|
||||
kwargs = {
|
||||
"total": len(futures),
|
||||
"unit": "it",
|
||||
"unit_scale": True,
|
||||
"leave": True,
|
||||
}
|
||||
# Print out the progress as tasks complete
|
||||
for f in tqdm(as_completed(futures), **kwargs):
|
||||
pass
|
||||
out = []
|
||||
# Get the results from the futures.
|
||||
for i, future in tqdm(enumerate(futures)):
|
||||
try:
|
||||
out.append(future.result())
|
||||
except Exception as e:
|
||||
out.append(e)
|
||||
return front + out
|
||||
249
ppstructure/table/table_metric/table_metric.py
Executable file
249
ppstructure/table/table_metric/table_metric.py
Executable file
@@ -0,0 +1,249 @@
|
||||
# Copyright 2020 IBM
|
||||
# Author: peter.zhong@au1.ibm.com
|
||||
#
|
||||
# This is free software; you can redistribute it and/or modify
|
||||
# it under the terms of the Apache 2.0 License.
|
||||
#
|
||||
# This software is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# Apache 2.0 License for more details.
|
||||
|
||||
from rapidfuzz.distance import Levenshtein
|
||||
from apted import APTED, Config
|
||||
from apted.helpers import Tree
|
||||
from collections import deque
|
||||
from .parallel import parallel_process
|
||||
from tqdm import tqdm
|
||||
from paddle.utils import try_import
|
||||
|
||||
|
||||
class TableTree(Tree):
|
||||
def __init__(self, tag, colspan=None, rowspan=None, content=None, *children):
|
||||
self.tag = tag
|
||||
self.colspan = colspan
|
||||
self.rowspan = rowspan
|
||||
self.content = content
|
||||
self.children = list(children)
|
||||
|
||||
def bracket(self):
|
||||
"""Show tree using brackets notation"""
|
||||
if self.tag == "td":
|
||||
result = '"tag": %s, "colspan": %d, "rowspan": %d, "text": %s' % (
|
||||
self.tag,
|
||||
self.colspan,
|
||||
self.rowspan,
|
||||
self.content,
|
||||
)
|
||||
else:
|
||||
result = '"tag": %s' % self.tag
|
||||
for child in self.children:
|
||||
result += child.bracket()
|
||||
return "{{{}}}".format(result)
|
||||
|
||||
|
||||
class CustomConfig(Config):
|
||||
def rename(self, node1, node2):
|
||||
"""Compares attributes of trees"""
|
||||
# print(node1.tag)
|
||||
if (
|
||||
(node1.tag != node2.tag)
|
||||
or (node1.colspan != node2.colspan)
|
||||
or (node1.rowspan != node2.rowspan)
|
||||
):
|
||||
return 1.0
|
||||
if node1.tag == "td":
|
||||
if node1.content or node2.content:
|
||||
# print(node1.content, )
|
||||
return Levenshtein.normalized_distance(node1.content, node2.content)
|
||||
return 0.0
|
||||
|
||||
|
||||
class CustomConfig_del_short(Config):
|
||||
def rename(self, node1, node2):
|
||||
"""Compares attributes of trees"""
|
||||
if (
|
||||
(node1.tag != node2.tag)
|
||||
or (node1.colspan != node2.colspan)
|
||||
or (node1.rowspan != node2.rowspan)
|
||||
):
|
||||
return 1.0
|
||||
if node1.tag == "td":
|
||||
if node1.content or node2.content:
|
||||
# print('before')
|
||||
# print(node1.content, node2.content)
|
||||
# print('after')
|
||||
node1_content = node1.content
|
||||
node2_content = node2.content
|
||||
if len(node1_content) < 3:
|
||||
node1_content = ["####"]
|
||||
if len(node2_content) < 3:
|
||||
node2_content = ["####"]
|
||||
return Levenshtein.normalized_distance(node1_content, node2_content)
|
||||
return 0.0
|
||||
|
||||
|
||||
class CustomConfig_del_block(Config):
|
||||
def rename(self, node1, node2):
|
||||
"""Compares attributes of trees"""
|
||||
if (
|
||||
(node1.tag != node2.tag)
|
||||
or (node1.colspan != node2.colspan)
|
||||
or (node1.rowspan != node2.rowspan)
|
||||
):
|
||||
return 1.0
|
||||
if node1.tag == "td":
|
||||
if node1.content or node2.content:
|
||||
node1_content = node1.content
|
||||
node2_content = node2.content
|
||||
while " " in node1_content:
|
||||
print(node1_content.index(" "))
|
||||
node1_content.pop(node1_content.index(" "))
|
||||
while " " in node2_content:
|
||||
print(node2_content.index(" "))
|
||||
node2_content.pop(node2_content.index(" "))
|
||||
return Levenshtein.normalized_distance(node1_content, node2_content)
|
||||
return 0.0
|
||||
|
||||
|
||||
class TEDS(object):
|
||||
"""Tree Edit Distance basead Similarity"""
|
||||
|
||||
def __init__(self, structure_only=False, n_jobs=1, ignore_nodes=None):
|
||||
assert isinstance(n_jobs, int) and (
|
||||
n_jobs >= 1
|
||||
), "n_jobs must be an integer greater than 1"
|
||||
self.structure_only = structure_only
|
||||
self.n_jobs = n_jobs
|
||||
self.ignore_nodes = ignore_nodes
|
||||
self.__tokens__ = []
|
||||
|
||||
def tokenize(self, node):
|
||||
"""Tokenizes table cells"""
|
||||
self.__tokens__.append("<%s>" % node.tag)
|
||||
if node.text is not None:
|
||||
self.__tokens__ += list(node.text)
|
||||
for n in node.getchildren():
|
||||
self.tokenize(n)
|
||||
if node.tag != "unk":
|
||||
self.__tokens__.append("</%s>" % node.tag)
|
||||
if node.tag != "td" and node.tail is not None:
|
||||
self.__tokens__ += list(node.tail)
|
||||
|
||||
def load_html_tree(self, node, parent=None):
|
||||
"""Converts HTML tree to the format required by apted"""
|
||||
global __tokens__
|
||||
if node.tag == "td":
|
||||
if self.structure_only:
|
||||
cell = []
|
||||
else:
|
||||
self.__tokens__ = []
|
||||
self.tokenize(node)
|
||||
cell = self.__tokens__[1:-1].copy()
|
||||
new_node = TableTree(
|
||||
node.tag,
|
||||
int(node.attrib.get("colspan", "1")),
|
||||
int(node.attrib.get("rowspan", "1")),
|
||||
cell,
|
||||
*deque(),
|
||||
)
|
||||
else:
|
||||
new_node = TableTree(node.tag, None, None, None, *deque())
|
||||
if parent is not None:
|
||||
parent.children.append(new_node)
|
||||
if node.tag != "td":
|
||||
for n in node.getchildren():
|
||||
self.load_html_tree(n, new_node)
|
||||
if parent is None:
|
||||
return new_node
|
||||
|
||||
def evaluate(self, pred, true):
|
||||
"""Computes TEDS score between the prediction and the ground truth of a
|
||||
given sample
|
||||
"""
|
||||
try_import("lxml")
|
||||
from lxml import etree, html
|
||||
|
||||
if (not pred) or (not true):
|
||||
return 0.0
|
||||
parser = html.HTMLParser(remove_comments=True, encoding="utf-8")
|
||||
pred = html.fromstring(pred, parser=parser)
|
||||
true = html.fromstring(true, parser=parser)
|
||||
if pred.xpath("body/table") and true.xpath("body/table"):
|
||||
pred = pred.xpath("body/table")[0]
|
||||
true = true.xpath("body/table")[0]
|
||||
if self.ignore_nodes:
|
||||
etree.strip_tags(pred, *self.ignore_nodes)
|
||||
etree.strip_tags(true, *self.ignore_nodes)
|
||||
n_nodes_pred = len(pred.xpath(".//*"))
|
||||
n_nodes_true = len(true.xpath(".//*"))
|
||||
n_nodes = max(n_nodes_pred, n_nodes_true)
|
||||
tree_pred = self.load_html_tree(pred)
|
||||
tree_true = self.load_html_tree(true)
|
||||
distance = APTED(
|
||||
tree_pred, tree_true, CustomConfig()
|
||||
).compute_edit_distance()
|
||||
return 1.0 - (float(distance) / n_nodes)
|
||||
else:
|
||||
return 0.0
|
||||
|
||||
def batch_evaluate(self, pred_json, true_json):
|
||||
"""Computes TEDS score between the prediction and the ground truth of
|
||||
a batch of samples
|
||||
@params pred_json: {'FILENAME': 'HTML CODE', ...}
|
||||
@params true_json: {'FILENAME': {'html': 'HTML CODE'}, ...}
|
||||
@output: {'FILENAME': 'TEDS SCORE', ...}
|
||||
"""
|
||||
samples = true_json.keys()
|
||||
if self.n_jobs == 1:
|
||||
scores = [
|
||||
self.evaluate(pred_json.get(filename, ""), true_json[filename]["html"])
|
||||
for filename in tqdm(samples)
|
||||
]
|
||||
else:
|
||||
inputs = [
|
||||
{
|
||||
"pred": pred_json.get(filename, ""),
|
||||
"true": true_json[filename]["html"],
|
||||
}
|
||||
for filename in samples
|
||||
]
|
||||
scores = parallel_process(
|
||||
inputs, self.evaluate, use_kwargs=True, n_jobs=self.n_jobs, front_num=1
|
||||
)
|
||||
scores = dict(zip(samples, scores))
|
||||
return scores
|
||||
|
||||
def batch_evaluate_html(self, pred_htmls, true_htmls):
|
||||
"""Computes TEDS score between the prediction and the ground truth of
|
||||
a batch of samples
|
||||
"""
|
||||
if self.n_jobs == 1:
|
||||
scores = [
|
||||
self.evaluate(pred_html, true_html)
|
||||
for (pred_html, true_html) in zip(pred_htmls, true_htmls)
|
||||
]
|
||||
else:
|
||||
inputs = [
|
||||
{"pred": pred_html, "true": true_html}
|
||||
for (pred_html, true_html) in zip(pred_htmls, true_htmls)
|
||||
]
|
||||
|
||||
scores = parallel_process(
|
||||
inputs, self.evaluate, use_kwargs=True, n_jobs=self.n_jobs, front_num=1
|
||||
)
|
||||
return scores
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import json
|
||||
import pprint
|
||||
|
||||
with open("sample_pred.json") as fp:
|
||||
pred_json = json.load(fp)
|
||||
with open("sample_gt.json") as fp:
|
||||
true_json = json.load(fp)
|
||||
teds = TEDS(n_jobs=4)
|
||||
scores = teds.batch_evaluate(pred_json, true_json)
|
||||
pp = pprint.PrettyPrinter()
|
||||
pp.pprint(scores)
|
||||
13
ppstructure/table/tablepyxl/__init__.py
Normal file
13
ppstructure/table/tablepyxl/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# 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.
|
||||
349
ppstructure/table/tablepyxl/style.py
Normal file
349
ppstructure/table/tablepyxl/style.py
Normal file
@@ -0,0 +1,349 @@
|
||||
# This is where we handle translating css styles into openpyxl styles
|
||||
# and cascading those from parent to child in the dom.
|
||||
|
||||
try:
|
||||
from openpyxl.cell import cell
|
||||
from openpyxl.styles import (
|
||||
Font,
|
||||
Alignment,
|
||||
PatternFill,
|
||||
NamedStyle,
|
||||
Border,
|
||||
Side,
|
||||
Color,
|
||||
)
|
||||
from openpyxl.styles.fills import FILL_SOLID
|
||||
from openpyxl.styles.numbers import FORMAT_CURRENCY_USD_SIMPLE, FORMAT_PERCENTAGE
|
||||
from openpyxl.styles.colors import BLACK
|
||||
except:
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"Can not import openpyxl, some functions in the ppstructure may not work. Please manually install openpyxl before using ppstructure."
|
||||
)
|
||||
|
||||
FORMAT_DATE_MMDDYYYY = "mm/dd/yyyy"
|
||||
|
||||
|
||||
def colormap(color):
|
||||
"""
|
||||
Convenience for looking up known colors
|
||||
"""
|
||||
cmap = {"black": BLACK}
|
||||
return cmap.get(color, color)
|
||||
|
||||
|
||||
def style_string_to_dict(style):
|
||||
"""
|
||||
Convert css style string to a python dictionary
|
||||
"""
|
||||
|
||||
def clean_split(string, delim):
|
||||
return (s.strip() for s in string.split(delim))
|
||||
|
||||
styles = [clean_split(s, ":") for s in style.split(";") if ":" in s]
|
||||
return dict(styles)
|
||||
|
||||
|
||||
def get_side(style, name):
|
||||
return {
|
||||
"border_style": style.get("border-{}-style".format(name)),
|
||||
"color": colormap(style.get("border-{}-color".format(name))),
|
||||
}
|
||||
|
||||
|
||||
known_styles = {}
|
||||
|
||||
|
||||
def style_dict_to_named_style(style_dict, number_format=None):
|
||||
"""
|
||||
Change css style (stored in a python dictionary) to openpyxl NamedStyle
|
||||
"""
|
||||
|
||||
style_and_format_string = str(
|
||||
{
|
||||
"style_dict": style_dict,
|
||||
"parent": style_dict.parent,
|
||||
"number_format": number_format,
|
||||
}
|
||||
)
|
||||
|
||||
if style_and_format_string not in known_styles:
|
||||
# Font
|
||||
font = Font(
|
||||
bold=style_dict.get("font-weight") == "bold",
|
||||
color=style_dict.get_color("color", None),
|
||||
size=style_dict.get("font-size"),
|
||||
)
|
||||
|
||||
# Alignment
|
||||
alignment = Alignment(
|
||||
horizontal=style_dict.get("text-align", "general"),
|
||||
vertical=style_dict.get("vertical-align"),
|
||||
wrap_text=style_dict.get("white-space", "nowrap") == "normal",
|
||||
)
|
||||
|
||||
# Fill
|
||||
bg_color = style_dict.get_color("background-color")
|
||||
fg_color = style_dict.get_color("foreground-color", Color())
|
||||
fill_type = style_dict.get("fill-type")
|
||||
if bg_color and bg_color != "transparent":
|
||||
fill = PatternFill(
|
||||
fill_type=fill_type or FILL_SOLID,
|
||||
start_color=bg_color,
|
||||
end_color=fg_color,
|
||||
)
|
||||
else:
|
||||
fill = PatternFill()
|
||||
|
||||
# Border
|
||||
border = Border(
|
||||
left=Side(**get_side(style_dict, "left")),
|
||||
right=Side(**get_side(style_dict, "right")),
|
||||
top=Side(**get_side(style_dict, "top")),
|
||||
bottom=Side(**get_side(style_dict, "bottom")),
|
||||
diagonal=Side(**get_side(style_dict, "diagonal")),
|
||||
diagonal_direction=None,
|
||||
outline=Side(**get_side(style_dict, "outline")),
|
||||
vertical=None,
|
||||
horizontal=None,
|
||||
)
|
||||
|
||||
name = "Style {}".format(len(known_styles) + 1)
|
||||
|
||||
pyxl_style = NamedStyle(
|
||||
name=name,
|
||||
font=font,
|
||||
fill=fill,
|
||||
alignment=alignment,
|
||||
border=border,
|
||||
number_format=number_format,
|
||||
)
|
||||
|
||||
known_styles[style_and_format_string] = pyxl_style
|
||||
|
||||
return known_styles[style_and_format_string]
|
||||
|
||||
|
||||
class StyleDict(dict):
|
||||
"""
|
||||
It's like a dictionary, but it looks for items in the parent dictionary
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.parent = kwargs.pop("parent", None)
|
||||
super(StyleDict, self).__init__(*args, **kwargs)
|
||||
|
||||
def __getitem__(self, item):
|
||||
if item in self:
|
||||
return super(StyleDict, self).__getitem__(item)
|
||||
elif self.parent:
|
||||
return self.parent[item]
|
||||
else:
|
||||
raise KeyError("{} not found".format(item))
|
||||
|
||||
def __hash__(self):
|
||||
return hash(tuple([(k, self.get(k)) for k in self._keys()]))
|
||||
|
||||
# Yielding the keys avoids creating unnecessary data structures
|
||||
# and happily works with both python2 and python3 where the
|
||||
# .keys() method is a dictionary_view in python3 and a list in python2.
|
||||
def _keys(self):
|
||||
yielded = set()
|
||||
for k in self.keys():
|
||||
yielded.add(k)
|
||||
yield k
|
||||
if self.parent:
|
||||
for k in self.parent._keys():
|
||||
if k not in yielded:
|
||||
yielded.add(k)
|
||||
yield k
|
||||
|
||||
def get(self, k, d=None):
|
||||
try:
|
||||
return self[k]
|
||||
except KeyError:
|
||||
return d
|
||||
|
||||
def get_color(self, k, d=None):
|
||||
"""
|
||||
Strip leading # off colors if necessary
|
||||
"""
|
||||
color = self.get(k, d)
|
||||
if hasattr(color, "startswith") and color.startswith("#"):
|
||||
color = color[1:]
|
||||
if (
|
||||
len(color) == 3
|
||||
): # Premailers reduces colors like #00ff00 to #0f0, openpyxl doesn't like that
|
||||
color = "".join(2 * c for c in color)
|
||||
return color
|
||||
|
||||
|
||||
class Element(object):
|
||||
"""
|
||||
Our base class for representing an html element along with a cascading style.
|
||||
The element is created along with a parent so that the StyleDict that we store
|
||||
can point to the parent's StyleDict.
|
||||
"""
|
||||
|
||||
def __init__(self, element, parent=None):
|
||||
self.element = element
|
||||
self.number_format = None
|
||||
parent_style = parent.style_dict if parent else None
|
||||
self.style_dict = StyleDict(
|
||||
style_string_to_dict(element.get("style", "")), parent=parent_style
|
||||
)
|
||||
self._style_cache = None
|
||||
|
||||
def style(self):
|
||||
"""
|
||||
Turn the css styles for this element into an openpyxl NamedStyle.
|
||||
"""
|
||||
if not self._style_cache:
|
||||
self._style_cache = style_dict_to_named_style(
|
||||
self.style_dict, number_format=self.number_format
|
||||
)
|
||||
return self._style_cache
|
||||
|
||||
def get_dimension(self, dimension_key):
|
||||
"""
|
||||
Extracts the dimension from the style dict of the Element and returns it as a float.
|
||||
"""
|
||||
dimension = self.style_dict.get(dimension_key)
|
||||
if dimension:
|
||||
if dimension[-2:] in ["px", "em", "pt", "in", "cm"]:
|
||||
dimension = dimension[:-2]
|
||||
dimension = float(dimension)
|
||||
return dimension
|
||||
|
||||
|
||||
class Table(Element):
|
||||
"""
|
||||
The concrete implementations of Elements are semantically named for the types of elements we are interested in.
|
||||
This defines a very concrete tree structure for html tables that we expect to deal with. I prefer this compared to
|
||||
allowing Element to have an arbitrary number of children and dealing with an abstract element tree.
|
||||
"""
|
||||
|
||||
def __init__(self, table):
|
||||
"""
|
||||
takes an html table object (from lxml)
|
||||
"""
|
||||
super(Table, self).__init__(table)
|
||||
table_head = table.find("thead")
|
||||
self.head = (
|
||||
TableHead(table_head, parent=self) if table_head is not None else None
|
||||
)
|
||||
table_body = table.find("tbody")
|
||||
self.body = TableBody(
|
||||
table_body if table_body is not None else table, parent=self
|
||||
)
|
||||
|
||||
|
||||
class TableHead(Element):
|
||||
"""
|
||||
This class maps to the `<th>` element of the html table.
|
||||
"""
|
||||
|
||||
def __init__(self, head, parent=None):
|
||||
super(TableHead, self).__init__(head, parent=parent)
|
||||
self.rows = [TableRow(tr, parent=self) for tr in head.findall("tr")]
|
||||
|
||||
|
||||
class TableBody(Element):
|
||||
"""
|
||||
This class maps to the `<tbody>` element of the html table.
|
||||
"""
|
||||
|
||||
def __init__(self, body, parent=None):
|
||||
super(TableBody, self).__init__(body, parent=parent)
|
||||
self.rows = [TableRow(tr, parent=self) for tr in body.findall("tr")]
|
||||
|
||||
|
||||
class TableRow(Element):
|
||||
"""
|
||||
This class maps to the `<tr>` element of the html table.
|
||||
"""
|
||||
|
||||
def __init__(self, tr, parent=None):
|
||||
super(TableRow, self).__init__(tr, parent=parent)
|
||||
self.cells = [
|
||||
TableCell(cell, parent=self) for cell in tr.findall("th") + tr.findall("td")
|
||||
]
|
||||
|
||||
|
||||
def element_to_string(el):
|
||||
return _element_to_string(el).strip()
|
||||
|
||||
|
||||
def _element_to_string(el):
|
||||
string = ""
|
||||
|
||||
for x in el.iterchildren():
|
||||
string += "\n" + _element_to_string(x)
|
||||
|
||||
text = el.text.strip() if el.text else ""
|
||||
tail = el.tail.strip() if el.tail else ""
|
||||
|
||||
return text + string + "\n" + tail
|
||||
|
||||
|
||||
class TableCell(Element):
|
||||
"""
|
||||
This class maps to the `<td>` element of the html table.
|
||||
"""
|
||||
|
||||
CELL_TYPES = {
|
||||
"TYPE_STRING",
|
||||
"TYPE_FORMULA",
|
||||
"TYPE_NUMERIC",
|
||||
"TYPE_BOOL",
|
||||
"TYPE_CURRENCY",
|
||||
"TYPE_PERCENTAGE",
|
||||
"TYPE_NULL",
|
||||
"TYPE_INLINE",
|
||||
"TYPE_ERROR",
|
||||
"TYPE_FORMULA_CACHE_STRING",
|
||||
"TYPE_INTEGER",
|
||||
}
|
||||
|
||||
def __init__(self, cell, parent=None):
|
||||
super(TableCell, self).__init__(cell, parent=parent)
|
||||
self.value = element_to_string(cell)
|
||||
self.number_format = self.get_number_format()
|
||||
|
||||
def data_type(self):
|
||||
cell_types = self.CELL_TYPES & set(self.element.get("class", "").split())
|
||||
if cell_types:
|
||||
if "TYPE_FORMULA" in cell_types:
|
||||
# Make sure TYPE_FORMULA takes precedence over the other classes in the set.
|
||||
cell_type = "TYPE_FORMULA"
|
||||
elif cell_types & {"TYPE_CURRENCY", "TYPE_INTEGER", "TYPE_PERCENTAGE"}:
|
||||
cell_type = "TYPE_NUMERIC"
|
||||
else:
|
||||
cell_type = cell_types.pop()
|
||||
else:
|
||||
cell_type = "TYPE_STRING"
|
||||
return getattr(cell, cell_type)
|
||||
|
||||
def get_number_format(self):
|
||||
if "TYPE_CURRENCY" in self.element.get("class", "").split():
|
||||
return FORMAT_CURRENCY_USD_SIMPLE
|
||||
if "TYPE_INTEGER" in self.element.get("class", "").split():
|
||||
return "#,##0"
|
||||
if "TYPE_PERCENTAGE" in self.element.get("class", "").split():
|
||||
return FORMAT_PERCENTAGE
|
||||
if "TYPE_DATE" in self.element.get("class", "").split():
|
||||
return FORMAT_DATE_MMDDYYYY
|
||||
if self.data_type() == cell.TYPE_NUMERIC:
|
||||
try:
|
||||
int(self.value)
|
||||
except ValueError:
|
||||
return "#,##0.##"
|
||||
else:
|
||||
return "#,##0"
|
||||
|
||||
def format(self, cell):
|
||||
cell.style = self.style()
|
||||
data_type = self.data_type()
|
||||
if data_type:
|
||||
cell.data_type = data_type
|
||||
135
ppstructure/table/tablepyxl/tablepyxl.py
Normal file
135
ppstructure/table/tablepyxl/tablepyxl.py
Normal file
@@ -0,0 +1,135 @@
|
||||
# Do imports like python3 so our package works for 2 and 3
|
||||
from __future__ import absolute_import
|
||||
|
||||
|
||||
from tablepyxl.style import Table
|
||||
from paddle.utils import try_import
|
||||
|
||||
|
||||
def string_to_int(s):
|
||||
if s.isdigit():
|
||||
return int(s)
|
||||
return 0
|
||||
|
||||
|
||||
def get_Tables(doc):
|
||||
try_import("lxml")
|
||||
from lxml import etree, html
|
||||
|
||||
tree = html.fromstring(doc)
|
||||
comments = tree.xpath("//comment()")
|
||||
for comment in comments:
|
||||
comment.drop_tag()
|
||||
return [Table(table) for table in tree.xpath("//table")]
|
||||
|
||||
|
||||
def write_rows(worksheet, elem, row, column=1):
|
||||
"""
|
||||
Writes every tr child element of elem to a row in the worksheet
|
||||
returns the next row after all rows are written
|
||||
"""
|
||||
try_import("openpyxl")
|
||||
from openpyxl.cell.cell import MergedCell
|
||||
from openpyxl.utils import get_column_letter
|
||||
|
||||
initial_column = column
|
||||
for table_row in elem.rows:
|
||||
for table_cell in table_row.cells:
|
||||
cell = worksheet.cell(row=row, column=column)
|
||||
while isinstance(cell, MergedCell):
|
||||
column += 1
|
||||
cell = worksheet.cell(row=row, column=column)
|
||||
|
||||
colspan = string_to_int(table_cell.element.get("colspan", "1"))
|
||||
rowspan = string_to_int(table_cell.element.get("rowspan", "1"))
|
||||
if rowspan > 1 or colspan > 1:
|
||||
worksheet.merge_cells(
|
||||
start_row=row,
|
||||
start_column=column,
|
||||
end_row=row + rowspan - 1,
|
||||
end_column=column + colspan - 1,
|
||||
)
|
||||
|
||||
cell.value = table_cell.value
|
||||
table_cell.format(cell)
|
||||
min_width = table_cell.get_dimension("min-width")
|
||||
max_width = table_cell.get_dimension("max-width")
|
||||
|
||||
if colspan == 1:
|
||||
# Initially, when iterating for the first time through the loop, the width of all the cells is None.
|
||||
# As we start filling in contents, the initial width of the cell (which can be retrieved by:
|
||||
# worksheet.column_dimensions[get_column_letter(column)].width) is equal to the width of the previous
|
||||
# cell in the same column (i.e. width of A2 = width of A1)
|
||||
width = max(
|
||||
worksheet.column_dimensions[get_column_letter(column)].width or 0,
|
||||
len(table_cell.value) + 2,
|
||||
)
|
||||
if max_width and width > max_width:
|
||||
width = max_width
|
||||
elif min_width and width < min_width:
|
||||
width = min_width
|
||||
worksheet.column_dimensions[get_column_letter(column)].width = width
|
||||
column += colspan
|
||||
row += 1
|
||||
column = initial_column
|
||||
return row
|
||||
|
||||
|
||||
def table_to_sheet(table, wb):
|
||||
"""
|
||||
Takes a table and workbook and writes the table to a new sheet.
|
||||
The sheet title will be the same as the table attribute name.
|
||||
"""
|
||||
ws = wb.create_sheet(title=table.element.get("name"))
|
||||
insert_table(table, ws, 1, 1)
|
||||
|
||||
|
||||
def document_to_workbook(doc, wb=None, base_url=None):
|
||||
"""
|
||||
Takes a string representation of an html document and writes one sheet for
|
||||
every table in the document.
|
||||
The workbook is returned
|
||||
"""
|
||||
try_import("premailer")
|
||||
try_import("openpyxl")
|
||||
from premailer import Premailer
|
||||
from openpyxl import Workbook
|
||||
|
||||
if not wb:
|
||||
wb = Workbook()
|
||||
wb.remove(wb.active)
|
||||
|
||||
inline_styles_doc = Premailer(
|
||||
doc, base_url=base_url, remove_classes=False
|
||||
).transform()
|
||||
tables = get_Tables(inline_styles_doc)
|
||||
|
||||
for table in tables:
|
||||
table_to_sheet(table, wb)
|
||||
|
||||
return wb
|
||||
|
||||
|
||||
def document_to_xl(doc, filename, base_url=None):
|
||||
"""
|
||||
Takes a string representation of an html document and writes one sheet for
|
||||
every table in the document. The workbook is written out to a file called filename
|
||||
"""
|
||||
wb = document_to_workbook(doc, base_url=base_url)
|
||||
wb.save(filename)
|
||||
|
||||
|
||||
def insert_table(table, worksheet, column, row):
|
||||
if table.head:
|
||||
row = write_rows(worksheet, table.head, row, column)
|
||||
if table.body:
|
||||
row = write_rows(worksheet, table.body, row, column)
|
||||
|
||||
|
||||
def insert_table_at_cell(table, cell):
|
||||
"""
|
||||
Inserts a table at the location of an openpyxl Cell object.
|
||||
"""
|
||||
ws = cell.parent
|
||||
column, row = cell.column, cell.row
|
||||
insert_table(table, ws, column, row)
|
||||
298
ppstructure/utility.py
Normal file
298
ppstructure/utility.py
Normal file
@@ -0,0 +1,298 @@
|
||||
# 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.
|
||||
import random
|
||||
import ast
|
||||
import PIL
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
import numpy as np
|
||||
from tools.infer.utility import (
|
||||
draw_ocr_box_txt,
|
||||
str2bool,
|
||||
str2int_tuple,
|
||||
init_args as infer_args,
|
||||
)
|
||||
import math
|
||||
|
||||
|
||||
def init_args():
|
||||
parser = infer_args()
|
||||
|
||||
# params for output
|
||||
parser.add_argument("--output", type=str, default="./output")
|
||||
# params for table structure
|
||||
parser.add_argument("--table_max_len", type=int, default=488)
|
||||
parser.add_argument("--table_algorithm", type=str, default="TableAttn")
|
||||
parser.add_argument("--table_model_dir", type=str)
|
||||
parser.add_argument("--merge_no_span_structure", type=str2bool, default=True)
|
||||
parser.add_argument(
|
||||
"--table_char_dict_path",
|
||||
type=str,
|
||||
default="../ppocr/utils/dict/table_structure_dict_ch.txt",
|
||||
)
|
||||
# params for formula recognition
|
||||
parser.add_argument("--formula_algorithm", type=str, default="LaTeXOCR")
|
||||
parser.add_argument("--formula_model_dir", type=str)
|
||||
parser.add_argument(
|
||||
"--formula_char_dict_path",
|
||||
type=str,
|
||||
default="../ppocr/utils/dict/latex_ocr_tokenizer.json",
|
||||
)
|
||||
parser.add_argument("--formula_batch_num", type=int, default=1)
|
||||
# params for layout
|
||||
parser.add_argument("--layout_model_dir", type=str)
|
||||
parser.add_argument(
|
||||
"--layout_dict_path",
|
||||
type=str,
|
||||
default="../ppocr/utils/dict/layout_dict/layout_publaynet_dict.txt",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layout_score_threshold", type=float, default=0.5, help="Threshold of score."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layout_nms_threshold", type=float, default=0.5, help="Threshold of nms."
|
||||
)
|
||||
# params for kie
|
||||
parser.add_argument("--kie_algorithm", type=str, default="LayoutXLM")
|
||||
parser.add_argument("--ser_model_dir", type=str)
|
||||
parser.add_argument("--re_model_dir", type=str)
|
||||
parser.add_argument("--use_visual_backbone", type=str2bool, default=True)
|
||||
parser.add_argument(
|
||||
"--ser_dict_path", type=str, default="../train_data/XFUND/class_list_xfun.txt"
|
||||
)
|
||||
# need to be None or tb-yx
|
||||
parser.add_argument("--ocr_order_method", type=str, default=None)
|
||||
# params for inference
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
choices=["structure", "kie"],
|
||||
default="structure",
|
||||
help="structure and kie is supported",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image_orientation",
|
||||
type=bool,
|
||||
default=False,
|
||||
help="Whether to enable image orientation recognition",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layout",
|
||||
type=str2bool,
|
||||
default=True,
|
||||
help="Whether to enable layout analysis",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--table",
|
||||
type=str2bool,
|
||||
default=True,
|
||||
help="In the forward, whether the table area uses table recognition",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--formula",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to enable formula recognition",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ocr",
|
||||
type=str2bool,
|
||||
default=True,
|
||||
help="In the forward, whether the non-table area is recognition by ocr",
|
||||
)
|
||||
# param for recovery
|
||||
parser.add_argument(
|
||||
"--recovery",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to enable layout of recovery",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--recovery_to_markdown",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to enable layout of recovery to markdown",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_pdf2docx_api",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to use pdf2docx api",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--invert",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to invert image before processing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--binarize",
|
||||
type=str2bool,
|
||||
default=False,
|
||||
help="Whether to threshold binarize image before processing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--alphacolor",
|
||||
type=str2int_tuple,
|
||||
default=(255, 255, 255),
|
||||
help="Replacement color for the alpha channel, if the latter is present; R,G,B integers",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = init_args()
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def draw_structure_result(image, result, font_path):
|
||||
if isinstance(image, np.ndarray):
|
||||
image = Image.fromarray(image)
|
||||
boxes, txts, scores = [], [], []
|
||||
|
||||
img_layout = image.copy()
|
||||
draw_layout = ImageDraw.Draw(img_layout)
|
||||
text_color = (255, 255, 255)
|
||||
text_background_color = (80, 127, 255)
|
||||
catid2color = {}
|
||||
font_size = 15
|
||||
font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
|
||||
|
||||
for region in result:
|
||||
if region["type"] not in catid2color:
|
||||
box_color = (
|
||||
random.randint(0, 255),
|
||||
random.randint(0, 255),
|
||||
random.randint(0, 255),
|
||||
)
|
||||
catid2color[region["type"]] = box_color
|
||||
else:
|
||||
box_color = catid2color[region["type"]]
|
||||
box_layout = region["bbox"]
|
||||
draw_layout.rectangle(
|
||||
[(box_layout[0], box_layout[1]), (box_layout[2], box_layout[3])],
|
||||
outline=box_color,
|
||||
width=3,
|
||||
)
|
||||
|
||||
if int(PIL.__version__.split(".")[0]) < 10:
|
||||
text_w, text_h = font.getsize(region["type"])
|
||||
else:
|
||||
left, top, right, bottom = font.getbbox(region["type"])
|
||||
text_w, text_h = right - left, bottom - top
|
||||
|
||||
draw_layout.rectangle(
|
||||
[
|
||||
(box_layout[0], box_layout[1]),
|
||||
(box_layout[0] + text_w, box_layout[1] + text_h),
|
||||
],
|
||||
fill=text_background_color,
|
||||
)
|
||||
draw_layout.text(
|
||||
(box_layout[0], box_layout[1]), region["type"], fill=text_color, font=font
|
||||
)
|
||||
|
||||
if region["type"] == "table" or (
|
||||
region["type"] == "equation" and "latex" in region["res"]
|
||||
):
|
||||
pass
|
||||
else:
|
||||
for text_result in region["res"]:
|
||||
boxes.append(np.array(text_result["text_region"]))
|
||||
txts.append(text_result["text"])
|
||||
scores.append(text_result["confidence"])
|
||||
|
||||
if "text_word_region" in text_result:
|
||||
for word_region in text_result["text_word_region"]:
|
||||
char_box = word_region
|
||||
box_height = int(
|
||||
math.sqrt(
|
||||
(char_box[0][0] - char_box[3][0]) ** 2
|
||||
+ (char_box[0][1] - char_box[3][1]) ** 2
|
||||
)
|
||||
)
|
||||
box_width = int(
|
||||
math.sqrt(
|
||||
(char_box[0][0] - char_box[1][0]) ** 2
|
||||
+ (char_box[0][1] - char_box[1][1]) ** 2
|
||||
)
|
||||
)
|
||||
if box_height == 0 or box_width == 0:
|
||||
continue
|
||||
boxes.append(word_region)
|
||||
txts.append("")
|
||||
scores.append(1.0)
|
||||
|
||||
im_show = draw_ocr_box_txt(
|
||||
img_layout, boxes, txts, scores, font_path=font_path, drop_score=0
|
||||
)
|
||||
return im_show
|
||||
|
||||
|
||||
def cal_ocr_word_box(rec_str, box, rec_word_info):
|
||||
"""Calculate the detection frame for each word based on the results of recognition and detection of ocr"""
|
||||
|
||||
col_num, word_list, word_col_list, state_list = rec_word_info
|
||||
box = box.tolist()
|
||||
bbox_x_start = box[0][0]
|
||||
bbox_x_end = box[1][0]
|
||||
bbox_y_start = box[0][1]
|
||||
bbox_y_end = box[2][1]
|
||||
|
||||
cell_width = (bbox_x_end - bbox_x_start) / col_num
|
||||
|
||||
word_box_list = []
|
||||
word_box_content_list = []
|
||||
cn_width_list = []
|
||||
cn_col_list = []
|
||||
for word, word_col, state in zip(word_list, word_col_list, state_list):
|
||||
if state == "cn":
|
||||
if len(word_col) != 1:
|
||||
char_seq_length = (word_col[-1] - word_col[0] + 1) * cell_width
|
||||
char_width = char_seq_length / (len(word_col) - 1)
|
||||
cn_width_list.append(char_width)
|
||||
cn_col_list += word_col
|
||||
word_box_content_list += word
|
||||
else:
|
||||
cell_x_start = bbox_x_start + int(word_col[0] * cell_width)
|
||||
cell_x_end = bbox_x_start + int((word_col[-1] + 1) * cell_width)
|
||||
cell = (
|
||||
(cell_x_start, bbox_y_start),
|
||||
(cell_x_end, bbox_y_start),
|
||||
(cell_x_end, bbox_y_end),
|
||||
(cell_x_start, bbox_y_end),
|
||||
)
|
||||
word_box_list.append(cell)
|
||||
word_box_content_list.append("".join(word))
|
||||
if len(cn_col_list) != 0:
|
||||
if len(cn_width_list) != 0:
|
||||
avg_char_width = np.mean(cn_width_list)
|
||||
else:
|
||||
avg_char_width = (bbox_x_end - bbox_x_start) / len(rec_str)
|
||||
for center_idx in cn_col_list:
|
||||
center_x = (center_idx + 0.5) * cell_width
|
||||
cell_x_start = max(int(center_x - avg_char_width / 2), 0) + bbox_x_start
|
||||
cell_x_end = (
|
||||
min(int(center_x + avg_char_width / 2), bbox_x_end - bbox_x_start)
|
||||
+ bbox_x_start
|
||||
)
|
||||
cell = (
|
||||
(cell_x_start, bbox_y_start),
|
||||
(cell_x_end, bbox_y_start),
|
||||
(cell_x_end, bbox_y_end),
|
||||
(cell_x_start, bbox_y_end),
|
||||
)
|
||||
word_box_list.append(cell)
|
||||
|
||||
return word_box_content_list, word_box_list
|
||||
Reference in New Issue
Block a user