Optim VisTable

demo.py

@@ -61,7 +61,9 @@ def demo_table():
     if not save_dir.exists():
         save_dir.mkdir(parents=True, exist_ok=True)
 
-    viser(img_path, save_dir, table_html_str, table_cell_bboxes)
+    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)
 
 

docs/README_Table.md

@@ -33,25 +33,30 @@ pip install rapid_table
 #### 使用方式
 1. python脚本运行
     ````python
+    from pathlib import Path
+
     from rapid_table import RapidTable
-    from rapidocr_onnxruntime import RapidOCR
+    from rapid_table import RapidTable, VisTable
 
     # RapidTable类提供model_path参数，可以自行指定上述2个模型，默认是en_ppstructure_mobile_v2_SLANet.onnx
     # table_engine = RapidTable(model_path='ch_ppstructure_mobile_v2_SLANet.onnx')
     table_engine = RapidTable()
     ocr_engine = RapidOCR()
+    viser = VisTable()
 
     img_path = 'test_images/table.jpg'
 
     ocr_result, _ = ocr_engine(img_path)
     table_html_str, table_cell_bboxes, elapse = 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_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)
 
-    # 可视化结果，如果不提供table_cell_bboxes，只会导出html文件
-    viser(img_path, save_dir, table_html_str, table_cell_bboxes)
     print(table_html_str)
     ````
 2. 终端运行
@@ -84,6 +89,9 @@ pip install rapid_table
     </div>
 
 #### 更新日志
+#### 2023-12-29 v0.1.3 update:
+- 优化可视化结果部分
+
 ##### 2023-12-27 v0.1.2 update:
 - 添加返回cell坐标框参数
 - 完善可视化函数

rapid_table/main.py

@@ -119,8 +119,17 @@ def main() -> None:
     viser = VisTable()
     if args.vis:
         img_path = Path(args.img_path)
+
         save_dir = img_path.resolve().parent
-        viser(img_path, save_dir, table_html_str, table_cell_bboxes)
+        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,
+        )
 
 
 if __name__ == "__main__":

rapid_table/utils.py

@@ -86,16 +86,17 @@ def __init__(
     def __call__(
         self,
         img_path: Union[str, Path],
-        save_dir: Union[str, Path],
         table_html_str: str,
+        save_html_path: Optional[str] = None,
         table_cell_bboxes: Optional[np.ndarray] = None,
+        save_drawed_path: Optional[str] = None,
     ) -> None:
-        save_html_path = save_dir / "table_res.html"
-        save_vis_img = save_dir / "table_res_vis.png"
-        self.export_html(table_html_str=table_html_str, save_path=save_html_path)
+        if save_html_path:
+            html_with_border = self.insert_border_style(table_html_str)
+            self.save_html(save_html_path, html_with_border)
 
         if table_cell_bboxes is None:
-            return
+            return None
 
         img = self.load_img(img_path)
 
@@ -107,17 +108,19 @@ def __call__(
         else:
             raise ValueError("Shape of table bounding boxes is not between in 4 or 8.")
 
-        self.save_img(save_vis_img, drawed_img)
+        if save_drawed_path:
+            self.save_img(save_drawed_path, drawed_img)
 
-    def export_html(self, table_html_str: str, save_path: Union[str, Path]):
+        return drawed_img
+
+    def insert_border_style(self, table_html_str: str):
         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_html_str.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)
+        html_with_border = f"{prefix_table}{style_res}<body>{suffix_table}"
+        return html_with_border
 
     @staticmethod
     def draw_rectangle(img: np.ndarray, boxes: np.ndarray) -> np.ndarray:
@@ -138,3 +141,8 @@ def draw_polylines(img: np.ndarray, points) -> np.ndarray:
     @staticmethod
     def save_img(save_path: Union[str, Path], img: np.ndarray):
         cv2.imwrite(str(save_path), img)
+
+    @staticmethod
+    def save_html(save_path: Union[str, Path], html: str):
+        with open(save_path, "w", encoding="utf-8") as f:
+            f.write(html)