mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-11 21:40:52 +08:00
Remove rapid_structure
This commit is contained in:
SWHL committed
1 parent
eb3db42f06
commit
5dcecf41a5
43 files changed
+10
-2830
No files matched your search
@@ -1,15 +0,0 @@
|
||||
## Rapid Structure
|
||||
- 该部分的功能主要针对文档类图像，包括文档图像分类、版面分析和表格识别。
|
||||
|
||||
### [版面分析](https://github.com/RapidAI/RapidOCR/blob/main/python/rapid_structure/docs/README_Layout.md)
|
||||
|
||||
### [表格识别](https://github.com/RapidAI/RapidOCR/blob/main/python/rapid_structure/docs/README_Table.md)
|
||||
|
||||
### [文档方向分类](https://github.com/RapidAI/RapidOCR/blob/main/python/rapid_structure/docs/README_Orientation.md)
|
||||
|
||||
### 整体结构
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[/文档图像/] --> B(文档方向分类 rapid_orientation) --> C(版面分析 rapid_layout) & D(表格识别 rapid_table) --> E(OCR识别 rapidocr_onnxruntime)
|
||||
E --> F(结构化输出)
|
||||
```
|
||||
@@ -1,85 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import copy
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from rapid_layout import RapidLayout
|
||||
from rapid_orientation import RapidOrientation
|
||||
from rapid_table import RapidTable
|
||||
|
||||
|
||||
def vis_layout(img: np.ndarray, layout_res: list) -> None:
|
||||
tmp_img = copy.deepcopy(img)
|
||||
for v in layout_res:
|
||||
bbox = np.round(v['bbox']).astype(np.int32)
|
||||
label = v['label']
|
||||
|
||||
start_point = (bbox[0], bbox[1])
|
||||
end_point = (bbox[2], bbox[3])
|
||||
|
||||
cv2.rectangle(tmp_img, start_point, end_point, (0, 255, 0), 2)
|
||||
cv2.putText(tmp_img, label, start_point,
|
||||
cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2)
|
||||
|
||||
draw_img_save = Path("./inference_results/")
|
||||
if not draw_img_save.exists():
|
||||
draw_img_save.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
image_save = str(draw_img_save / 'layout_result.jpg')
|
||||
cv2.imwrite(image_save, tmp_img)
|
||||
print(f'The infer result has saved in {image_save}')
|
||||
|
||||
|
||||
def vis_table(table_res):
|
||||
style_res = '''<style>td {border-left: 1px solid;border-bottom:1px solid;}
|
||||
table, th {border-top:1px solid;font-size: 10px;
|
||||
border-collapse: collapse;border-right: 1px solid;}
|
||||
</style>'''
|
||||
prefix_table, suffix_table = table_res.split('<body>')
|
||||
new_table_res = f'{prefix_table}{style_res}<body>{suffix_table}'
|
||||
|
||||
draw_img_save = Path("./inference_results/")
|
||||
if not draw_img_save.exists():
|
||||
draw_img_save.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
html_path = str(draw_img_save / 'table_result.html')
|
||||
with open(html_path, 'w', encoding='utf-8') as f:
|
||||
f.write(new_table_res)
|
||||
print(f'The infer result has saved in {html_path}')
|
||||
|
||||
|
||||
def demo_layout():
|
||||
layout_engine = RapidLayout()
|
||||
|
||||
img = cv2.imread('test_images/layout.png')
|
||||
|
||||
layout_res, _ = layout_engine(img)
|
||||
|
||||
vis_layout(img, layout_res)
|
||||
print(layout_res)
|
||||
|
||||
|
||||
def demo_table():
|
||||
table_engine = RapidTable()
|
||||
img = cv2.imread('test_images/table.jpg')
|
||||
table_html_str, _ = table_engine(img)
|
||||
|
||||
vis_table(table_html_str)
|
||||
print(table_html_str)
|
||||
|
||||
|
||||
def demo_orientation():
|
||||
orientation_engine = RapidOrientation()
|
||||
img = cv2.imread('tests/test_files/img_rot180_demo.jpg')
|
||||
cls_result, _ = orientation_engine(img)
|
||||
print(cls_result)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
demo_layout()
|
||||
# demo_table()
|
||||
# demo_orientation()
|
||||
@@ -1,73 +0,0 @@
|
||||
## Rapid Layout
|
||||
<p align="left">
|
||||
<a href=""><img src="https://img.shields.io/badge/Python->=3.7,<=3.10-aff.svg"></a>
|
||||
<a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
|
||||
<a href="https://pypi.org/project/rapid-layout/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-layout"></a>
|
||||
<a href="https://pepy.tech/project/rapid-layout"><img src="https://static.pepy.tech/personalized-badge/rapid-layout?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
|
||||
</p>
|
||||
|
||||
#### 简介和说明
|
||||
- 该部分主要是做文档类图像的版面分析。模型来源：[PaddleOCR 版面分析](https://github.com/PaddlePaddle/PaddleOCR/blob/133d67f27dc8a241d6b2e30a9f047a0fb75bebbe/ppstructure/layout/README_ch.md)
|
||||
- 具体来说，就是分析给定的文档类别图像（论文截图等），定位其中类别和位置，如标题、段落、表格和图片等各个部分。
|
||||
- 目前支持三种类别的版面分析模型：中文、英文和表格版面分析模型，具体可参见下面表格：
|
||||
|
||||
| 模型类型 | 模型名称 | 模型大小 | 支持种类 |
|
||||
|:---:|:---:|:---:|:---:|
|
||||
| 表格 | `layout_table.onnx` | 7.06M | `table` |
|
||||
| 英文 | `layout_publaynet.onnx` | 7.06M | `text title list table figure` |
|
||||
| 中文 | `layout_cdla.onnx` | 7.07M | `text title figure figure_caption table table_caption` <br> `header footer reference equation` |
|
||||
- 模型下载地址为：[百度网盘](https://pan.baidu.com/s/1PI9fksW6F6kQfJhwUkewWg?pwd=p29g) | [Google Drive](https://drive.google.com/drive/folders/1DAPWSN2zGQ-ED_Pz7RaJGTjfkN2-Mvsf?usp=sharing)
|
||||
|
||||
#### 使用方式
|
||||
1. pip安装
|
||||
- 由于模型较小，预先将中文版面分析模型(`layout_cdla.onnx`)打包进了whl包内，如果做中文版面分析，可直接安装使用
|
||||
```bash
|
||||
$ pip install rapid-layout
|
||||
```
|
||||
2. python脚本运行
|
||||
```python
|
||||
import cv2
|
||||
from rapid_layout import RapidLayout
|
||||
|
||||
# RapidLayout类提供model_path参数，可以自行指定上述3个模型，默认是layout_cdla.onnx
|
||||
# layout_engine = RapidLayout(model_path='layout_publaynet.onnx')
|
||||
layout_engine = RapidLayout()
|
||||
|
||||
img = cv2.imread('test_images/layout.png')
|
||||
|
||||
layout_res, elapse = layout_engine(img)
|
||||
```
|
||||
|
||||
3. 终端运行
|
||||
- 用法:
|
||||
```bash
|
||||
$ rapid_layout -h
|
||||
usage: rapid_layout [-h] [-v] -img IMG_PATH [-m MODEL_PATH]
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
-v, --vis Wheter to visualize the layout results.
|
||||
-img IMG_PATH, --img_path IMG_PATH
|
||||
Path to image for layout.
|
||||
-m MODEL_PATH, --model_path MODEL_PATH
|
||||
The model path used for inference.
|
||||
```
|
||||
- 示例:
|
||||
```bash
|
||||
$ rapid_layout -v -img test_images/layout.png
|
||||
```
|
||||
|
||||
4. 结果
|
||||
- 返回结果
|
||||
```python
|
||||
# bbox: [左上角x0,左上角y0, 右下角x0, 右上角x1]
|
||||
# label: 类别
|
||||
[
|
||||
{'bbox': array([321.4160495, 91.53214898, 562.06141263, 199.85522603]), 'label': 'text'},
|
||||
{'bbox': array([58.67292211, 107.29000663, 300.25448676, 199.68142]), 'label': 'table_caption'}
|
||||
]
|
||||
```
|
||||
- 可视化结果
|
||||
<div align="center">
|
||||
<img src="https://warehouse-camo.ingress.cmh1.psfhosted.org/bbf4e9220b4d38ccefad2f86d1d894d6f2e8e98e/68747470733a2f2f7261772e67697468756275736572636f6e74656e742e636f6d2f526170696441492f52617069644f43522f393437633639353864333066343763376337623031366637646333303866323335616365633365652f707974686f6e2f72617069645f7374727563747572652f746573745f696d616765732f6c61796f75745f726573756c742e6a7067" width="80%" height="80%">
|
||||
</div>
|
||||
@@ -1,56 +0,0 @@
|
||||
## Rapid Orientation
|
||||
<p align="left">
|
||||
<a href=""><img src="https://img.shields.io/badge/Python->=3.7,<=3.10-aff.svg"></a>
|
||||
<a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
|
||||
<a href="https://pypi.org/project/rapid-orientation/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-orientation"></a>
|
||||
<a href="https://pepy.tech/project/rapid-orientation"><img src="https://static.pepy.tech/personalized-badge/rapid-orientation?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
|
||||
</p>
|
||||
|
||||
#### 简介和说明
|
||||
- 该部分主要是做含文字图像方向分类模型。模型来源：[PaddleClas 含文字图像方向分类模型](https://github.com/PaddlePaddle/PaddleClas/blob/177e4be74639c0960efeae2c5166d3226c9a02eb/docs/zh_CN/models/PULC/PULC_text_image_orientation.md)
|
||||
|
||||
| 模型类型 | 模型名称 | 模型大小 | 支持种类 |
|
||||
|:---:|:---:|:---:|:---:|
|
||||
| 四方向分类 | `rapid_orientation.onnx` | 6.5M | `0 90 180 270`|
|
||||
|
||||
#### 使用方式
|
||||
1. pip安装
|
||||
- 由于模型较小，已经将分类模型(`rapid_orientation.onnx`)打包进了whl包内：
|
||||
```bash
|
||||
$ pip install rapid-orientation
|
||||
```
|
||||
2. python脚本运行
|
||||
```python
|
||||
import cv2
|
||||
from rapid_orientation import RapidOrientation
|
||||
|
||||
orientation_engine = RapidOrientation()
|
||||
|
||||
img = cv2.imread('test_images/layout.png')
|
||||
|
||||
orientation_res, elapse = orientation_engine(img)
|
||||
print(orientation_res)
|
||||
```
|
||||
|
||||
3. 终端运行
|
||||
- 用法:
|
||||
```bash
|
||||
$ rapid_orientation -h
|
||||
usage: rapid_orientation [-h] -img IMG_PATH [-m MODEL_PATH]
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
-img IMG_PATH, --img_path IMG_PATH
|
||||
Path to image for layout.
|
||||
-m MODEL_PATH, --model_path MODEL_PATH
|
||||
The model path used for inference.
|
||||
```
|
||||
- 示例:
|
||||
```bash
|
||||
$ rapid_orientation -img test_images/layout.png
|
||||
```
|
||||
|
||||
4. 结果
|
||||
```python
|
||||
# 返回结果为str类型，有四类：0 | 90 | 180 | 270
|
||||
```
|
||||
@@ -1,66 +0,0 @@
|
||||
## Rapid Table
|
||||
<p align="left">
|
||||
<a href=""><img src="https://img.shields.io/badge/Python->=3.7,<=3.10-aff.svg"></a>
|
||||
<a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
|
||||
<a href="https://pypi.org/project/rapid-table/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-table"></a>
|
||||
<a href="https://pepy.tech/project/rapid-table"><img src="https://static.pepy.tech/personalized-badge/rapid-table?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
|
||||
</p>
|
||||
|
||||
#### 简介和说明
|
||||
- 该部分主要是做文档类图像的表格结构还原。模型来源：[PaddleOCR 表格识别](https://github.com/PaddlePaddle/PaddleOCR/blob/133d67f27dc8a241d6b2e30a9f047a0fb75bebbe/ppstructure/table/README_ch.md)
|
||||
- 具体来说，就是分析给定的表格图像，将表格还原为对应的html格式。
|
||||
- 目前支持两种类别的表格识别模型：中文和英文表格识别模型，具体可参见下面表格：
|
||||
|
||||
| 模型类型 | 模型名称 | 模型大小 |
|
||||
|:---:|:---:|:---:|
|
||||
| 英文 | `en_ppstructure_mobile_v2_SLANet.onnx` | 7.3M |
|
||||
| 中文 | `ch_ppstructure_mobile_v2_SLANet.onnx` | 7.4M |
|
||||
- 模型下载地址为：[百度网盘](https://pan.baidu.com/s/1PI9fksW6F6kQfJhwUkewWg?pwd=p29g) | [Google Drive](https://drive.google.com/drive/folders/1DAPWSN2zGQ-ED_Pz7RaJGTjfkN2-Mvsf?usp=sharing)
|
||||
|
||||
#### 使用方式
|
||||
1. pip安装
|
||||
- 由于模型较小，预先将英文表格识别模型(`en_ppstructure_mobile_v2_SLANet.onnx`)打包进了whl包内，如果做英文表格识别，可直接安装使用
|
||||
```bash
|
||||
$ pip install rapid-table
|
||||
```
|
||||
2. python脚本运行
|
||||
```python
|
||||
import cv2
|
||||
from rapid_table import RapidTable
|
||||
|
||||
# RapidTable类提供model_path参数，可以自行指定上述2个模型，默认是en_ppstructure_mobile_v2_SLANet.onnx
|
||||
# table_engine = RapidTable(model_path='ch_ppstructure_mobile_v2_SLANet.onnx')
|
||||
table_engine = RapidTable()
|
||||
|
||||
img = cv2.imread('test_images/table.jpg')
|
||||
table_html_str, _ = table_engine(img)
|
||||
print(table_html_str)
|
||||
```
|
||||
3. 终端运行
|
||||
- 用法:
|
||||
```bash
|
||||
$ rapid_table -h
|
||||
usage: rapid_table [-h] [-v] -img IMG_PATH [-m MODEL_PATH]
|
||||
|
||||
optional arguments:
|
||||
-h, --help show this help message and exit
|
||||
-v, --vis Wheter to visualize the layout results.
|
||||
-img IMG_PATH, --img_path IMG_PATH
|
||||
Path to image for layout.
|
||||
-m MODEL_PATH, --model_path MODEL_PATH
|
||||
The model path used for inference.
|
||||
```
|
||||
- 示例:
|
||||
```bash
|
||||
$ rapid_table -v -img test_images/table.jpg
|
||||
```
|
||||
|
||||
4. 结果
|
||||
- 返回结果
|
||||
```html
|
||||
<html><body><table><thead><tr><td>Methods</td><td></td><td></td><td></td><td>FPS</td></tr></thead><tbody><tr><td>SegLink [26]</td><td>70.0</td><td>86d><td.0</td><td>77.0</td><td>8.9</td></tr><tr><td>PixelLink [4]</td><td>73.2</td><td>83.0</td><td>77.8</td><td></td></tr><tr><td>TextSnake [18]</td><td>73.9</td><td>83.2</td><td>78.3</td><td>1.1</td></tr><tr><td>TextField [37]</td><td>75.9</td><td>87.4</td><td>81.3</td><td>5.2</td></tr><tr><td>MSR[38]</td><td>76.7</td><td>87.87.4</td><td>81.7</td><td></td></tr><tr><td>FTSN [3]</td><td>77.1</td><td>87.6</td><td>82.0</td><td></td></tr><tr><td>LSE[30]</td><td>81.7</td><td>84.2</td><td>82.9</td><><ttd></td></tr><tr><td>CRAFT [2]</td><td>78.2</td><td>88.2</td><td>82.9</td><td>8.6</td></tr><tr><td>MCN[16]</td><td>79</td><td>88</td><td>83</td><td></td></tr><tr><td>ATRR</>[35]</td><td>82.1</td><td>85.2</td><td>83.6</td><td></td></tr><tr><td>PAN [34]</td><td>83.8</td><td>84.4</td><td>84.1</td><td>30.2</td></tr><tr><td>DB[12]</td><td>79.2</t91/d><td>91.5</td><td>84.9</td><td>32.0</td></tr><tr><td>DRRG[41]</td><td>82.30</td><td>88.05</td><td>85.08</td><td></td></tr><tr><td>Ours (SynText)</td><td>80.68</td><td>85<t..40</td><td>82.97</td><td>12.68</td></tr><tr><td>Ours (MLT-17)</td><td>84.54</td><td>86.62</td><td>85.57</td><td>12.31</td></tr></tbody></table></body></html>
|
||||
```
|
||||
- 可视化结果
|
||||
<div align="center">
|
||||
<table><thead><tr><td>Methods</td><td></td><td></td><td></td><td>FPS</td></tr></thead><tbody><tr><td>SegLink [26]</td><td>70.0</td><td>86d><td.0</td><td>77.0</td><td>8.9</td></tr><tr><td>PixelLink [4]</td><td>73.2</td><td>83.0</td><td>77.8</td><td></td></tr><tr><td>TextSnake [18]</td><td>73.9</td><td>83.2</td><td>78.3</td><td>1.1</td></tr><tr><td>TextField [37]</td><td>75.9</td><td>87.4</td><td>81.3</td><td>5.2</td></tr><tr><td>MSR[38]</td><td>76.7</td><td>87.87.4</td><td>81.7</td><td></td></tr><tr><td>FTSN [3]</td><td>77.1</td><td>87.6</td><td>82.0</td><td></td></tr><tr><td>LSE[30]</td><td>81.7</td><td>84.2</td><td>82.9</td><><ttd></td></tr><tr><td>CRAFT [2]</td><td>78.2</td><td>88.2</td><td>82.9</td><td>8.6</td></tr><tr><td>MCN[16]</td><td>79</td><td>88</td><td>83</td><td></td></tr><tr><td>ATRR</>[35]</td><td>82.1</td><td>85.2</td><td>83.6</td><td></td></tr><tr><td>PAN [34]</td><td>83.8</td><td>84.4</td><td>84.1</td><td>30.2</td></tr><tr><td>DB[12]</td><td>79.2</t91/d><td>91.5</td><td>84.9</td><td>32.0</td></tr><tr><td>DRRG[41]</td><td>82.30</td><td>88.05</td><td>85.08</td><td></td></tr><tr><td>Ours (SynText)</td><td>80.68</td><td>85<t..40</td><td>82.97</td><td>12.68</td></tr><tr><td>Ours (MLT-17)</td><td>84.54</td><td>86.62</td><td>85.57</td><td>12.31</td></tr></tbody></table>
|
||||
</div>
|
||||
@@ -1,14 +0,0 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from .rapid_layout import RapidLayout
|
||||
@@ -1,24 +0,0 @@
|
||||
model_path: models/layout_cdla.onnx
|
||||
|
||||
use_cuda: false
|
||||
CUDAExecutionProvider:
|
||||
device_id: 0
|
||||
arena_extend_strategy: kNextPowerOfTwo
|
||||
cudnn_conv_algo_search: EXHAUSTIVE
|
||||
do_copy_in_default_stream: true
|
||||
|
||||
pre_process:
|
||||
Resize:
|
||||
size: [800, 608]
|
||||
NormalizeImage:
|
||||
std: [0.229, 0.224, 0.225]
|
||||
mean: [0.485, 0.456, 0.406]
|
||||
scale: 1./255.
|
||||
order: hwc
|
||||
ToCHWImage:
|
||||
KeepKeys:
|
||||
keep_keys: ['image']
|
||||
|
||||
post_process:
|
||||
score_threshold: 0.5
|
||||
nms_threshold: 0.5
|
||||
Whitespace-only changes.
@@ -1,95 +0,0 @@
|
||||
# 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 argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .utils import (OrtInferSession, PicoDetPostProcess, create_operators,
|
||||
read_yaml, transform, vis_layout)
|
||||
|
||||
root_dir = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
class RapidLayout():
|
||||
def __init__(self,
|
||||
model_path: str = None):
|
||||
config_path = str(root_dir / 'config.yaml')
|
||||
config = read_yaml(config_path)
|
||||
if model_path is None:
|
||||
model_path = str(root_dir / 'models' / 'layout_cdla.onnx')
|
||||
config['model_path'] = model_path
|
||||
|
||||
self.session = OrtInferSession(config)
|
||||
labels = self.session.get_metadata()['character'].splitlines()
|
||||
|
||||
self.preprocess_op = create_operators(config['pre_process'])
|
||||
self.postprocess_op = PicoDetPostProcess(labels,
|
||||
**config['post_process'])
|
||||
|
||||
def __call__(self, img):
|
||||
ori_im = img.copy()
|
||||
data = {'image': img}
|
||||
data = transform(data, self.preprocess_op)
|
||||
img = data[0]
|
||||
|
||||
if img is None:
|
||||
return None, 0
|
||||
|
||||
img = np.expand_dims(img, axis=0)
|
||||
img = img.copy()
|
||||
|
||||
preds, elapse = 0, 1
|
||||
starttime = time.time()
|
||||
|
||||
preds = self.session(img)
|
||||
|
||||
score_list, boxes_list = [], []
|
||||
num_outs = int(len(preds) / 2)
|
||||
for out_idx in range(num_outs):
|
||||
score_list.append(preds[out_idx])
|
||||
boxes_list.append(preds[out_idx + num_outs])
|
||||
preds = dict(boxes=score_list, boxes_num=boxes_list)
|
||||
post_preds = self.postprocess_op(ori_im, img, preds)
|
||||
elapse = time.time() - starttime
|
||||
return post_preds, elapse
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-v', '--vis', action='store_true',
|
||||
help='Wheter to visualize the layout results.')
|
||||
parser.add_argument('-img', '--img_path', type=str, required=True,
|
||||
help='Path to image for layout.')
|
||||
parser.add_argument('-m', '--model_path', type=str,
|
||||
default=str(root_dir / 'models' / 'layout_cdla.onnx'),
|
||||
help='The model path used for inference.')
|
||||
args = parser.parse_args()
|
||||
|
||||
layout_engine = RapidLayout(args.model_path)
|
||||
|
||||
img = cv2.imread(args.img_path)
|
||||
layout_res, elapse = layout_engine(img)
|
||||
print(layout_res)
|
||||
|
||||
if args.vis:
|
||||
img_path = Path(args.img_path)
|
||||
save_path = img_path.resolve().parent / f'vis_{img_path.name}'
|
||||
vis_layout(img, layout_res, str(save_path))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,443 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import copy
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
import cv2
|
||||
import numpy as np
|
||||
from onnxruntime import (get_available_providers, get_device,
|
||||
SessionOptions, InferenceSession,
|
||||
GraphOptimizationLevel)
|
||||
|
||||
|
||||
class OrtInferSession():
|
||||
def __init__(self, config):
|
||||
sess_opt = SessionOptions()
|
||||
sess_opt.log_severity_level = 4
|
||||
sess_opt.enable_cpu_mem_arena = False
|
||||
sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
cuda_ep = 'CUDAExecutionProvider'
|
||||
cpu_ep = 'CPUExecutionProvider'
|
||||
cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested",}
|
||||
|
||||
EP_list = []
|
||||
if config['use_cuda'] and get_device() == 'GPU' \
|
||||
and cuda_ep in get_available_providers():
|
||||
EP_list = [(cuda_ep, config[cuda_ep])]
|
||||
EP_list.append((cpu_ep, cpu_provider_options))
|
||||
|
||||
self._verify_model(config['model_path'])
|
||||
self.session = InferenceSession(config['model_path'],
|
||||
sess_options=sess_opt,
|
||||
providers=EP_list)
|
||||
|
||||
if config['use_cuda'] and cuda_ep not in self.session.get_providers():
|
||||
warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n'
|
||||
'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, '
|
||||
'you can check their relations from the offical web site: '
|
||||
'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
|
||||
RuntimeWarning)
|
||||
|
||||
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
||||
input_dict = dict(zip(self.get_input_names(), [input_content]))
|
||||
try:
|
||||
return self.session.run(self.get_output_names(), input_dict)
|
||||
except Exception as e:
|
||||
raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
|
||||
|
||||
def get_input_names(self, ):
|
||||
return [v.name for v in self.session.get_inputs()]
|
||||
|
||||
def get_output_names(self,):
|
||||
return [v.name for v in self.session.get_outputs()]
|
||||
|
||||
def get_metadata(self):
|
||||
meta_dict = self.session.get_modelmeta().custom_metadata_map
|
||||
return meta_dict
|
||||
|
||||
@staticmethod
|
||||
def _verify_model(model_path):
|
||||
model_path = Path(model_path)
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(f'{model_path} does not exists.')
|
||||
if not model_path.is_file():
|
||||
raise FileExistsError(f'{model_path} is not a file.')
|
||||
|
||||
|
||||
class ONNXRuntimeError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def transform(data, ops=None):
|
||||
""" transform """
|
||||
if ops is None:
|
||||
ops = []
|
||||
|
||||
for op in ops:
|
||||
data = op(data)
|
||||
if data is None:
|
||||
return None
|
||||
return data
|
||||
|
||||
|
||||
def create_operators(op_param_dict):
|
||||
ops = []
|
||||
for op_name, param in op_param_dict.items():
|
||||
if param is None:
|
||||
param = {}
|
||||
op = eval(op_name)(**param)
|
||||
ops.append(op)
|
||||
return ops
|
||||
|
||||
|
||||
class Resize():
|
||||
def __init__(self, size=(640, 640), **kwargs):
|
||||
self.size = size
|
||||
|
||||
def resize_image(self, img):
|
||||
resize_h, resize_w = self.size
|
||||
ori_h, ori_w = img.shape[:2] # (h, w, c)
|
||||
ratio_h = float(resize_h) / ori_h
|
||||
ratio_w = float(resize_w) / ori_w
|
||||
img = cv2.resize(img, (int(resize_w), int(resize_h)))
|
||||
return img, [ratio_h, ratio_w]
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
if 'polys' in data:
|
||||
text_polys = data['polys']
|
||||
|
||||
img_resize, [ratio_h, ratio_w] = self.resize_image(img)
|
||||
if 'polys' in data:
|
||||
new_boxes = []
|
||||
for box in text_polys:
|
||||
new_box = []
|
||||
for cord in box:
|
||||
new_box.append([cord[0] * ratio_w, cord[1] * ratio_h])
|
||||
new_boxes.append(new_box)
|
||||
data['polys'] = np.array(new_boxes, dtype=np.float32)
|
||||
data['image'] = img_resize
|
||||
return data
|
||||
|
||||
|
||||
class NormalizeImage():
|
||||
def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
|
||||
if isinstance(scale, str):
|
||||
scale = eval(scale)
|
||||
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
|
||||
mean = mean if mean is not None else [0.485, 0.456, 0.406]
|
||||
std = std if std is not None else [0.229, 0.224, 0.225]
|
||||
|
||||
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
|
||||
self.mean = np.array(mean).reshape(shape).astype('float32')
|
||||
self.std = np.array(std).reshape(shape).astype('float32')
|
||||
|
||||
def __call__(self, data):
|
||||
img = np.array(data['image'])
|
||||
assert isinstance(img,
|
||||
np.ndarray), "invalid input 'img' in NormalizeImage"
|
||||
data['image'] = (
|
||||
img.astype('float32') * self.scale - self.mean) / self.std
|
||||
return data
|
||||
|
||||
|
||||
class ToCHWImage():
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = np.array(data['image'])
|
||||
data['image'] = img.transpose((2, 0, 1))
|
||||
return data
|
||||
|
||||
|
||||
class KeepKeys():
|
||||
def __init__(self, keep_keys, **kwargs):
|
||||
self.keep_keys = keep_keys
|
||||
|
||||
def __call__(self, data):
|
||||
data_list = []
|
||||
for key in self.keep_keys:
|
||||
data_list.append(data[key])
|
||||
return data_list
|
||||
|
||||
|
||||
def read_yaml(yaml_path):
|
||||
with open(yaml_path, 'rb') as f:
|
||||
data = yaml.load(f, Loader=yaml.Loader)
|
||||
return data
|
||||
|
||||
|
||||
class PicoDetPostProcess():
|
||||
def __init__(self,
|
||||
labels,
|
||||
strides=[8, 16, 32, 64],
|
||||
score_threshold=0.4,
|
||||
nms_threshold=0.5,
|
||||
nms_top_k=1000,
|
||||
keep_top_k=100):
|
||||
self.labels = labels
|
||||
self.strides = strides
|
||||
self.score_threshold = score_threshold
|
||||
self.nms_threshold = nms_threshold
|
||||
self.nms_top_k = nms_top_k
|
||||
self.keep_top_k = keep_top_k
|
||||
|
||||
def __call__(self, ori_img, img, preds):
|
||||
scores, raw_boxes = preds['boxes'], preds['boxes_num']
|
||||
batch_size = raw_boxes[0].shape[0]
|
||||
reg_max = int(raw_boxes[0].shape[-1] / 4 - 1)
|
||||
out_boxes_num = []
|
||||
out_boxes_list = []
|
||||
results = []
|
||||
ori_shape, input_shape, scale_factor = self.img_info(ori_img, img)
|
||||
|
||||
for batch_id in range(batch_size):
|
||||
# generate centers
|
||||
decode_boxes = []
|
||||
select_scores = []
|
||||
for stride, box_distribute, score in zip(self.strides, raw_boxes,
|
||||
scores):
|
||||
box_distribute = box_distribute[batch_id]
|
||||
score = score[batch_id]
|
||||
# centers
|
||||
fm_h = input_shape[0] / stride
|
||||
fm_w = input_shape[1] / stride
|
||||
h_range = np.arange(fm_h)
|
||||
w_range = np.arange(fm_w)
|
||||
ww, hh = np.meshgrid(w_range, h_range)
|
||||
ct_row = (hh.flatten() + 0.5) * stride
|
||||
ct_col = (ww.flatten() + 0.5) * stride
|
||||
center = np.stack((ct_col, ct_row, ct_col, ct_row), axis=1)
|
||||
|
||||
# box distribution to distance
|
||||
reg_range = np.arange(reg_max + 1)
|
||||
box_distance = box_distribute.reshape((-1, reg_max + 1))
|
||||
box_distance = self.softmax(box_distance, axis=1)
|
||||
box_distance = box_distance * np.expand_dims(reg_range, axis=0)
|
||||
box_distance = np.sum(box_distance, axis=1).reshape((-1, 4))
|
||||
box_distance = box_distance * stride
|
||||
|
||||
# top K candidate
|
||||
topk_idx = np.argsort(score.max(axis=1))[::-1]
|
||||
topk_idx = topk_idx[:self.nms_top_k]
|
||||
center = center[topk_idx]
|
||||
score = score[topk_idx]
|
||||
box_distance = box_distance[topk_idx]
|
||||
|
||||
# decode box
|
||||
decode_box = center + [-1, -1, 1, 1] * box_distance
|
||||
|
||||
select_scores.append(score)
|
||||
decode_boxes.append(decode_box)
|
||||
|
||||
# nms
|
||||
bboxes = np.concatenate(decode_boxes, axis=0)
|
||||
confidences = np.concatenate(select_scores, axis=0)
|
||||
picked_box_probs = []
|
||||
picked_labels = []
|
||||
for class_index in range(0, confidences.shape[1]):
|
||||
probs = confidences[:, class_index]
|
||||
mask = probs > self.score_threshold
|
||||
probs = probs[mask]
|
||||
if probs.shape[0] == 0:
|
||||
continue
|
||||
subset_boxes = bboxes[mask, :]
|
||||
box_probs = np.concatenate(
|
||||
[subset_boxes, probs.reshape(-1, 1)], axis=1)
|
||||
box_probs = self.hard_nms(
|
||||
box_probs,
|
||||
iou_threshold=self.nms_threshold,
|
||||
top_k=self.keep_top_k, )
|
||||
picked_box_probs.append(box_probs)
|
||||
picked_labels.extend([class_index] * box_probs.shape[0])
|
||||
|
||||
if len(picked_box_probs) == 0:
|
||||
out_boxes_list.append(np.empty((0, 4)))
|
||||
out_boxes_num.append(0)
|
||||
|
||||
else:
|
||||
picked_box_probs = np.concatenate(picked_box_probs)
|
||||
|
||||
# resize output boxes
|
||||
picked_box_probs[:, :4] = self.warp_boxes(
|
||||
picked_box_probs[:, :4], ori_shape[batch_id])
|
||||
im_scale = np.concatenate([
|
||||
scale_factor[batch_id][::-1], scale_factor[batch_id][::-1]
|
||||
])
|
||||
picked_box_probs[:, :4] /= im_scale
|
||||
# clas score box
|
||||
out_boxes_list.append(
|
||||
np.concatenate(
|
||||
[
|
||||
np.expand_dims(
|
||||
np.array(picked_labels),
|
||||
axis=-1), np.expand_dims(
|
||||
picked_box_probs[:, 4], axis=-1),
|
||||
picked_box_probs[:, :4]
|
||||
],
|
||||
axis=1))
|
||||
out_boxes_num.append(len(picked_labels))
|
||||
|
||||
out_boxes_list = np.concatenate(out_boxes_list, axis=0)
|
||||
out_boxes_num = np.asarray(out_boxes_num).astype(np.int32)
|
||||
|
||||
for dt in out_boxes_list:
|
||||
clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
|
||||
label = self.labels[clsid]
|
||||
result = {'bbox': bbox, 'label': label}
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def load_layout_dict(self, layout_dict_path):
|
||||
with open(layout_dict_path, 'r', encoding='utf-8') as fp:
|
||||
labels = fp.readlines()
|
||||
return [label.strip('\n') for label in labels]
|
||||
|
||||
def warp_boxes(self, boxes, ori_shape):
|
||||
"""Apply transform to boxes
|
||||
"""
|
||||
width, height = ori_shape[1], ori_shape[0]
|
||||
n = len(boxes)
|
||||
if n:
|
||||
# warp points
|
||||
xy = np.ones((n * 4, 3))
|
||||
xy[:, :2] = boxes[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(
|
||||
n * 4, 2) # x1y1, x2y2, x1y2, x2y1
|
||||
# xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3]).reshape(n, 8) # rescale
|
||||
# create new boxes
|
||||
x = xy[:, [0, 2, 4, 6]]
|
||||
y = xy[:, [1, 3, 5, 7]]
|
||||
xy = np.concatenate(
|
||||
(x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
|
||||
# clip boxes
|
||||
xy[:, [0, 2]] = xy[:, [0, 2]].clip(0, width)
|
||||
xy[:, [1, 3]] = xy[:, [1, 3]].clip(0, height)
|
||||
return xy.astype(np.float32)
|
||||
else:
|
||||
return boxes
|
||||
|
||||
def img_info(self, ori_img, img):
|
||||
origin_shape = ori_img.shape
|
||||
resize_shape = img.shape
|
||||
im_scale_y = resize_shape[2] / float(origin_shape[0])
|
||||
im_scale_x = resize_shape[3] / float(origin_shape[1])
|
||||
scale_factor = np.array([im_scale_y, im_scale_x], dtype=np.float32)
|
||||
img_shape = np.array(img.shape[2:], dtype=np.float32)
|
||||
|
||||
input_shape = np.array(img).astype('float32').shape[2:]
|
||||
ori_shape = np.array((img_shape, )).astype('float32')
|
||||
scale_factor = np.array((scale_factor, )).astype('float32')
|
||||
return ori_shape, input_shape, scale_factor
|
||||
|
||||
@staticmethod
|
||||
def softmax(x, axis=None):
|
||||
def logsumexp(a, axis=None, b=None, keepdims=False):
|
||||
a_max = np.amax(a, axis=axis, keepdims=True)
|
||||
|
||||
if a_max.ndim > 0:
|
||||
a_max[~np.isfinite(a_max)] = 0
|
||||
elif not np.isfinite(a_max):
|
||||
a_max = 0
|
||||
|
||||
tmp = np.exp(a - a_max)
|
||||
|
||||
# suppress warnings about log of zero
|
||||
with np.errstate(divide='ignore'):
|
||||
s = np.sum(tmp, axis=axis, keepdims=keepdims)
|
||||
out = np.log(s)
|
||||
|
||||
if not keepdims:
|
||||
a_max = np.squeeze(a_max, axis=axis)
|
||||
out += a_max
|
||||
return out
|
||||
|
||||
return np.exp(x - logsumexp(x, axis=axis, keepdims=True))
|
||||
|
||||
def hard_nms(self, box_scores, iou_threshold, top_k=-1,
|
||||
candidate_size=200):
|
||||
"""
|
||||
Args:
|
||||
box_scores (N, 5): boxes in corner-form and probabilities.
|
||||
iou_threshold: intersection over union threshold.
|
||||
top_k: keep top_k results. If k <= 0, keep all the results.
|
||||
candidate_size: only consider the candidates with the highest scores.
|
||||
Returns:
|
||||
picked: a list of indexes of the kept boxes
|
||||
"""
|
||||
scores = box_scores[:, -1]
|
||||
boxes = box_scores[:, :-1]
|
||||
picked = []
|
||||
indexes = np.argsort(scores)
|
||||
indexes = indexes[-candidate_size:]
|
||||
while len(indexes) > 0:
|
||||
current = indexes[-1]
|
||||
picked.append(current)
|
||||
if 0 < top_k == len(picked) or len(indexes) == 1:
|
||||
break
|
||||
current_box = boxes[current, :]
|
||||
indexes = indexes[:-1]
|
||||
rest_boxes = boxes[indexes, :]
|
||||
iou = self.iou_of(rest_boxes, np.expand_dims(current_box, axis=0),)
|
||||
indexes = indexes[iou <= iou_threshold]
|
||||
|
||||
return box_scores[picked, :]
|
||||
|
||||
def iou_of(self, boxes0, boxes1, eps=1e-5):
|
||||
"""Return intersection-over-union (Jaccard index) of boxes.
|
||||
Args:
|
||||
boxes0 (N, 4): ground truth boxes.
|
||||
boxes1 (N or 1, 4): predicted boxes.
|
||||
eps: a small number to avoid 0 as denominator.
|
||||
Returns:
|
||||
iou (N): IoU values.
|
||||
"""
|
||||
overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
|
||||
overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
|
||||
|
||||
overlap_area = self.area_of(overlap_left_top, overlap_right_bottom)
|
||||
area0 = self.area_of(boxes0[..., :2], boxes0[..., 2:])
|
||||
area1 = self.area_of(boxes1[..., :2], boxes1[..., 2:])
|
||||
return overlap_area / (area0 + area1 - overlap_area + eps)
|
||||
|
||||
@staticmethod
|
||||
def area_of(left_top, right_bottom):
|
||||
"""Compute the areas of rectangles given two corners.
|
||||
Args:
|
||||
left_top (N, 2): left top corner.
|
||||
right_bottom (N, 2): right bottom corner.
|
||||
Returns:
|
||||
area (N): return the area.
|
||||
"""
|
||||
hw = np.clip(right_bottom - left_top, 0.0, None)
|
||||
return hw[..., 0] * hw[..., 1]
|
||||
|
||||
|
||||
def vis_layout(img: np.ndarray, layout_res: list, save_path: str) -> None:
|
||||
font = cv2.FONT_HERSHEY_COMPLEX
|
||||
font_scale = 1
|
||||
font_color = (0, 0, 255)
|
||||
font_thickness = 1
|
||||
|
||||
tmp_img = copy.deepcopy(img)
|
||||
for v in layout_res:
|
||||
bbox = np.round(v['bbox']).astype(np.int32)
|
||||
label = v['label']
|
||||
|
||||
start_point = (bbox[0], bbox[1])
|
||||
end_point = (bbox[2], bbox[3])
|
||||
|
||||
cv2.rectangle(tmp_img, start_point, end_point, (0, 255, 0), 2)
|
||||
|
||||
(w, h), _ = cv2.getTextSize(label, font, font_scale, font_thickness)
|
||||
put_point = start_point[0], start_point[1] + h
|
||||
cv2.putText(tmp_img, label, put_point, font, font_scale,
|
||||
font_color, font_thickness)
|
||||
|
||||
cv2.imwrite(save_path, tmp_img)
|
||||
print(f'The infer result has saved in {save_path}')
|
||||
@@ -1,4 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
from .rapid_orientation import RapidOrientation
|
||||
@@ -1,16 +0,0 @@
|
||||
model_path: models/rapid_orientation.onnx
|
||||
|
||||
use_cuda: false
|
||||
CUDAExecutionProvider:
|
||||
device_id: 0
|
||||
arena_extend_strategy: kNextPowerOfTwo
|
||||
cudnn_conv_algo_search: EXHAUSTIVE
|
||||
do_copy_in_default_stream: true
|
||||
|
||||
PreProcess:
|
||||
- ResizeImage:
|
||||
resize_short: 256
|
||||
- CropImage:
|
||||
size: 224
|
||||
- NormalizeImage:
|
||||
- ToCHWImage:
|
||||
Whitespace-only changes.
@@ -1,82 +0,0 @@
|
||||
# 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.
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
from .utils import OrtInferSession, create_operators
|
||||
|
||||
root_dir = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
class RapidOrientation():
|
||||
def __init__(self, model_path: str = None):
|
||||
config_path = str(root_dir / 'config.yaml')
|
||||
config = self.read_yaml(config_path)
|
||||
if model_path is None:
|
||||
model_path = str(root_dir / 'models' / 'rapid_orientation.onnx')
|
||||
config['model_path'] = model_path
|
||||
|
||||
self.session = OrtInferSession(config)
|
||||
self.labels = self.session.get_metadata()['character'].splitlines()
|
||||
|
||||
self.preprocess_ops = create_operators(config["PreProcess"])
|
||||
|
||||
def __call__(self, images):
|
||||
s = time.time()
|
||||
for ops in self.preprocess_ops:
|
||||
images = ops(images)
|
||||
image = np.array(images)[None, ...]
|
||||
|
||||
pred_output = self.session(image)[0]
|
||||
|
||||
pred_output = pred_output.squeeze()
|
||||
pred_idx = np.argmax(pred_output)
|
||||
pred_txt = self.labels[pred_idx]
|
||||
elapse = time.time() - s
|
||||
return pred_txt, elapse
|
||||
|
||||
@staticmethod
|
||||
def read_yaml(yaml_path):
|
||||
with open(yaml_path, 'rb') as f:
|
||||
data = yaml.load(f, Loader=yaml.Loader)
|
||||
return data
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-img', '--img_path', type=str, required=True,
|
||||
help='Path to image for layout.')
|
||||
parser.add_argument('-m', '--model_path', type=str,
|
||||
default=str(root_dir / 'models' /
|
||||
'rapid_orientation.onnx'),
|
||||
help='The model path used for inference.')
|
||||
args = parser.parse_args()
|
||||
|
||||
orientation_engine = RapidOrientation(args.model_path)
|
||||
|
||||
img = cv2.imread(args.img_path)
|
||||
orientaion_result, _ = orientation_engine(img)
|
||||
print(orientaion_result)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,169 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import warnings
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from onnxruntime import (GraphOptimizationLevel, InferenceSession,
|
||||
SessionOptions, get_available_providers, get_device)
|
||||
|
||||
|
||||
class OrtInferSession():
|
||||
def __init__(self, config):
|
||||
sess_opt = SessionOptions()
|
||||
sess_opt.log_severity_level = 4
|
||||
sess_opt.enable_cpu_mem_arena = False
|
||||
sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
cuda_ep = 'CUDAExecutionProvider'
|
||||
cpu_ep = 'CPUExecutionProvider'
|
||||
cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested", }
|
||||
|
||||
EP_list = []
|
||||
if config['use_cuda'] and get_device() == 'GPU' \
|
||||
and cuda_ep in get_available_providers():
|
||||
EP_list = [(cuda_ep, config[cuda_ep])]
|
||||
EP_list.append((cpu_ep, cpu_provider_options))
|
||||
|
||||
self._verify_model(config['model_path'])
|
||||
self.session = InferenceSession(config['model_path'],
|
||||
sess_options=sess_opt,
|
||||
providers=EP_list)
|
||||
|
||||
has_cuda_ep = cuda_ep not in self.session.get_providers()
|
||||
if config['use_cuda'] and has_cuda_ep:
|
||||
warnings.warn(f'{cuda_ep} is not avaiable for current env,'
|
||||
f'the inference part is automatically shifted to '
|
||||
f'be executed under {cpu_ep}. '
|
||||
f'Please ensure the installed onnxruntime-gpu '
|
||||
f' version matches your cuda and cudnn version, '
|
||||
f'you can check their relations from the offical web site: '
|
||||
f'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
|
||||
RuntimeWarning)
|
||||
|
||||
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
||||
input_dict = dict(zip(self.get_input_names(), [input_content]))
|
||||
try:
|
||||
return self.session.run(self.get_output_names(), input_dict)
|
||||
except Exception as e:
|
||||
raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
|
||||
|
||||
def get_input_names(self, ):
|
||||
return [v.name for v in self.session.get_inputs()]
|
||||
|
||||
def get_output_names(self,):
|
||||
return [v.name for v in self.session.get_outputs()]
|
||||
|
||||
def get_metadata(self):
|
||||
meta_dict = self.session.get_modelmeta().custom_metadata_map
|
||||
return meta_dict
|
||||
|
||||
@staticmethod
|
||||
def _verify_model(model_path):
|
||||
model_path = Path(model_path)
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(f'{model_path} does not exists.')
|
||||
if not model_path.is_file():
|
||||
raise FileExistsError(f'{model_path} is not a file.')
|
||||
|
||||
|
||||
class ONNXRuntimeError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def create_operators(params):
|
||||
"""
|
||||
create operators based on the config
|
||||
|
||||
Args:
|
||||
params(list): a dict list, used to create some operators
|
||||
"""
|
||||
assert isinstance(params, list), ('operator config should be a list')
|
||||
mod = importlib.import_module(__name__)
|
||||
ops = []
|
||||
for operator in params:
|
||||
assert isinstance(operator,
|
||||
dict) and len(operator) == 1, "yaml format error"
|
||||
op_name = list(operator)[0]
|
||||
param = {} if operator[op_name] is None else operator[op_name]
|
||||
op = getattr(mod, op_name)(**param)
|
||||
ops.append(op)
|
||||
return ops
|
||||
|
||||
|
||||
class ResizeImage():
|
||||
def __init__(self, size=None, resize_short=None):
|
||||
if resize_short is not None and resize_short > 0:
|
||||
self.resize_short = resize_short
|
||||
self.w, self.h = None, None
|
||||
elif size is not None:
|
||||
self.resize_short = None
|
||||
self.w = size if isinstance(size, int) else size[0]
|
||||
self.h = size if isinstance(size, int) else size[1]
|
||||
else:
|
||||
raise ValueError("invalid params for ReisizeImage for '\
|
||||
'both 'size' and 'resize_short' are None")
|
||||
|
||||
def __call__(self, img: np.ndarray):
|
||||
img_h, img_w = img.shape[:2]
|
||||
|
||||
if self.resize_short:
|
||||
percent = float(self.resize_short) / min(img_w, img_h)
|
||||
w = int(round(img_w * percent))
|
||||
h = int(round(img_h * percent))
|
||||
else:
|
||||
w = self.w
|
||||
h = self.h
|
||||
return cv2.resize(img, (w, h))
|
||||
|
||||
|
||||
class NormalizeImage():
|
||||
def __init__(self,):
|
||||
self.scale = np.float32(1.0 / 255.0)
|
||||
mean = [0.485, 0.456, 0.406]
|
||||
std = [0.229, 0.224, 0.225]
|
||||
|
||||
shape = 1, 1, 3
|
||||
self.mean = np.array(mean).reshape(shape).astype('float32')
|
||||
self.std = np.array(std).reshape(shape).astype('float32')
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img).astype(np.float32)
|
||||
img = (img * self.scale - self.mean) / self.std
|
||||
return img.astype(np.float32)
|
||||
|
||||
|
||||
class ToCHWImage():
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, img):
|
||||
img = np.array(img)
|
||||
return img.transpose((2, 0, 1))
|
||||
|
||||
|
||||
class CropImage():
|
||||
def __init__(self, size):
|
||||
self.size = size
|
||||
if isinstance(size, int):
|
||||
self.size = (size, size)
|
||||
|
||||
def __call__(self, img):
|
||||
w, h = self.size
|
||||
img_h, img_w = img.shape[:2]
|
||||
|
||||
if img_h < h or img_w < w:
|
||||
raise ValueError(
|
||||
f"The size({h}, {w}) of CropImage must be greater than "
|
||||
f"size({img_h}, {img_w}) of image."
|
||||
)
|
||||
|
||||
w_start = (img_w - w) // 2
|
||||
h_start = (img_h - h) // 2
|
||||
|
||||
w_end = w_start + w
|
||||
h_end = h_start + h
|
||||
return img[h_start:h_end, w_start:w_end, :]
|
||||
@@ -1,4 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
from .rapid_table import RapidTable
|
||||
Whitespace-only changes.
@@ -1,95 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import argparse
|
||||
import copy
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from rapidocr_onnxruntime import RapidOCR
|
||||
|
||||
from .table_matcher import TableMatch
|
||||
from .table_structure import TableStructurer
|
||||
|
||||
root_dir = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
class RapidTable():
|
||||
def __init__(self, model_path: str = None):
|
||||
self.ocr_sys = RapidOCR()
|
||||
|
||||
if model_path is None:
|
||||
model_path = str(root_dir / 'models' / 'en_ppstructure_mobile_v2_SLANet.onnx')
|
||||
|
||||
self.table_structure = TableStructurer(model_path)
|
||||
self.table_matcher = TableMatch()
|
||||
|
||||
def __call__(self, img):
|
||||
s = time.time()
|
||||
dt_boxes, rec_res = self._ocr(copy.deepcopy(img))
|
||||
|
||||
structure_res, _ = self.table_structure(copy.deepcopy(img))
|
||||
pred_html = self.table_matcher(structure_res, dt_boxes, rec_res)
|
||||
|
||||
elapse = time.time() - s
|
||||
return pred_html, elapse
|
||||
|
||||
def _ocr(self, img):
|
||||
h, w = img.shape[:2]
|
||||
|
||||
ocr_result, _ = self.ocr_sys(img)
|
||||
dt_boxes, rec_res, scores = list(zip(*ocr_result))
|
||||
rec_res = list(zip(rec_res, scores))
|
||||
|
||||
r_boxes = []
|
||||
for box in dt_boxes:
|
||||
box = np.array(box)
|
||||
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)
|
||||
return dt_boxes, rec_res
|
||||
|
||||
|
||||
def vis_table(table_res: str, save_path: str) -> None:
|
||||
style_res = '''<style>td {border-left: 1px solid;border-bottom:1px solid;}
|
||||
table, th {border-top:1px solid;font-size: 10px;
|
||||
border-collapse: collapse;border-right: 1px solid;}
|
||||
</style>'''
|
||||
prefix_table, suffix_table = table_res.split('<body>')
|
||||
new_table_res = f'{prefix_table}{style_res}<body>{suffix_table}'
|
||||
with open(save_path, 'w', encoding='utf-8') as f:
|
||||
f.write(new_table_res)
|
||||
print(f'The infer result has saved in {save_path}')
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-v', '--vis', action='store_true',
|
||||
help='Wheter to visualize the layout results.')
|
||||
parser.add_argument('-img', '--img_path', type=str, required=True,
|
||||
help='Path to image for layout.')
|
||||
parser.add_argument('-m', '--model_path', type=str,
|
||||
default=str(root_dir / 'models' / 'en_ppstructure_mobile_v2_SLANet.onnx'),
|
||||
help='The model path used for inference.')
|
||||
args = parser.parse_args()
|
||||
|
||||
rapid_table = RapidTable(args.model_path)
|
||||
|
||||
img = cv2.imread(args.img_path)
|
||||
table_html_str, elapse = rapid_table(img)
|
||||
print(table_html_str)
|
||||
|
||||
if args.vis:
|
||||
img_path = Path(args.img_path)
|
||||
save_path = img_path.resolve().parent / f'vis_{img_path.stem}.html'
|
||||
vis_table(table_html_str, str(save_path))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -1,4 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
from .matcher import TableMatch
|
||||
@@ -1,158 +0,0 @@
|
||||
# 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.
|
||||
# -*- encoding: utf-8 -*-
|
||||
import numpy as np
|
||||
|
||||
from .utils import compute_iou, deal_bb, distance
|
||||
|
||||
|
||||
class TableMatch():
|
||||
def __init__(self, filter_ocr_result=True, 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)
|
||||
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. - 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 += ' '
|
||||
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
|
||||
@@ -1,244 +0,0 @@
|
||||
# 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.
|
||||
# -*- encoding: utf-8 -*-
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import re
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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 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
|
||||
@@ -1,14 +0,0 @@
|
||||
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
from .table_structure import TableStructurer
|
||||
@@ -1,57 +0,0 @@
|
||||
# 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 time
|
||||
import numpy as np
|
||||
|
||||
from .utils import (OrtInferSession, TableLabelDecode, TablePreprocess)
|
||||
|
||||
|
||||
class TableStructurer():
|
||||
def __init__(self, model_path: str):
|
||||
self.preprocess_op = TablePreprocess()
|
||||
|
||||
self.session = OrtInferSession(model_path)
|
||||
|
||||
self.character = self.session.get_metadata()
|
||||
self.postprocess_op = TableLabelDecode(self.character)
|
||||
|
||||
def __call__(self, img):
|
||||
starttime = time.time()
|
||||
data = {'image': img}
|
||||
data = self.preprocess_op(data)
|
||||
img = data[0]
|
||||
if img is None:
|
||||
return None, 0
|
||||
img = np.expand_dims(img, axis=0)
|
||||
img = img.copy()
|
||||
|
||||
outputs = self.session(img)
|
||||
|
||||
preds = {
|
||||
'loc_preds': outputs[0],
|
||||
'structure_probs': outputs[1]
|
||||
}
|
||||
|
||||
shape_list = np.expand_dims(data[-1], axis=0)
|
||||
post_result = self.postprocess_op(preds, [shape_list])
|
||||
|
||||
bbox_list = post_result['bbox_batch_list'][0]
|
||||
|
||||
structure_str_list = post_result['structure_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
|
||||
return (structure_str_list, bbox_list), elapse
|
||||
@@ -1,349 +0,0 @@
|
||||
# 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.
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from onnxruntime import (GraphOptimizationLevel, InferenceSession,
|
||||
SessionOptions)
|
||||
|
||||
|
||||
class OrtInferSession():
|
||||
def __init__(self, onnx_path: str):
|
||||
self._verify_model(onnx_path)
|
||||
|
||||
sess_opt = SessionOptions()
|
||||
sess_opt.log_severity_level = 4
|
||||
sess_opt.enable_cpu_mem_arena = False
|
||||
sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
cpu_ep = 'CPUExecutionProvider'
|
||||
cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested",}
|
||||
EP_list = [(cpu_ep, cpu_provider_options)]
|
||||
|
||||
self._verify_model(onnx_path)
|
||||
self.session = InferenceSession(onnx_path,
|
||||
sess_options=sess_opt,
|
||||
providers=EP_list)
|
||||
|
||||
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
||||
input_dict = dict(zip(self.get_input_names(), [input_content]))
|
||||
try:
|
||||
return self.session.run(self.get_output_names(), input_dict)
|
||||
except Exception as e:
|
||||
raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
|
||||
|
||||
def get_input_names(self, ):
|
||||
return [v.name for v in self.session.get_inputs()]
|
||||
|
||||
def get_output_names(self,):
|
||||
return [v.name for v in self.session.get_outputs()]
|
||||
|
||||
def get_metadata(self, key: str = 'character') -> list:
|
||||
meta_dict = self.session.get_modelmeta().custom_metadata_map
|
||||
content_list = meta_dict[key].splitlines()
|
||||
return content_list
|
||||
|
||||
@staticmethod
|
||||
def _verify_model(model_path):
|
||||
model_path = Path(model_path)
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(f'{model_path} does not exists.')
|
||||
if not model_path.is_file():
|
||||
raise FileExistsError(f'{model_path} is not a file.')
|
||||
|
||||
|
||||
class ONNXRuntimeError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class TableLabelDecode():
|
||||
""" """
|
||||
def __init__(self,
|
||||
dict_character,
|
||||
merge_no_span_structure=True,
|
||||
**kwargs):
|
||||
if merge_no_span_structure:
|
||||
if "<td></td>" not in dict_character:
|
||||
dict_character.append("<td></td>")
|
||||
if "<td>" in dict_character:
|
||||
dict_character.remove("<td>")
|
||||
|
||||
dict_character = self.add_special_char(dict_character)
|
||||
self.dict = {}
|
||||
for i, char in enumerate(dict_character):
|
||||
self.dict[char] = i
|
||||
self.character = dict_character
|
||||
self.td_token = ['<td>', '<td', '<td></td>']
|
||||
|
||||
def __call__(self, preds, batch=None):
|
||||
structure_probs = preds['structure_probs']
|
||||
bbox_preds = preds['loc_preds']
|
||||
shape_list = batch[-1]
|
||||
result = self.decode(structure_probs, bbox_preds, shape_list)
|
||||
if len(batch) == 1: # only contains shape
|
||||
return result
|
||||
|
||||
label_decode_result = self.decode_label(batch)
|
||||
return result, label_decode_result
|
||||
|
||||
def decode(self, structure_probs, bbox_preds, shape_list):
|
||||
"""convert text-label into text-index.
|
||||
"""
|
||||
ignored_tokens = self.get_ignored_tokens()
|
||||
end_idx = self.dict[self.end_str]
|
||||
|
||||
structure_idx = structure_probs.argmax(axis=2)
|
||||
structure_probs = structure_probs.max(axis=2)
|
||||
|
||||
structure_batch_list = []
|
||||
bbox_batch_list = []
|
||||
batch_size = len(structure_idx)
|
||||
for batch_idx in range(batch_size):
|
||||
structure_list = []
|
||||
bbox_list = []
|
||||
score_list = []
|
||||
for idx in range(len(structure_idx[batch_idx])):
|
||||
char_idx = int(structure_idx[batch_idx][idx])
|
||||
if idx > 0 and char_idx == end_idx:
|
||||
break
|
||||
|
||||
if char_idx in ignored_tokens:
|
||||
continue
|
||||
|
||||
text = self.character[char_idx]
|
||||
if text in self.td_token:
|
||||
bbox = bbox_preds[batch_idx, idx]
|
||||
bbox = self._bbox_decode(bbox, shape_list[batch_idx])
|
||||
bbox_list.append(bbox)
|
||||
structure_list.append(text)
|
||||
score_list.append(structure_probs[batch_idx, idx])
|
||||
structure_batch_list.append([structure_list, np.mean(score_list)])
|
||||
bbox_batch_list.append(np.array(bbox_list))
|
||||
result = {
|
||||
'bbox_batch_list': bbox_batch_list,
|
||||
'structure_batch_list': structure_batch_list,
|
||||
}
|
||||
return result
|
||||
|
||||
def decode_label(self, batch):
|
||||
"""convert text-label into text-index.
|
||||
"""
|
||||
structure_idx = batch[1]
|
||||
gt_bbox_list = batch[2]
|
||||
shape_list = batch[-1]
|
||||
ignored_tokens = self.get_ignored_tokens()
|
||||
end_idx = self.dict[self.end_str]
|
||||
|
||||
structure_batch_list = []
|
||||
bbox_batch_list = []
|
||||
batch_size = len(structure_idx)
|
||||
for batch_idx in range(batch_size):
|
||||
structure_list = []
|
||||
bbox_list = []
|
||||
for idx in range(len(structure_idx[batch_idx])):
|
||||
char_idx = int(structure_idx[batch_idx][idx])
|
||||
if idx > 0 and char_idx == end_idx:
|
||||
break
|
||||
if char_idx in ignored_tokens:
|
||||
continue
|
||||
structure_list.append(self.character[char_idx])
|
||||
|
||||
bbox = gt_bbox_list[batch_idx][idx]
|
||||
if bbox.sum() != 0:
|
||||
bbox = self._bbox_decode(bbox, shape_list[batch_idx])
|
||||
bbox_list.append(bbox)
|
||||
structure_batch_list.append(structure_list)
|
||||
bbox_batch_list.append(bbox_list)
|
||||
result = {
|
||||
'bbox_batch_list': bbox_batch_list,
|
||||
'structure_batch_list': structure_batch_list,
|
||||
}
|
||||
return result
|
||||
|
||||
def _bbox_decode(self, bbox, shape):
|
||||
h, w, ratio_h, ratio_w, pad_h, pad_w = shape
|
||||
bbox[0::2] *= w
|
||||
bbox[1::2] *= h
|
||||
return bbox
|
||||
|
||||
def get_ignored_tokens(self):
|
||||
beg_idx = self.get_beg_end_flag_idx("beg")
|
||||
end_idx = self.get_beg_end_flag_idx("end")
|
||||
return [beg_idx, end_idx]
|
||||
|
||||
def get_beg_end_flag_idx(self, beg_or_end):
|
||||
if beg_or_end == "beg":
|
||||
idx = np.array(self.dict[self.beg_str])
|
||||
elif beg_or_end == "end":
|
||||
idx = np.array(self.dict[self.end_str])
|
||||
else:
|
||||
assert False, "unsupport type %s in get_beg_end_flag_idx" \
|
||||
% beg_or_end
|
||||
return idx
|
||||
|
||||
def add_special_char(self, dict_character):
|
||||
self.beg_str = "sos"
|
||||
self.end_str = "eos"
|
||||
dict_character = dict_character
|
||||
dict_character = [self.beg_str] + dict_character + [self.end_str]
|
||||
return dict_character
|
||||
|
||||
|
||||
class TablePreprocess():
|
||||
def __init__(self):
|
||||
self.table_max_len = 488
|
||||
self.build_pre_process_list()
|
||||
self.ops = self.create_operators()
|
||||
|
||||
def __call__(self, data):
|
||||
""" transform """
|
||||
if self.ops is None:
|
||||
self.ops = []
|
||||
|
||||
for op in self.ops:
|
||||
data = op(data)
|
||||
if data is None:
|
||||
return None
|
||||
return data
|
||||
|
||||
def create_operators(self,):
|
||||
"""
|
||||
create operators based on the config
|
||||
|
||||
Args:
|
||||
params(list): a dict list, used to create some operators
|
||||
"""
|
||||
assert isinstance(self.pre_process_list, list), ('operator config should be a list')
|
||||
ops = []
|
||||
for operator in self.pre_process_list:
|
||||
assert isinstance(operator,
|
||||
dict) and len(operator) == 1, "yaml format error"
|
||||
op_name = list(operator)[0]
|
||||
param = {} if operator[op_name] is None else operator[op_name]
|
||||
op = eval(op_name)(**param)
|
||||
ops.append(op)
|
||||
return ops
|
||||
|
||||
def build_pre_process_list(self):
|
||||
resize_op = {'ResizeTableImage': {'max_len': self.table_max_len, }}
|
||||
pad_op = {
|
||||
'PaddingTableImage': {
|
||||
'size': [self.table_max_len, self.table_max_len]
|
||||
}
|
||||
}
|
||||
normalize_op = {
|
||||
'NormalizeImage': {
|
||||
'std': [0.229, 0.224, 0.225],
|
||||
'mean': [0.485, 0.456, 0.406],
|
||||
'scale': '1./255.',
|
||||
'order': 'hwc'
|
||||
}
|
||||
}
|
||||
to_chw_op = {'ToCHWImage': None}
|
||||
keep_keys_op = {'KeepKeys': {'keep_keys': ['image', 'shape']}}
|
||||
self.pre_process_list = [
|
||||
resize_op, normalize_op, pad_op, to_chw_op, keep_keys_op
|
||||
]
|
||||
|
||||
|
||||
class ResizeTableImage():
|
||||
def __init__(self, max_len, resize_bboxes=False, infer_mode=False,
|
||||
**kwargs):
|
||||
super(ResizeTableImage, self).__init__()
|
||||
self.max_len = max_len
|
||||
self.resize_bboxes = resize_bboxes
|
||||
self.infer_mode = infer_mode
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
height, width = img.shape[0:2]
|
||||
ratio = self.max_len / (max(height, width) * 1.0)
|
||||
resize_h = int(height * ratio)
|
||||
resize_w = int(width * ratio)
|
||||
resize_img = cv2.resize(img, (resize_w, resize_h))
|
||||
if self.resize_bboxes and not self.infer_mode:
|
||||
data['bboxes'] = data['bboxes'] * ratio
|
||||
data['image'] = resize_img
|
||||
data['src_img'] = img
|
||||
data['shape'] = np.array([height, width, ratio, ratio])
|
||||
data['max_len'] = self.max_len
|
||||
return data
|
||||
|
||||
|
||||
class PaddingTableImage():
|
||||
def __init__(self, size, **kwargs):
|
||||
super(PaddingTableImage, self).__init__()
|
||||
self.size = size
|
||||
|
||||
def __call__(self, data):
|
||||
img = data['image']
|
||||
pad_h, pad_w = self.size
|
||||
padding_img = np.zeros((pad_h, pad_w, 3), dtype=np.float32)
|
||||
height, width = img.shape[0:2]
|
||||
padding_img[0:height, 0:width, :] = img.copy()
|
||||
data['image'] = padding_img
|
||||
shape = data['shape'].tolist()
|
||||
shape.extend([pad_h, pad_w])
|
||||
data['shape'] = np.array(shape)
|
||||
return data
|
||||
|
||||
|
||||
class NormalizeImage():
|
||||
""" normalize image such as substract mean, divide std
|
||||
"""
|
||||
def __init__(self, scale=None, mean=None, std=None, order='chw', **kwargs):
|
||||
if isinstance(scale, str):
|
||||
scale = eval(scale)
|
||||
self.scale = np.float32(scale if scale is not None else 1.0 / 255.0)
|
||||
mean = mean if mean is not None else [0.485, 0.456, 0.406]
|
||||
std = std if std is not None else [0.229, 0.224, 0.225]
|
||||
|
||||
shape = (3, 1, 1) if order == 'chw' else (1, 1, 3)
|
||||
self.mean = np.array(mean).reshape(shape).astype('float32')
|
||||
self.std = np.array(std).reshape(shape).astype('float32')
|
||||
|
||||
def __call__(self, data):
|
||||
img = np.array(data['image'])
|
||||
assert isinstance(img,
|
||||
np.ndarray), "invalid input 'img' in NormalizeImage"
|
||||
data['image'] = (
|
||||
img.astype('float32') * self.scale - self.mean) / self.std
|
||||
return data
|
||||
|
||||
|
||||
class ToCHWImage():
|
||||
""" convert hwc image to chw image
|
||||
"""
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
def __call__(self, data):
|
||||
img = np.array(data['image'])
|
||||
data['image'] = img.transpose((2, 0, 1))
|
||||
return data
|
||||
|
||||
|
||||
class KeepKeys():
|
||||
def __init__(self, keep_keys, **kwargs):
|
||||
self.keep_keys = keep_keys
|
||||
|
||||
def __call__(self, data):
|
||||
data_list = []
|
||||
for key in self.keep_keys:
|
||||
data_list.append(data[key])
|
||||
return data_list
|
||||
@@ -1,8 +0,0 @@
|
||||
numpy>=1.21.6
|
||||
opencv_python>=4.6.0.66
|
||||
setuptools>=61.2.0
|
||||
PyYAML
|
||||
pytest>=7.2.1
|
||||
pip>=22.3.1
|
||||
onnxruntime
|
||||
wheel
|
||||
@@ -1,87 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import setuptools
|
||||
|
||||
|
||||
def get_latest_version(package_name):
|
||||
output = subprocess.run(["pip", "index", "versions", package_name],
|
||||
capture_output=True)
|
||||
output = output.stdout.decode('utf-8')
|
||||
if output:
|
||||
return extract_version(output)
|
||||
return None
|
||||
|
||||
|
||||
def version_add_one(version, add_loc=-1):
|
||||
if version:
|
||||
version_list = version.split('.')
|
||||
mini_version = str(int(version_list[add_loc]) + 1)
|
||||
version_list[add_loc] = mini_version
|
||||
new_version = '.'.join(version_list)
|
||||
return new_version
|
||||
return '0.0.1'
|
||||
|
||||
|
||||
def get_readme():
|
||||
root_dir = Path(__file__).resolve().parent.parent.parent
|
||||
readme_path = str(root_dir / 'docs' / 'doc_whl_rapid_layout.md')
|
||||
with open(readme_path, 'r', encoding='utf-8') as f:
|
||||
readme = f.read()
|
||||
return readme
|
||||
|
||||
|
||||
def extract_version(message: str) -> str:
|
||||
pattern = r'\d+\.(?:\d+\.)*\d+'
|
||||
matched_versions = re.findall(pattern, message)
|
||||
if matched_versions:
|
||||
return matched_versions[0]
|
||||
return ''
|
||||
|
||||
|
||||
MODULE_NAME = 'rapid_layout'
|
||||
latest_version = get_latest_version(MODULE_NAME)
|
||||
VERSION_NUM = version_add_one(latest_version)
|
||||
|
||||
# 优先提取commit message中的语义化版本号，如无，则自动加1
|
||||
if len(sys.argv) > 2:
|
||||
match_str = ' '.join(sys.argv[2:])
|
||||
matched_versions = extract_version(match_str)
|
||||
if matched_versions:
|
||||
VERSION_NUM = matched_versions
|
||||
sys.argv = sys.argv[:2]
|
||||
|
||||
setuptools.setup(
|
||||
name=MODULE_NAME,
|
||||
version=VERSION_NUM,
|
||||
platforms="Any",
|
||||
long_description=get_readme(),
|
||||
long_description_content_type='text/markdown',
|
||||
description='Tools for document layout analysis based ONNXRuntime.',
|
||||
author="SWHL",
|
||||
author_email="liekkaskono@163.com",
|
||||
url="https://github.com/RapidAI/RapidOCR",
|
||||
license='Apache-2.0',
|
||||
include_package_data=True,
|
||||
install_requires=["onnxruntime>=1.7.0", "PyYAML>=6.0",
|
||||
"opencv_python>=4.5.1.48", "numpy>=1.21.6"],
|
||||
packages=[MODULE_NAME, f'{MODULE_NAME}.models'],
|
||||
package_data={'': ['layout_cdla.onnx', '*.yaml']},
|
||||
keywords=[
|
||||
'ppstructure,layout,rapidocr,rapid_layout'
|
||||
],
|
||||
classifiers=[
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
],
|
||||
entry_points={
|
||||
'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.{MODULE_NAME}:main']
|
||||
}
|
||||
)
|
||||
@@ -1,87 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import setuptools
|
||||
|
||||
|
||||
def get_latest_version(package_name):
|
||||
output = subprocess.run(["pip", "index", "versions", package_name],
|
||||
capture_output=True)
|
||||
output = output.stdout.decode('utf-8')
|
||||
if output:
|
||||
return extract_version(output)
|
||||
return None
|
||||
|
||||
|
||||
def version_add_one(version, add_loc=-1):
|
||||
if version:
|
||||
version_list = version.split('.')
|
||||
mini_version = str(int(version_list[add_loc]) + 1)
|
||||
version_list[add_loc] = mini_version
|
||||
new_version = '.'.join(version_list)
|
||||
return new_version
|
||||
return '0.0.1'
|
||||
|
||||
|
||||
def get_readme():
|
||||
root_dir = Path(__file__).resolve().parent.parent.parent
|
||||
readme_path = str(root_dir / 'docs' / 'doc_whl_rapid_orientation.md')
|
||||
with open(readme_path, 'r', encoding='utf-8') as f:
|
||||
readme = f.read()
|
||||
return readme
|
||||
|
||||
|
||||
def extract_version(message: str) -> str:
|
||||
pattern = r'\d+\.(?:\d+\.)*\d+'
|
||||
matched_versions = re.findall(pattern, message)
|
||||
if matched_versions:
|
||||
return matched_versions[0]
|
||||
return ''
|
||||
|
||||
|
||||
MODULE_NAME = 'rapid_orientation'
|
||||
latest_version = get_latest_version(MODULE_NAME)
|
||||
VERSION_NUM = version_add_one(latest_version)
|
||||
|
||||
# 优先提取commit message中的语义化版本号，如无，则自动加1
|
||||
if len(sys.argv) > 2:
|
||||
match_str = ' '.join(sys.argv[2:])
|
||||
matched_versions = extract_version(match_str)
|
||||
if matched_versions:
|
||||
VERSION_NUM = matched_versions
|
||||
sys.argv = sys.argv[:2]
|
||||
|
||||
setuptools.setup(
|
||||
name=MODULE_NAME,
|
||||
version=VERSION_NUM,
|
||||
platforms="Any",
|
||||
long_description=get_readme(),
|
||||
long_description_content_type='text/markdown',
|
||||
description='Tools for classifying the direction of images containing text based ONNXRuntime.',
|
||||
author="SWHL",
|
||||
author_email="liekkaskono@163.com",
|
||||
url="https://github.com/RapidAI/RapidOCR",
|
||||
license='Apache-2.0',
|
||||
include_package_data=True,
|
||||
install_requires=["onnxruntime>=1.7.0", "PyYAML>=6.0",
|
||||
"opencv_python>=4.5.1.48", "numpy>=1.21.6"],
|
||||
packages=[MODULE_NAME, f'{MODULE_NAME}.models'],
|
||||
package_data={'': ['*.onnx', '*.yaml']},
|
||||
keywords=[
|
||||
'ppstructure,layout,rapidocr,rapid_orientation'
|
||||
],
|
||||
classifiers=[
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
],
|
||||
entry_points={
|
||||
'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.{MODULE_NAME}:main']
|
||||
}
|
||||
)
|
||||
@@ -1,91 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import setuptools
|
||||
|
||||
|
||||
def get_latest_version(package_name):
|
||||
output = subprocess.run(["pip", "index", "versions", package_name],
|
||||
capture_output=True)
|
||||
output = output.stdout.decode('utf-8')
|
||||
if output:
|
||||
return extract_version(output)
|
||||
return None
|
||||
|
||||
|
||||
def version_add_one(version, add_loc=-1):
|
||||
if version:
|
||||
version_list = version.split('.')
|
||||
mini_version = str(int(version_list[add_loc]) + 1)
|
||||
version_list[add_loc] = mini_version
|
||||
new_version = '.'.join(version_list)
|
||||
return new_version
|
||||
return '0.0.1'
|
||||
|
||||
|
||||
def get_readme():
|
||||
root_dir = Path(__file__).resolve().parent.parent.parent
|
||||
readme_path = str(root_dir / 'docs' / 'doc_whl_rapid_table.md')
|
||||
with open(readme_path, 'r', encoding='utf-8') as f:
|
||||
readme = f.read()
|
||||
return readme
|
||||
|
||||
|
||||
def extract_version(message: str) -> str:
|
||||
pattern = r'\d+\.(?:\d+\.)*\d+'
|
||||
matched_versions = re.findall(pattern, message)
|
||||
if matched_versions:
|
||||
return matched_versions[0]
|
||||
return ''
|
||||
|
||||
|
||||
MODULE_NAME = 'rapid_table'
|
||||
latest_version = get_latest_version(MODULE_NAME)
|
||||
VERSION_NUM = version_add_one(latest_version)
|
||||
|
||||
# 优先提取commit message中的语义化版本号，如无，则自动加1
|
||||
if len(sys.argv) > 2:
|
||||
match_str = ' '.join(sys.argv[2:])
|
||||
matched_versions = extract_version(match_str)
|
||||
if matched_versions:
|
||||
VERSION_NUM = matched_versions
|
||||
sys.argv = sys.argv[:2]
|
||||
|
||||
setuptools.setup(
|
||||
name=MODULE_NAME,
|
||||
version=VERSION_NUM,
|
||||
platforms="Any",
|
||||
long_description=get_readme(),
|
||||
long_description_content_type='text/markdown',
|
||||
description='Tools for parsing table structures based ONNXRuntime.',
|
||||
author="SWHL",
|
||||
author_email="liekkaskono@163.com",
|
||||
url="https://github.com/RapidAI/RapidOCR",
|
||||
license='Apache-2.0',
|
||||
include_package_data=True,
|
||||
install_requires=["onnxruntime>=1.7.0", "PyYAML>=6.0",
|
||||
"opencv_python>=4.5.1.48", "numpy>=1.21.6",
|
||||
"rapidocr_onnxruntime>=1.1.18"],
|
||||
packages=[MODULE_NAME,
|
||||
f'{MODULE_NAME}.models',
|
||||
f'{MODULE_NAME}.table_matcher',
|
||||
f'{MODULE_NAME}.table_structure'],
|
||||
package_data={'': ['en_ppstructure_mobile_v2_SLANet.onnx']},
|
||||
keywords=[
|
||||
'ppstructure,table,rapidocr,rapid_table'
|
||||
],
|
||||
classifiers=[
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
],
|
||||
entry_points={
|
||||
'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.{MODULE_NAME}:main']
|
||||
}
|
||||
)
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 164 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 58 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 79 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 92 KiB |
@@ -1,24 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
|
||||
root_dir = Path(__file__).resolve().parent.parent
|
||||
|
||||
sys.path.append(str(root_dir))
|
||||
|
||||
from rapid_layout import RapidLayout
|
||||
|
||||
|
||||
def test_layout():
|
||||
img_path = str(root_dir / 'test_images' / 'layout.png')
|
||||
img = cv2.imread(img_path)
|
||||
|
||||
layout_engine = RapidLayout()
|
||||
|
||||
layout_res, elapse = layout_engine(img)
|
||||
|
||||
assert len(layout_res) == 13
|
||||
@@ -1,28 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import pytest
|
||||
|
||||
root_dir = Path(__file__).resolve().parent
|
||||
sys.path.append(str(root_dir.parent))
|
||||
|
||||
from rapid_orientation import RapidOrientation
|
||||
text_orientation = RapidOrientation()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'img_path, result',
|
||||
[
|
||||
(f'{root_dir}/test_files/img_rot0_demo.jpg', '0'),
|
||||
(f'{root_dir}/test_files/img_rot180_demo.jpg', '180'),
|
||||
]
|
||||
)
|
||||
def test_img(img_path, result):
|
||||
img = cv2.imread(str(img_path))
|
||||
pred_result, _ = text_orientation(img)
|
||||
|
||||
assert pred_result == result
|
||||
@@ -1,23 +0,0 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
# @Author: SWHL
|
||||
# @Contact: liekkaskono@163.com
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
|
||||
root_dir = Path(__file__).resolve().parent.parent
|
||||
|
||||
sys.path.append(str(root_dir))
|
||||
|
||||
from rapid_table import RapidTable
|
||||
|
||||
|
||||
def test_table():
|
||||
img_path = str(root_dir / 'test_images' / 'table.jpg')
|
||||
img = cv2.imread(img_path)
|
||||
|
||||
rapid_table = RapidTable()
|
||||
|
||||
table_html_str, elapse = rapid_table(img)
|
||||
assert table_html_str.count('<tr>') == 16
|
||||
Reference in new issue
Block a user