From cae1eefc68dfdd2777cdaa9ae20f533adbc6ce3b Mon Sep 17 00:00:00 2001 From: SWHL Date: Thu, 19 Sep 2024 13:01:23 +0800 Subject: [PATCH] feat(wired_table_rec): Add format_html --- demo_wired.py | 18 ++++++++++----- wired_table_rec/utils_table_recover.py | 31 +++++++++++++++++++++++++- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/demo_wired.py b/demo_wired.py index db2f23d..eb99145 100644 --- a/demo_wired.py +++ b/demo_wired.py @@ -1,26 +1,34 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -import os -from lineless_table_rec.utils_table_recover import format_html +from pathlib import Path + from wired_table_rec import WiredTableRecognition from wired_table_rec.utils_table_recover import ( + format_html, plot_rec_box, plot_rec_box_with_logic_info, ) -output_dir = "outputs" +output_dir = Path("outputs") +output_dir.mkdir(parents=True, exist_ok=True) + table_rec = WiredTableRecognition() img_path = "tests/test_files/wired/table1.png" html, elasp, polygons, logic_points, ocr_res = table_rec(img_path) + print(f"cost: {elasp:.5f}") + complete_html = format_html(html) -os.makedirs(os.path.dirname(f"{output_dir}/table.html"), exist_ok=True) -with open(f"{output_dir}/table.html", "w", encoding="utf-8") as file: + +save_table_path = output_dir / "table.html" +with open(save_table_path, "w", encoding="utf-8") as file: file.write(complete_html) plot_rec_box_with_logic_info( img_path, f"{output_dir}/table_rec_box.jpg", logic_points, polygons ) plot_rec_box(img_path, f"{output_dir}/ocr_box.jpg", ocr_res) + +print(f"The results has been saved under {output_dir}") diff --git a/wired_table_rec/utils_table_recover.py b/wired_table_rec/utils_table_recover.py index ca2cae5..4af788c 100644 --- a/wired_table_rec/utils_table_recover.py +++ b/wired_table_rec/utils_table_recover.py @@ -3,7 +3,7 @@ # @Contact: liekkaskono@163.com import os import random -from typing import Dict, List, Union, Any +from typing import Any, Dict, List, Union import cv2 import numpy as np @@ -638,3 +638,32 @@ def vis_table(img: np.ndarray, polygons: np.ndarray) -> np.ndarray: font = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(img, str(i), poly[0], font, 1, (0, 0, 255), 1) return img + + +def format_html(html): + return f""" + + + + + Complex Table Example + + + + {html} + + + """