Browse Source

update pre_cut labels&fix table formula bug

liushuai35 8 months ago
parent
commit
816051ec73

+ 21 - 1
paddlex/inference/pipelines/layout_parsing/pipeline_v2.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 from typing import Optional, Union, Tuple, Iterator
 import numpy as np
 import re
+import copy
 
 from ....utils import logging
 from ...common.batch_sampler import ImageBatchSampler
@@ -536,6 +537,25 @@ class LayoutParsingPipelineV2(BasePipeline):
                 overall_ocr_res = {}
 
             if model_settings["use_table_recognition"]:
+                table_overall_ocr_res = copy.deepcopy(overall_ocr_res)
+                for formula_res in formula_res_list:
+                    x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
+                    poly_points = [
+                        (x_min, y_min),
+                        (x_max, y_min),
+                        (x_max, y_max),
+                        (x_min, y_max),
+                    ]
+                    table_overall_ocr_res["dt_polys"].append(poly_points)
+                    table_overall_ocr_res["rec_texts"].append(
+                        f"${formula_res['rec_formula']}$"
+                    )
+                    table_overall_ocr_res["rec_boxes"] = np.vstack(
+                        (table_overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
+                    )
+                    table_overall_ocr_res["rec_polys"].append(poly_points)
+                    table_overall_ocr_res["rec_scores"].append(1)
+
                 table_res_all = next(
                     self.table_recognition_pipeline(
                         doc_preprocessor_image,
@@ -543,7 +563,7 @@ class LayoutParsingPipelineV2(BasePipeline):
                         use_doc_unwarping=False,
                         use_layout_detection=False,
                         use_ocr_model=False,
-                        overall_ocr_res=overall_ocr_res,
+                        overall_ocr_res=table_overall_ocr_res,
                         layout_det_res=layout_det_res,
                         cell_sort_by_y_projection=True,
                     ),

+ 10 - 1
paddlex/inference/pipelines/layout_parsing/utils.py

@@ -1105,7 +1105,16 @@ def _get_sub_category(
     vision_labels = ["image", "table", "chart", "figure"]
     vision_title_labels = ["figure_title", "chart_title", "table_title"]
     all_labels = title_labels + sub_title_labels + vision_labels + vision_title_labels
-    special_pre_cut_labels = title_labels + sub_title_labels
+    special_pre_cut_labels = sub_title_labels
+
+    # single doc title is irregular,pre cut not applicable
+    num_doc_title = 0
+    for block in blocks:
+        if block["block_label"] == "doc_title":
+            num_doc_title += 1
+            if num_doc_title == 2:
+                special_pre_cut_labels = title_labels + sub_title_labels
+                break
 
     min_x = min(block["block_bbox"][0] for block in blocks)
     min_y = min(block["block_bbox"][1] for block in blocks)