chore: Optim code

Author: SWHL

.gitignore

@@ -157,4 +157,4 @@ long1.jpg
 
 .DS_Store
 *.npy
-/lineless_table_rec/output/
+outputs/

demo_lineless.py

@@ -1,19 +1,29 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-from pathlib import Path
+import os
 
 from lineless_table_rec import LinelessTableRecognition
+from lineless_table_rec.utils_table_recover import (
+    format_html,
+    plot_rec_box,
+    plot_rec_box_with_logic_info,
+)
 
-engine = LinelessTableRecognition()
-
+output_dir = "outputs"
 img_path = "tests/test_files/lineless_table_recognition.jpg"
-table_str, elapse = engine(img_path)
+table_rec = LinelessTableRecognition()
+
+html, elasp, polygons, logic_points, ocr_res = table_rec(img_path)
+print(f"cost: {elasp:.5f}")
 
-print(table_str)
-print(elapse)
+complete_html = format_html(html)
+os.makedirs(os.path.dirname(f"{output_dir}/table.html"), exist_ok=True)
 
-with open(f"{Path(img_path).stem}.html", "w", encoding="utf-8") as f:
-    f.write(table_str)
+with open(f"{output_dir}/table.html", "w", encoding="utf-8") as file:
+    file.write(complete_html)
 
-print("ok")
+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)

lineless_table_rec/main.py

@@ -1,9 +1,7 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-import argparse
 import logging
-import os
 import time
 import traceback
 from pathlib import Path
@@ -13,13 +11,16 @@
 import numpy as np
 from rapidocr_onnxruntime import RapidOCR
 
-from lineless_table_process import DetProcess, get_affine_transform_upper_left
-from utils import InputType, LoadImage, OrtInferSession
-from utils_table_recover import (
+from .process import DetProcess, get_affine_transform_upper_left
+from .utils import InputType, LoadImage, OrtInferSession
+from .utils_table_recover import (
+    box_4_2_poly_to_box_4_1,
+    filter_duplicated_box,
+    gather_ocr_list_by_row,
     get_rotate_crop_image,
+    match_ocr_cell,
     plot_html_table,
-    sorted_ocr_boxes, box_4_2_poly_to_box_4_1, match_ocr_cell,
-    filter_duplicated_box, gather_ocr_list_by_row, plot_rec_box_with_logic_info, plot_rec_box, format_html,
+    sorted_ocr_boxes,
 )
 
 cur_dir = Path(__file__).resolve().parent
@@ -29,9 +30,9 @@
 
 class LinelessTableRecognition:
     def __init__(
-            self,
-            detect_model_path: Union[str, Path] = detect_model_path,
-            process_model_path: Union[str, Path] = process_model_path,
+        self,
+        detect_model_path: Union[str, Path] = detect_model_path,
+        process_model_path: Union[str, Path] = process_model_path,
     ):
         self.mean = np.array([0.408, 0.447, 0.470], dtype=np.float32).reshape(1, 1, 3)
         self.std = np.array([0.289, 0.274, 0.278], dtype=np.float32).reshape(1, 1, 3)
@@ -61,36 +62,56 @@ def __call__(self, content: InputType):
             # 转换为中间格式，修正识别框坐标,将物理识别框，逻辑识别框，ocr识别框整合为dict，方便后续处理
             t_rec_ocr_list = self.transform_res(cell_box_det_map, polygons, logi_points)
             # 拆分包含和重叠的识别框
-            deleted_idx_set = filter_duplicated_box([table_box_ocr['t_box'] for table_box_ocr in t_rec_ocr_list])
-            t_rec_ocr_list = [t_rec_ocr_list[i] for i in range(len(t_rec_ocr_list)) if i not in deleted_idx_set]
+            deleted_idx_set = filter_duplicated_box(
+                [table_box_ocr["t_box"] for table_box_ocr in t_rec_ocr_list]
+            )
+            t_rec_ocr_list = [
+                t_rec_ocr_list[i]
+                for i in range(len(t_rec_ocr_list))
+                if i not in deleted_idx_set
+            ]
             # 生成行列对应的二维表格, 合并同行同列识别框中的的ocr识别框
             t_rec_ocr_list, grid = self.handle_overlap_row_col(t_rec_ocr_list)
             # todo 根据grid 及 not_match_orc_boxes，尝试将ocr识别填入单行单列中
             # 将同一个识别框中的ocr结果排序并同行合并
             t_rec_ocr_list = self.sort_and_gather_ocr_res(t_rec_ocr_list)
             # 渲染为html
-            logi_points = [t_box_ocr['t_logic_box'] for t_box_ocr in t_rec_ocr_list]
+            logi_points = [t_box_ocr["t_logic_box"] for t_box_ocr in t_rec_ocr_list]
             cell_box_det_map = {
-                i: [ocr_box_and_text[1] for ocr_box_and_text in t_box_ocr['t_ocr_res']]
+                i: [ocr_box_and_text[1] for ocr_box_and_text in t_box_ocr["t_ocr_res"]]
                 for i, t_box_ocr in enumerate(t_rec_ocr_list)
             }
             table_str = plot_html_table(logi_points, cell_box_det_map)
 
             # 输出可视化排序,用于验证结果，生产版本可以去掉
-            _, idx_list = sorted_ocr_boxes([t_box_ocr['t_box'] for t_box_ocr in t_rec_ocr_list])
+            _, idx_list = sorted_ocr_boxes(
+                [t_box_ocr["t_box"] for t_box_ocr in t_rec_ocr_list]
+            )
             t_rec_ocr_list = [t_rec_ocr_list[i] for i in idx_list]
-            sorted_polygons = [t_box_ocr['t_box'] for t_box_ocr in t_rec_ocr_list]
-            sorted_logi_points = [t_box_ocr['t_logic_box'] for t_box_ocr in t_rec_ocr_list]
+            sorted_polygons = [t_box_ocr["t_box"] for t_box_ocr in t_rec_ocr_list]
+            sorted_logi_points = [
+                t_box_ocr["t_logic_box"] for t_box_ocr in t_rec_ocr_list
+            ]
             ocr_boxes_res = [box_4_2_poly_to_box_4_1(ori_ocr[0]) for ori_ocr in ocr_res]
             sorted_ocr_boxes_res, _ = sorted_ocr_boxes(ocr_boxes_res)
             table_elapse = time.perf_counter() - ss
-            return table_str, table_elapse, sorted_polygons, sorted_logi_points, sorted_ocr_boxes_res
+            return (
+                table_str,
+                table_elapse,
+                sorted_polygons,
+                sorted_logi_points,
+                sorted_ocr_boxes_res,
+            )
         except Exception:
             logging.warning(traceback.format_exc())
             return "", 0.0, None, None, None
 
-    def transform_res(self, cell_box_det_map: dict[int, List[any]], polygons: np.ndarray,
-                      logi_points: list[np.ndarray]) -> list[dict[str, any]]:
+    def transform_res(
+        self,
+        cell_box_det_map: dict[int, List[any]],
+        polygons: np.ndarray,
+        logi_points: list[np.ndarray],
+    ) -> list[dict[str, any]]:
         res = []
         for i in range(len(polygons)):
             ocr_res_list = cell_box_det_map.get(i)
@@ -102,11 +123,14 @@ def transform_res(self, cell_box_det_map: dict[int, List[any]], polygons: np.nda
             ymax = max([ocr_box[0][2][1] for ocr_box in ocr_res_list])
             dict_res = {
                 # xmin,xmax,ymin,ymax
-                't_box': [xmin, ymin, xmax, ymax],
+                "t_box": [xmin, ymin, xmax, ymax],
                 # row_start,row_end,col_start,col_end
-                't_logic_box': logi_points[i].tolist(),
+                "t_logic_box": logi_points[i].tolist(),
                 # [[xmin,xmax,ymin,ymax], text]
-                't_ocr_res': [[box_4_2_poly_to_box_4_1(ocr_det[0]), ocr_det[1]] for ocr_det in ocr_res_list]
+                "t_ocr_res": [
+                    [box_4_2_poly_to_box_4_1(ocr_det[0]), ocr_det[1]]
+                    for ocr_det in ocr_res_list
+                ],
             }
             res.append(dict_res)
         return res
@@ -156,24 +180,30 @@ def infer(self, input_content: Dict[str, Any]) -> Tuple[np.ndarray, np.ndarray]:
 
     def sort_and_gather_ocr_res(self, res):
         for i, dict_res in enumerate(res):
-            dict_res['t_ocr_res'] = gather_ocr_list_by_row(dict_res['t_ocr_res'])
-            _, sorted_idx = sorted_ocr_boxes([ocr_det[0] for ocr_det in dict_res['t_ocr_res']])
-            dict_res['t_ocr_res'] = [dict_res['t_ocr_res'][i] for i in sorted_idx]
+            dict_res["t_ocr_res"] = gather_ocr_list_by_row(dict_res["t_ocr_res"])
+            _, sorted_idx = sorted_ocr_boxes(
+                [ocr_det[0] for ocr_det in dict_res["t_ocr_res"]]
+            )
+            dict_res["t_ocr_res"] = [dict_res["t_ocr_res"][i] for i in sorted_idx]
         return res
 
     def handle_overlap_row_col(self, res):
         max_row, max_col = 0, 0
         for dict_res in res:
-            max_row = max(max_row, dict_res['t_logic_box'][1] + 1)  # 加1是因为结束下标是包含在内的
-            max_col = max(max_col, dict_res['t_logic_box'][3] + 1)  # 加1是因为结束下标是包含在内的
+            max_row = max(
+                max_row, dict_res["t_logic_box"][1] + 1
+            )  # 加1是因为结束下标是包含在内的
+            max_col = max(
+                max_col, dict_res["t_logic_box"][3] + 1
+            )  # 加1是因为结束下标是包含在内的
         # 创建一个二维数组来存储 sorted_logi_points 中的元素
         grid = [[None] * max_col for _ in range(max_row)]
         # 将 sorted_logi_points 中的元素填充到 grid 中
         deleted_idx = set()
         for i, dict_res in enumerate(res):
             if i in deleted_idx:
                 continue
-            row_start, row_end, col_start, col_end = dict_res['t_logic_box']
+            row_start, row_end, col_start, col_end = dict_res["t_logic_box"]
             for row in range(row_start, row_end + 1):
                 if i in deleted_idx:
                     continue
@@ -184,15 +214,16 @@ def handle_overlap_row_col(self, res):
                     if not exist_dict_res:
                         grid[row][col] = dict_res
                         continue
-                    if exist_dict_res['t_logic_box'] == dict_res['t_logic_box']:
-                        exist_dict_res['t_ocr_res'].extend(dict_res['t_ocr_res'])
+                    if exist_dict_res["t_logic_box"] == dict_res["t_logic_box"]:
+                        exist_dict_res["t_ocr_res"].extend(dict_res["t_ocr_res"])
                         deleted_idx.add(i)
                         # 修正识别框坐标
-                        exist_dict_res['t_box'] = [min(exist_dict_res['t_box'][0], dict_res['t_box'][0]),
-                                                   min(exist_dict_res['t_box'][1], dict_res['t_box'][1]),
-                                                   max(exist_dict_res['t_box'][2], dict_res['t_box'][2]),
-                                                   max(exist_dict_res['t_box'][3], dict_res['t_box'][3]),
-                                                   ]
+                        exist_dict_res["t_box"] = [
+                            min(exist_dict_res["t_box"][0], dict_res["t_box"][0]),
+                            min(exist_dict_res["t_box"][1], dict_res["t_box"][1]),
+                            max(exist_dict_res["t_box"][2], dict_res["t_box"][2]),
+                            max(exist_dict_res["t_box"][3], dict_res["t_box"][3]),
+                        ]
                         continue
 
         #  去掉重叠框
@@ -217,10 +248,10 @@ def filter_logi_points(slct_logi: np.ndarray) -> list[np.ndarray]:
         return slct_logi[0].astype(np.int32)
 
     def re_rec(
-            self,
-            img: np.ndarray,
-            sorted_polygons: np.ndarray,
-            cell_box_map: Dict[int, List[str]],
+        self,
+        img: np.ndarray,
+        sorted_polygons: np.ndarray,
+        cell_box_map: Dict[int, List[str]],
     ) -> Dict[int, List[any]]:
         """找到poly对应为空的框，尝试将直接将poly框直接送到识别中"""
         #
@@ -237,25 +268,3 @@ def re_rec(
             scores = [rec[1] for rec in rec_res]
             cell_box_map[i] = [[box, "".join(text), min(scores)]]
         return cell_box_map
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-img", "--img_path", type=str, required=True)
-    parser.add_argument( "--output_dir", default= "./output", type=str)
-    args = parser.parse_args()
-    # args.img_path = '../images/image (78).png'
-    table_rec = LinelessTableRecognition()
-    html, elasp, polygons, logic_points, ocr_res = table_rec(args.img_path)
-    print(f"cost: {elasp:.5f}")
-    complete_html = format_html(html)
-    os.makedirs(os.path.dirname(f'{args.output_dir}/table.html'), exist_ok=True)
-    with open(f'{args.output_dir}/table.html', 'w', encoding='utf-8') as file:
-        file.write(complete_html)
-    plot_rec_box_with_logic_info(args.img_path, f'{args.output_dir}/table_rec_box.jpg', logic_points, polygons)
-    plot_rec_box(args.img_path, f'{args.output_dir}/ocr_box.jpg', ocr_res)
-
-
-
-if __name__ == "__main__":
-    main()