This commit is contained in:
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)
|
||||
Reference in New Issue
Block a user