Update app.py

app.py

@@ -33,7 +33,7 @@ table_engine_list = [
 # 示例图片路径
 example_images = [
     "images/wired1.png",
-    "images/wired2.png",
+    "images/wired2.jpg",
     "images/wired3.png",
     "images/lineless1.png",
     "images/wired4.jpg",
@@ -67,6 +67,17 @@ for det_model in det_model_dir.keys():
             rec_model_dir=rec_model_path
         )
 
+def trans_char_ocr_res(ocr_res):
+    word_result = []
+    for res in ocr_res:
+        score = res[2]
+        for word_box, word in zip(res[3], res[4]):
+            word_res = []
+            word_res.append(word_box)
+            word_res.append(word)
+            word_res.append(score)
+            word_result.append(word_res)
+    return word_result
 
 def select_ocr_model(det_model, rec_model):
     return ocr_engine_dict[f"{det_model}_{rec_model}"]
@@ -94,8 +105,10 @@ 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, small_box_cut_enhance):
-    img = img_loader(img)
+def process_image(img_input, small_box_cut_enhance, table_engine_type, char_ocr, rotated_fix, col_threshold, row_threshold):
+    det_model="mobile_det"
+    rec_model="mobile_rec"
+    img = img_loader(img_input)
     start = time.time()
     table_engine, talbe_type = select_table_model(img, table_engine_type, det_model, rec_model)
     ocr_engine = select_ocr_model(det_model, rec_model)
@@ -108,24 +121,20 @@ def process_image(img, table_engine_type, det_model, rec_model, small_box_cut_en
         ocr_boxes = result[0]['res']['boxes']
         all_elapse = f"- `table all cost: {time.time() - start:.5f}`"
     else:
-        ocr_res, ocr_infer_elapse = ocr_engine(img)
+        ocr_res, ocr_infer_elapse = ocr_engine(img, return_word_box=char_ocr)
         det_cost, cls_cost, rec_cost = ocr_infer_elapse
+        if char_ocr:
+            ocr_res = trans_char_ocr_res(ocr_res)
         ocr_boxes = [box_4_2_poly_to_box_4_1(ori_ocr[0]) for ori_ocr in ocr_res]
         if isinstance(table_engine, RapidTable):
             html, polygons, table_rec_elapse = table_engine(img, ocr_result=ocr_res)
             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
-                )
-
+            html, table_rec_elapse, polygons, logic_points, ocr_res = table_engine(img, ocr_result=ocr_res,
+                                                                                   enhance_box_line=small_box_cut_enhance,
+                                                                                   rotated_fix=rotated_fix,
+                                                                                   col_threshold=col_threshold,
+                                                                                   row_threshold=row_threshold)
         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}"
 
@@ -191,10 +200,33 @@ def main():
                         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",
-                                            value=rec_models_labels[0])
+                    char_ocr = gr.Checkbox(
+                        label="char rec ocr",
+                        value=False
+                    )
+                    rotate_adapt = gr.Checkbox(
+                        label="Table Rotate Rec Enhancement",
+                        value=False
+                    )
+                    col_threshold = gr.Slider(
+                        label="col threshold(determine same col)",
+                        minimum=5,
+                        maximum=100,
+                        value=15,
+                        step=5
+                    )
+                    row_threshold = gr.Slider(
+                        label="row threshold(determine same row)",
+                        minimum=5,
+                        maximum=100,
+                        value=10,
+                        step=5
+                    )
+
+                    # 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",
+                    #                         value=rec_models_labels[0])
 
                     run_button = gr.Button("Run")
                     gr.Markdown("# Elapsed Time")
@@ -210,7 +242,7 @@ def main():
 
         run_button.click(
             fn=process_image,
-            inputs=[img_input, table_engine_type, det_model, rec_model, small_box_cut_enhance],
+            inputs=[img_input, small_box_cut_enhance, table_engine_type, char_ocr, rotate_adapt, col_threshold, row_threshold],
             outputs=[html_output, table_boxes_output, ocr_boxes_output, elapse_text]
         )