first commit
Some checks are pending
Build/Publish Develop Docs / deploy (push) Waiting to run

This commit is contained in:
2025-07-02 08:57:16 +03:00
commit 56532cc9a9
1901 changed files with 457695 additions and 0 deletions

159
ppstructure/table/README.md Normal file
View 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
![tableocr_pipeline](../docs/table/tableocr_pipeline_en.jpg)
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
![](../docs/imgs/table_ch_result1.jpg)
![](../docs/imgs/table_ch_result2.jpg)
![](../docs/imgs/table_ch_result3.jpg)
## 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

View 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
具体流程图如下
![tableocr_pipeline](../docs/table/tableocr_pipeline.jpg)
流程说明:
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. 效果演示
![](../docs/imgs/table_ch_result1.jpg)
![](../docs/imgs/table_ch_result2.jpg)
![](../docs/imgs/table_ch_result3.jpg)
## 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

View 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.

View 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
View 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
View 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

View 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())

View 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)

View 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

View 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

View 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

View 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)

View 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.

View 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

View 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)