upadate model

app.py

@@ -1,4 +1,3 @@
-import threading
 import time
 
 import cv2
@@ -11,6 +10,7 @@ from table_cls import TableCls
 from wired_table_rec import WiredTableRecognition
 
 from utils import plot_rec_box, LoadImage, format_html, box_4_2_poly_to_box_4_1
+
 img_loader = LoadImage()
 table_rec_path = "models/table_rec/ch_ppstructure_mobile_v2_SLANet.onnx"
 det_model_dir = {
@@ -39,8 +39,9 @@ example_images = [
     "images/wired4.jpg",
     "images/lineless2.png",
     "images/wired5.jpg",
-    "images/lineless3.jpg",
-    "images/wired6.jpg",
+    "images/lineless4.jpg",
+    "images/wired7.jpg",
+    "images/wired9.jpg",
 ]
 rapid_table_engine = RapidTable(model_path=table_rec_path)
 SLANet_plus_table_Engine = RapidTable()
@@ -93,7 +94,7 @@ def select_table_model(img, table_engine_type, det_model, rec_model):
         return lineless_table_engine, "lineless_table"
 
 
-def process_image(img, table_engine_type, det_model, rec_model):
+def process_image(img, table_engine_type, det_model, rec_model, small_box_cut_enhance):
     img = img_loader(img)
     start = time.time()
     table_engine, talbe_type = select_table_model(img, table_engine_type, det_model, rec_model)
@@ -115,6 +116,15 @@ def process_image(img, table_engine_type, det_model, rec_model):
             polygons = [[polygon[0], polygon[1], polygon[4], polygon[5]] for polygon in polygons]
         elif isinstance(table_engine, (WiredTableRecognition, LinelessTableRecognition)):
             html, table_rec_elapse, polygons, _, _ = table_engine(img, ocr_result=ocr_res)
+            if not small_box_cut_enhance:
+                html, table_rec_elapse, polygons, logic_points, ocr_res = table_engine(
+                    img, ocr_result=ocr_res,
+                    morph_close=False, more_h_lines=False, more_v_lines=False, extend_line=False
+                )
+            else:
+                html, table_rec_elapse, polygons, logic_points, ocr_res = table_engine(
+                    img, ocr_result=ocr_res
+                )
 
         sum_elapse = time.time() - start
         all_elapse = f"- table_type: {talbe_type}\n table all cost: {sum_elapse:.5f}\n - table rec cost: {table_rec_elapse:.5f}\n - ocr cost: {det_cost + cls_cost + rec_cost:.5f}"
@@ -136,23 +146,51 @@ def main():
             overflow-x: auto;
             white-space: nowrap;
         }
+        .header-links {
+            text-align: center;
+        }
+        .header-links a {
+            display: inline-block;
+            text-align: center;
+            margin-right: 10px;  /* 调整间距 */
+        }
     """) as demo:
+        gr.HTML(
+            "<h1 style='text-align: center;'><a href='https://github.com/RapidAI/TableStructureRec?tab=readme-ov-file'>TableStructureRec</a></h1>"
+        )
+        gr.HTML('''
+                                        <div class="header-links">
+                                          <a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.12-aff.svg"></a>
+                                          <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Mac%2C%20Win-pink.svg"></a>
+                                          <a href="https://pypi.org/project/lineless-table-rec/"><img alt="PyPI" src="https://img.shields.io/pypi/v/lineless-table-rec"></a>
+                                          <a href="https://pepy.tech/project/lineless-table-rec"><img src="https://static.pepy.tech/personalized-badge/lineless-table-rec?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20Lineless"></a>
+                                          <a href="https://pepy.tech/project/wired-table-rec"><img src="https://static.pepy.tech/personalized-badge/wired-table-rec?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads%20Wired"></a>
+                                          <a href="https://semver.org/"><img alt="SemVer2.0" src="https://img.shields.io/badge/SemVer-2.0-brightgreen"></a>
+                                          <a href="https://github.com/psf/black"><img src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
+                                          <a href="https://github.com/RapidAI/TableStructureRec/blob/c41bbd23898cb27a957ed962b0ffee3c74dfeff1/LICENSE"><img alt="GitHub" src="https://img.shields.io/badge/license-Apache 2.0-blue"></a>
+                                        </div>
+                                        ''')
         with gr.Row():  # 两列布局
             with gr.Tab("Options"):
                 with gr.Column(variant="panel", scale=1):  # 侧边栏，宽度比例为1
-                    img_input = gr.Image(label="Upload or Select Image",  sources="upload", value="images/lineless3.jpg")
+                    img_input = gr.Image(label="Upload or Select Image", sources="upload", value="images/lineless3.jpg")
 
                     # 示例图片选择器
                     examples = gr.Examples(
                         examples=example_images,
+                        examples_per_page=len(example_images),
                         inputs=img_input,
                         fn=lambda x: x,  # 简单返回图片路径
                         outputs=img_input,
-                        cache_examples=True
+                        cache_examples=False
                     )
 
                     table_engine_type = gr.Dropdown(table_engine_list, label="Select Recognition Table Engine",
                                                     value=table_engine_list[0])
+                    small_box_cut_enhance = gr.Checkbox(
+                        label="Box Cutting Enhancement (Disable to avoid excessive cutting, Enable to reduce missed cutting)",
+                        value=True
+                    )
                     det_model = gr.Dropdown(det_models_labels, label="Select OCR Detection Model",
                                             value=det_models_labels[0])
                     rec_model = gr.Dropdown(rec_models_labels, label="Select OCR Recognition Model",
@@ -172,7 +210,7 @@ def main():
 
         run_button.click(
             fn=process_image,
-            inputs=[img_input, table_engine_type, det_model, rec_model],
+            inputs=[img_input, table_engine_type, det_model, rec_model, small_box_cut_enhance],
             outputs=[html_output, table_boxes_output, ocr_boxes_output, elapse_text]
         )
 
@@ -180,4 +218,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()