Merge pull request #26 from RapidAI/optim_wired_rec

feat(wired_table_rec): Add format_html
RapidAI · Sep 19, 2024 · 8f4eb21 · 8f4eb21
2 parents c65caf9 + cae1eef
commit 8f4eb21
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 6 deletions.
diff --git a/demo_wired.py b/demo_wired.py
@@ -1,26 +1,34 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-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
@@ -3,7 +3,7 @@
 # @Contact: [email protected]
 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"""
+    <!DOCTYPE html>
+    <html lang="zh-CN">
+    <head>
+    <meta charset="UTF-8">
+    <title>Complex Table Example</title>
+    <style>
+        table {{
+            border-collapse: collapse;
+            width: 100%;
+        }}
+        th, td {{
+            border: 1px solid black;
+            padding: 8px;
+            text-align: center;
+        }}
+        th {{
+            background-color: #f2f2f2;
+        }}
+    </style>
+    </head>
+    <body>
+    {html}
+    </body>
+    </html>
+    """