Jelajahi Sumber

support to save tabel and formula as image when the table rec pipeline or formula rec pipeline is not used

gaotingquan 6 bulan lalu
induk
melakukan
68d1d19d63

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

@@ -791,7 +791,7 @@ class _LayoutParsingPipelineV2(BasePipeline):
                     text_rec_score_thresh=text_rec_score_thresh,
                 )
 
-            if label in ["chart", "image", "seal"]:
+            if label in ["chart", "image", "seal", "table", "formula"]:
                 x_min, y_min, x_max, y_max = list(map(int, block_bbox))
                 img_path = f"imgs/img_in_table_box_{x_min}_{y_min}_{x_max}_{y_max}.jpg"
                 img = Image.fromarray(image[y_min:y_max, x_min:x_max, ::-1])

+ 17 - 4
paddlex/inference/pipelines/layout_parsing/result_v2.py

@@ -421,11 +421,9 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
                     original_image_width=original_image_width,
                 )
             )
-            format_table = lambda block: "\n" + format_text_func(block)
         else:
             format_text_func = lambda block: block.content
             format_image_func = format_image_plain_func
-            format_table = lambda block: simplify_table_func("\n" + block.content)
 
         if self["model_settings"].get("use_chart_recognition", False):
             format_chart_func = format_chart2table_func
@@ -439,6 +437,21 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
         else:
             format_seal_func = format_image_func
 
+        if self["model_settings"].get("use_table_recognition", False):
+            if pretty:
+                format_table_func = lambda block: "\n" + format_text_func(block)
+            else:
+                format_table_func = lambda block: simplify_table_func(
+                    "\n" + block.content
+                )
+        else:
+            format_table_func = format_image_func
+
+        if self["model_settings"].get("use_formula_recognition", False):
+            format_formula_func = lambda block: f"$${block.content}$$"
+        else:
+            format_formula_func = format_image_func
+
         handle_funcs_dict = {
             "paragraph_title": format_title_func,
             "abstract_title": format_title_func,
@@ -465,8 +478,8 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
             ),
             "image": format_image_func,
             "chart": format_chart_func,
-            "formula": lambda block: f"$${block.content}$$",
-            "table": format_table,
+            "formula": format_formula_func,
+            "table": format_table_func,
             "reference": partial(
                 format_first_line_func,
                 templates=["参考文献", "references"],

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

@@ -607,7 +607,7 @@ def remove_extra_space(input_text: str) -> str:
 def gather_imgs(original_img, layout_det_objs):
     imgs_in_doc = []
     for det_obj in layout_det_objs:
-        if det_obj["label"] in ("image", "chart", "seal"):
+        if det_obj["label"] in ("image", "chart", "seal", "formula", "table"):
             x_min, y_min, x_max, y_max = list(map(int, det_obj["coordinate"]))
             img_path = f"imgs/img_in_table_box_{x_min}_{y_min}_{x_max}_{y_max}.jpg"
             img = Image.fromarray(original_img[y_min:y_max, x_min:x_max, ::-1])