-
-
Notifications
You must be signed in to change notification settings - Fork 11
/
demo.py
101 lines (73 loc) · 2.97 KB
/
demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
# -*- 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, VisTable
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():
from rapidocr_onnxruntime import RapidOCR
table_engine = RapidTable()
img = cv2.imread("tests/test_files/table.jpg")
table_html_str, _ = table_engine(img)
ocr_engine = RapidOCR()
viser = VisTable()
img_path = "tests/test_files/table.jpg"
ocr_result, _ = ocr_engine(img_path)
table_html_str, table_cell_bboxes, _ = table_engine(img_path, ocr_result)
save_dir = Path("./inference_results/")
if not save_dir.exists():
save_dir.mkdir(parents=True, exist_ok=True)
save_html_path = save_dir / f"{Path(img_path).stem}.html"
save_drawed_path = save_dir / f"vis_{Path(img_path).name}"
viser(img_path, table_html_str, save_html_path, table_cell_bboxes, save_drawed_path)
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()