fix: fix wired unet model rec

Author: Joker1212

demo_all.py

@@ -0,0 +1,41 @@
+from table_cls import TableCls
+from wired_table_rec.main import WiredTableInput, WiredTableRecognition
+from lineless_table_rec.main import LinelessTableInput, LinelessTableRecognition
+
+if __name__ == "__main__":
+    # Init
+    wired_input = WiredTableInput()
+    lineless_input = LinelessTableInput()
+    wired_engine = WiredTableRecognition(wired_input)
+    lineless_engine = LinelessTableRecognition(lineless_input)
+    # 默认小yolo模型(0.1s)，可切换为精度更高yolox(0.25s),更快的qanything(0.07s)模型或paddle模型(0.03s)
+    table_cls = TableCls()
+    img_path = f"tests/test_files/table.jpg"
+
+    cls, elasp = table_cls(img_path)
+    if cls == "wired":
+        table_engine = wired_engine
+    else:
+        table_engine = lineless_engine
+
+    table_results = table_engine(img_path, enhance_box_line=False)
+    # 使用RapidOCR输入
+    # ocr_engine = RapidOCR()
+    # ocr_result, _ = ocr_engine(img_path)
+    # table_results = table_engine(img_path, ocr_result=ocr_result)
+
+    # Visualize table rec result
+    # save_dir = Path("outputs")
+    # save_dir.mkdir(parents=True, exist_ok=True)
+    #
+    # save_html_path = f"outputs/{Path(img_path).stem}.html"
+    # save_drawed_path = f"outputs/{Path(img_path).stem}_table_vis{Path(img_path).suffix}"
+    # save_logic_path = (
+    #     f"outputs/{Path(img_path).stem}_table_vis_logic{Path(img_path).suffix}"
+    # )
+
+    #
+    # vis_table = VisTable()
+    # vis_imged = vis_table(
+    #     img_path, table_results, save_html_path, save_drawed_path, save_logic_path
+    # )

demo_lineless.py

@@ -6,12 +6,12 @@
 from rapidocr_onnxruntime import RapidOCR
 
 from lineless_table_rec import LinelessTableRecognition
-from lineless_table_rec.main import RapidTableInput
+from lineless_table_rec.main import LinelessTableInput
 from lineless_table_rec.utils.utils import VisTable
 
 output_dir = Path("outputs")
 output_dir.mkdir(parents=True, exist_ok=True)
-input_args = RapidTableInput()
+input_args = LinelessTableInput()
 table_engine = LinelessTableRecognition(input_args)
 ocr_engine = RapidOCR()
 viser = VisTable()

demo_wired.py

@@ -6,12 +6,12 @@
 from rapidocr_onnxruntime import RapidOCR
 
 from wired_table_rec import WiredTableRecognition
-from wired_table_rec.main import RapidTableInput, ModelType
+from wired_table_rec.main import WiredTableInput
 from wired_table_rec.utils.utils import VisTable
 
 output_dir = Path("outputs")
 output_dir.mkdir(parents=True, exist_ok=True)
-input_args = RapidTableInput(model_type=ModelType.CYCLE_CENTER_NET.value)
+input_args = WiredTableInput()
 table_engine = WiredTableRecognition(input_args)
 ocr_engine = RapidOCR()
 viser = VisTable()

lineless_table_rec/main.py

@@ -24,6 +24,7 @@
     match_ocr_cell,
     plot_html_table,
     sorted_ocr_boxes,
+    box_4_1_poly_to_box_4_2,
 )
 
 
@@ -41,23 +42,23 @@ class ModelType(Enum):
 
 
 @dataclass
-class RapidTableInput:
+class LinelessTableInput:
     model_type: Optional[str] = ModelType.LORE.value
     model_path: Union[str, Path, None, Dict[str, str]] = None
     use_cuda: bool = False
     device: str = "cpu"
 
 
 @dataclass
-class RapidTableOutput:
+class LinelessTableOutput:
     pred_html: Optional[str] = None
     cell_bboxes: Optional[np.ndarray] = None
     logic_points: Optional[np.ndarray] = None
     elapse: Optional[float] = None
 
 
 class LinelessTableRecognition:
-    def __init__(self, config: RapidTableInput):
+    def __init__(self, config: LinelessTableInput):
         self.model_type = config.model_type
         if self.model_type not in KEY_TO_MODEL_URL:
             model_list = ",".join(KEY_TO_MODEL_URL)
@@ -78,7 +79,7 @@ def __call__(
         content: InputType,
         ocr_result: Optional[List[Union[List[List[float]], str, str]]] = None,
         **kwargs,
-    ) -> RapidTableOutput:
+    ) -> LinelessTableOutput:
         s = time.perf_counter()
         rec_again = True
         need_ocr = True
@@ -92,7 +93,7 @@ def __call__(
                 sorted_polygons, idx_list = sorted_ocr_boxes(
                     [box_4_2_poly_to_box_4_1(box) for box in polygons]
                 )
-                return RapidTableOutput(
+                return LinelessTableOutput(
                     "",
                     sorted_polygons,
                     logi_points[idx_list],
@@ -121,6 +122,10 @@ def __call__(
             # 将同一个识别框中的ocr结果排序并同行合并
             t_rec_ocr_list = self.sort_and_gather_ocr_res(t_rec_ocr_list)
             # 渲染为html
+            polygons = [
+                box_4_1_poly_to_box_4_2(t_box_ocr["t_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"]]
@@ -132,13 +137,13 @@ def __call__(
             _, idx_list = sorted_ocr_boxes(
                 [t_box_ocr["t_box"] for t_box_ocr in t_rec_ocr_list]
             )
-            polygons = polygons.reshape(-1, 8)
+            polygons = np.array(polygons).reshape(-1, 8)
             logi_points = np.array(logi_points)
             elapse = time.perf_counter() - s
         except Exception:
             logging.warning(traceback.format_exc())
-            return RapidTableOutput("", None, None, 0.0)
-        return RapidTableOutput(pred_html, polygons, logi_points, elapse)
+            return LinelessTableOutput("", None, None, 0.0)
+        return LinelessTableOutput(pred_html, polygons, logi_points, elapse)
 
     def transform_res(
         self,

tests/test_lineless_table_rec.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 import pytest
 
-from lineless_table_rec.main import RapidTableInput, ModelType
+from lineless_table_rec.main import LinelessTableInput, ModelType
 
 cur_dir = Path(__file__).resolve().parent
 root_dir = cur_dir.parent
@@ -16,7 +16,7 @@
 from lineless_table_rec import LinelessTableRecognition
 
 test_file_dir = cur_dir / "test_files"
-input_args = RapidTableInput(model_type=ModelType.LORE.value)
+input_args = LinelessTableInput(model_type=ModelType.LORE.value)
 table_recog = LinelessTableRecognition(input_args)
 
 

tests/test_wired_table_rec.py

@@ -8,7 +8,7 @@
 from bs4 import BeautifulSoup
 from rapidocr_onnxruntime import RapidOCR
 
-from wired_table_rec.main import RapidTableInput, ModelType
+from wired_table_rec.main import WiredTableInput, ModelType
 from wired_table_rec.utils.utils import rescale_size
 from wired_table_rec.utils.utils_table_recover import (
     plot_html_table,
@@ -26,7 +26,7 @@
 from wired_table_rec import WiredTableRecognition
 
 test_file_dir = cur_dir / "test_files" / "wired"
-input_args = RapidTableInput(model_type=ModelType.UNET.value)
+input_args = WiredTableInput(model_type=ModelType.UNET.value)
 table_recog = WiredTableRecognition(input_args)
 ocr_engine = RapidOCR()
 

wired_table_rec/main.py

@@ -41,23 +41,23 @@ class ModelType(Enum):
 
 
 @dataclass
-class RapidTableInput:
+class WiredTableInput:
     model_type: Optional[str] = ModelType.UNET.value
     model_path: Union[str, Path, None, Dict[str, str]] = None
     use_cuda: bool = False
     device: str = "cpu"
 
 
 @dataclass
-class RapidTableOutput:
+class WiredTableOutput:
     pred_html: Optional[str] = None
     cell_bboxes: Optional[np.ndarray] = None
     logic_points: Optional[np.ndarray] = None
     elapse: Optional[float] = None
 
 
 class WiredTableRecognition:
-    def __init__(self, config: RapidTableInput):
+    def __init__(self, config: WiredTableInput):
         self.model_type = config.model_type
         if self.model_type not in KEY_TO_MODEL_URL:
             model_list = ",".join(KEY_TO_MODEL_URL)
@@ -85,7 +85,7 @@ def __call__(
         img: InputType,
         ocr_result: Optional[List[Union[List[List[float]], str, str]]] = None,
         **kwargs,
-    ) -> RapidTableOutput:
+    ) -> WiredTableOutput:
         s = time.perf_counter()
         rec_again = True
         need_ocr = True
@@ -100,7 +100,7 @@ def __call__(
         polygons, rotated_polygons = self.table_structure(img, **kwargs)
         if polygons is None:
             logging.warning("polygons is None.")
-            return RapidTableOutput("", None, None, 0.0)
+            return WiredTableOutput("", None, None, 0.0)
 
         try:
             table_res, logi_points = self.table_recover(
@@ -115,7 +115,7 @@ def __call__(
                 sorted_polygons, idx_list = sorted_ocr_boxes(
                     [box_4_2_poly_to_box_4_1(box) for box in polygons]
                 )
-                return RapidTableOutput(
+                return WiredTableOutput(
                     "",
                     sorted_polygons,
                     logi_points[idx_list],
@@ -137,14 +137,14 @@ def __call__(
                 for i, t_box_ocr in enumerate(t_rec_ocr_list)
             }
             pred_html = plot_html_table(logi_points, cell_box_det_map)
-            polygons = polygons.reshape(-1, 8)
+            polygons = np.array(polygons).reshape(-1, 8)
             logi_points = np.array(logi_points)
             elapse = time.perf_counter() - s
 
         except Exception:
             logging.warning(traceback.format_exc())
-            return RapidTableOutput("", None, None, 0.0)
-        return RapidTableOutput(pred_html, polygons, logi_points, elapse)
+            return WiredTableOutput("", None, None, 0.0)
+        return WiredTableOutput(pred_html, polygons, logi_points, elapse)
 
     def transform_res(
         self,
@@ -276,12 +276,12 @@ def main():
         raise ModuleNotFoundError(
             "Please install the rapidocr_onnxruntime by pip install rapidocr_onnxruntime."
         ) from exc
-
-    table_rec = WiredTableRecognition()
+    input_args = WiredTableInput()
+    table_rec = WiredTableRecognition(input_args)
     ocr_result, _ = ocr_engine(args.img_path)
-    table_str, elapse = table_rec(args.img_path, ocr_result)
-    print(table_str)
-    print(f"cost: {elapse:.5f}")
+    table_results = table_rec(args.img_path, ocr_result)
+    print(table_results.pred_html)
+    print(f"cost: {table_results.elapse:.5f}")
 
 
 if __name__ == "__main__":