Forráskód Böngészése

update v2 reference to v1 & adjust reference format (#2948)

* support inline formula embedding & update reference format

* update markdown format and table_recognition add cell sorting

* layout_parsing_v2 save_to_img remove formula_res & update table match

* layout_parsing_v2 add detection param

* update save_order name and segment split

* update v2 reference to v1 & adjust reference format

* update markdown default save_path.
shuai.liu 10 hónapja
szülő
commit
81cc313635

+ 1 - 1
api_examples/pipelines/test_layout_parsing_v2.py

@@ -14,7 +14,7 @@
 
 from paddlex import create_pipeline
 
-pipeline = create_pipeline(pipeline="layout_parsing_v2")
+pipeline = create_pipeline(pipeline_name="layout_parsing_v2")
 
 output = pipeline.predict(
     "./test_samples/demo_paper.png",

+ 1 - 1
paddlex/configs/pipelines/layout_parsing_v2.yaml

@@ -41,7 +41,7 @@ SubPipelines:
     SubModules:
       TextDetection:
         module_name: text_detection
-        model_name: PP-OCRv4_server_det
+        model_name: PP-OCRv4_mobile_det
         model_dir: null
         limit_side_len: 960
         limit_type: max

+ 1 - 0
paddlex/inference/common/result/mixin.py

@@ -601,6 +601,7 @@ class MarkdownMixin:
     def __init__(self, *args: list, **kwargs: dict):
         self._markdown_writer = MarkdownWriter(*args, **kwargs)
         self._save_funcs.append(self.save_to_markdown)
+        self.save_path = None
 
     @abstractmethod
     def _to_markdown(self):

+ 67 - 21
paddlex/inference/pipelines_new/layout_parsing/pipeline_v2.py

@@ -13,11 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-import os
-import sys
-from typing import Any, Dict, Optional, Union
-
-import cv2
+from typing import Optional, Union, Tuple
 import numpy as np
 
 from ....utils import logging
@@ -26,10 +22,9 @@ from ...common.reader import ReadImage
 from ...models_new.object_detection.result import DetResult
 from ...utils.pp_option import PaddlePredictorOption
 from ..base import BasePipeline
-from ..components import convert_points_to_boxes
 from ..ocr.result import OCRResult
 from .result_v2 import LayoutParsingResultV2
-from .utils import get_structure_res
+from .utils import get_single_block_parsing_res
 from .utils import get_sub_regions_ocr_res
 
 # [TODO] 待更新models_new到models
@@ -183,7 +178,10 @@ class LayoutParsingPipelineV2(BasePipeline):
             if box_info["label"].lower() in ["formula", "table", "seal"]:
                 object_boxes.append(box_info["coordinate"])
         object_boxes = np.array(object_boxes)
-        return get_sub_regions_ocr_res(overall_ocr_res, object_boxes, flag_within=False)
+        sub_regions_ocr_res = get_sub_regions_ocr_res(
+            overall_ocr_res, object_boxes, flag_within=False
+        )
+        return sub_regions_ocr_res
 
     def check_model_settings_valid(self, input_params: dict) -> bool:
         """
@@ -222,6 +220,49 @@ class LayoutParsingPipelineV2(BasePipeline):
 
         return True
 
+    def get_layout_parsing_res(
+        self,
+        image: list,
+        layout_det_res: DetResult,
+        overall_ocr_res: OCRResult,
+        table_res_list: list,
+        seal_res_list: list,
+    ) -> list:
+        """
+        Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
+        Args:
+            image (list): The input image.
+            overall_ocr_res (OCRResult): An object containing the overall OCR results, including detected text boxes and recognized text. The structure is expected to have:
+            - "input_img": The image on which OCR was performed.
+            - "dt_boxes": A list of detected text box coordinates.
+            - "rec_texts": A list of recognized text corresponding to the detected boxes.
+
+            layout_det_res (DetResult): An object containing the layout detection results, including detected layout boxes and their labels. The structure is expected to have:
+                - "boxes": A list of dictionaries with keys "coordinate" for box coordinates and "label" for the type of content.
+
+            table_res_list (list): A list of table detection results, where each item is a dictionary containing:
+                - "layout_bbox": The bounding box of the table layout.
+                - "pred_html": The predicted HTML representation of the table.
+        Returns:
+            list: A list of dictionaries representing the layout parsing result.
+        """
+        layout_parsing_res = get_single_block_parsing_res(
+            overall_ocr_res=overall_ocr_res,
+            layout_det_res=layout_det_res,
+            table_res_list=table_res_list,
+            seal_res_list=seal_res_list,
+        )
+
+        parsing_res_list = [
+            {
+                "block_bbox": [0, 0, 2550, 2550],
+                "block_size": [image.shape[1], image.shape[0]],
+                "sub_blocks": layout_parsing_res,
+            },
+        ]
+
+        return parsing_res_list
+
     def get_model_settings(
         self,
         use_doc_orientation_classify: bool | None,
@@ -293,6 +334,10 @@ class LayoutParsingPipelineV2(BasePipeline):
         seal_det_box_thresh: float | None = None,
         seal_det_unclip_ratio: float | None = None,
         seal_rec_score_thresh: float | None = None,
+        layout_threshold: Optional[Union[float, dict]] = None,
+        layout_nms: Optional[bool] = None,
+        layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
+        layout_merge_bboxes_mode: Optional[str] = None,
         **kwargs,
     ) -> LayoutParsingResultV2:
         """
@@ -340,7 +385,13 @@ class LayoutParsingPipelineV2(BasePipeline):
             doc_preprocessor_image = doc_preprocessor_res["output_img"]
 
             layout_det_res = next(
-                self.layout_det_model(doc_preprocessor_image),
+                self.layout_det_model(
+                    doc_preprocessor_image,
+                    threshold=layout_threshold,
+                    layout_nms=layout_nms,
+                    layout_unclip_ratio=layout_unclip_ratio,
+                    layout_merge_bboxes_mode=layout_merge_bboxes_mode,
+                )
             )
 
             if model_settings["use_formula_recognition"]:
@@ -448,18 +499,13 @@ class LayoutParsingPipelineV2(BasePipeline):
                     "input_img"
                 ]
 
-            structure_res = get_structure_res(
-                overall_ocr_res,
-                layout_det_res,
-                table_res_list,
+            parsing_res_list = self.get_layout_parsing_res(
+                doc_preprocessor_image,
+                layout_det_res=layout_det_res,
+                overall_ocr_res=overall_ocr_res,
+                table_res_list=table_res_list,
+                seal_res_list=seal_res_list,
             )
-            structure_res_list = [
-                {
-                    "block_bbox": [0, 0, 2550, 2550],
-                    "block_size": [image_array.shape[1], image_array.shape[0]],
-                    "sub_blocks": structure_res,
-                },
-            ]
 
             single_img_res = {
                 "input_path": batch_data.input_paths[0],
@@ -471,7 +517,7 @@ class LayoutParsingPipelineV2(BasePipeline):
                 "table_res_list": table_res_list,
                 "seal_res_list": seal_res_list,
                 "formula_res_list": formula_res_list,
-                "layout_parsing_result": structure_res_list,
+                "parsing_res_list": parsing_res_list,
                 "model_settings": model_settings,
             }
             yield LayoutParsingResultV2(single_img_res)

+ 30 - 13
paddlex/inference/pipelines_new/layout_parsing/result_v2.py

@@ -248,13 +248,19 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
         Returns:
             None
         """
-        input_name = self["input_path"]
+        input_path = Path(self["input_path"])
+        page_index = self["page_index"]
         save_path = Path(save_path)
         if save_path.suffix.lower() not in (".jpg", ".png"):
-            save_path = save_path / f"{input_name}.jpg"
+            if input_path.suffix.lower() == ".pdf":
+                save_path = save_path / f"page_{page_index}.jpg"
+            else:
+                save_path = save_path / f"{input_path.stem}.jpg"
         else:
             save_path = save_path.with_suffix("")
-        ordering_image_path = save_path.parent / f"{save_path.stem}_ordering.jpg"
+        ordering_image_path = (
+            save_path.parent / f"{save_path.stem}_layout_order_res.jpg"
+        )
 
         try:
             image = Image.fromarray(self["doc_preprocessor_res"]["output_img"])
@@ -264,7 +270,7 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
 
         draw = ImageDraw.Draw(image, "RGBA")
 
-        parsing_result = self["layout_parsing_result"]
+        parsing_result = self["parsing_res_list"]
         for block in parsing_result:
             if self.already_sorted == False:
                 block = get_layout_ordering(
@@ -304,9 +310,14 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
         Returns:
             Dict
         """
-        save_path = Path(self.save_path)
+        if self.save_path == None:
+            is_save_mk_img = False
+        else:
+            is_save_mk_img = True
+            save_path = Path(self.save_path)
+
+        parsing_result = self["parsing_res_list"]
 
-        parsing_result = self["layout_parsing_result"]
         for block in parsing_result:
             if self.already_sorted == False:
                 block = get_layout_ordering(
@@ -322,11 +333,13 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
                     already_sorted=self.already_sorted,
                 )
         self.already_sorted == True
-        recursive_img_array2path(
-            self["layout_parsing_result"],
-            save_path.parent,
-            labels=["img"],
-        )
+
+        if is_save_mk_img:
+            recursive_img_array2path(
+                self["parsing_res_list"],
+                save_path.parent,
+                labels=["img"],
+            )
 
         def _format_data(obj):
 
@@ -355,6 +368,9 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
                 )
 
             def format_image():
+                if is_save_mk_img is False:
+                    return ""
+
                 img_tags = []
                 if "img" in sub_block["image"]:
                     img_tags.append(
@@ -429,11 +445,12 @@ class LayoutParsingResultV2(BaseCVResult, HtmlMixin, XlsxMixin, MarkdownMixin):
                 "chart": format_chart,
                 "formula": lambda: f"$${sub_block['formula']}$$",
                 "table": format_table,
-                "reference": format_reference,
+                # "reference": format_reference,
+                "reference": lambda: sub_block["reference"],
                 "algorithm": lambda: sub_block["algorithm"].strip("\n"),
                 "seal": lambda: sub_block["seal"].strip("\n"),
             }
-            parsing_result = obj["layout_parsing_result"]
+            parsing_result = obj["parsing_res_list"]
             markdown_content = ""
             for block in parsing_result:  # for each block show ordering results
                 sub_blocks = block["sub_blocks"]

+ 70 - 19
paddlex/inference/pipelines_new/layout_parsing/utils.py

@@ -15,6 +15,7 @@
 __all__ = [
     "get_sub_regions_ocr_res",
     "get_layout_ordering",
+    "get_single_block_parsing_res",
     "recursive_img_array2path",
     "get_show_color",
     "sorted_layout_boxes",
@@ -233,7 +234,55 @@ def _whether_y_overlap_exceeds_threshold(bbox1, bbox2, overlap_ratio_threshold=0
     return (overlap / min_height) > overlap_ratio_threshold
 
 
-def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.7):
+def _adjust_span_text(span, prepend=False, append=False):
+    """
+    Adjust the text of a span by prepending or appending a newline.
+
+    Args:
+        span (list): A list where the second element is the text of the span.
+        prepend (bool): If True, prepend a newline to the text.
+        append (bool): If True, append a newline to the text.
+
+    Returns:
+        None: The function modifies the span in place.
+    """
+    if prepend:
+        span[1] = "\n" + span[1]
+    if append:
+        span[1] = span[1] + "\n"
+
+
+def _format_line(line, layout_min, layout_max, is_reference=False):
+    """
+    Format a line of text spans based on layout constraints.
+
+    Args:
+        line (list): A list of spans, where each span is a list containing a bounding box and text.
+        layout_min (int): The minimum x-coordinate of the layout bounding box.
+        layout_max (int): The maximum x-coordinate of the layout bounding box.
+        is_reference (bool): A flag indicating whether the line is a reference line, which affects formatting rules.
+
+    Returns:
+        None: The function modifies the line in place.
+    """
+    first_span = line[0]
+    end_span = line[-1]
+
+    if not is_reference:
+        if first_span[0][0] - layout_min > 10:
+            _adjust_span_text(first_span, prepend=True)
+        if layout_max - end_span[0][2] > 10:
+            _adjust_span_text(end_span, append=True)
+    else:
+        if first_span[0][0] - layout_min < 5:
+            _adjust_span_text(first_span, prepend=True)
+        if layout_max - end_span[0][2] > 20:
+            _adjust_span_text(end_span, append=True)
+
+
+def _sort_ocr_res_by_y_projection(
+    label, layout_bbox, ocr_res, line_height_iou_threshold=0.7
+):
     """
     Sorts OCR results based on their spatial arrangement, grouping them into lines and blocks.
 
@@ -257,6 +306,8 @@ def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.
     rec_texts = ocr_res["rec_texts"]
 
     x_min, _, x_max, _ = layout_bbox
+    inline_x_min = min([box[0] for box in boxes])
+    inline_x_max = max([box[2] for box in boxes])
 
     spans = list(zip(boxes, rec_texts))
 
@@ -287,13 +338,10 @@ def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.
 
     for line in lines:
         line.sort(key=lambda span: span[0][0])
-        first_span = line[0]
-        end_span = line[-1]
-
-        if first_span[0][0] - x_min > 15:
-            first_span[1] = "\n" + first_span[1]
-        if x_max - end_span[0][2] > 15:
-            end_span[1] = end_span[1] + "\n"
+        if label == "reference":
+            line = _format_line(line, inline_x_min, inline_x_max, is_reference=True)
+        else:
+            line = _format_line(line, x_min, x_max)
 
     # Flatten lines back into a single list for boxes and texts
     ocr_res["boxes"] = [span[0] for line in lines for span in line]
@@ -302,10 +350,11 @@ def _sort_box_by_y_projection(layout_bbox, ocr_res, line_height_iou_threshold=0.
     return ocr_res
 
 
-def get_structure_res(
+def get_single_block_parsing_res(
     overall_ocr_res: OCRResult,
     layout_det_res: DetResult,
-    table_res_list,
+    table_res_list: list,
+    seal_res_list: list,
 ) -> OCRResult:
     """
     Extract structured information from OCR and layout detection results.
@@ -330,7 +379,7 @@ def get_structure_res(
             - "layout_bbox": The coordinates of the layout box.
     """
 
-    structure_boxes = []
+    single_block_layout_parsing_res = []
     input_img = overall_ocr_res["doc_preprocessor_res"]["output_img"]
 
     for box_info in layout_det_res["boxes"]:
@@ -348,7 +397,7 @@ def get_structure_res(
                     )
                     > 0.5
                 ):
-                    structure_boxes.append(
+                    single_block_layout_parsing_res.append(
                         {
                             "label": label,
                             f"{label}": table_res["pred_html"],
@@ -370,12 +419,14 @@ def get_structure_res(
                     rec_res["flag"] = True
 
             if rec_res["flag"]:
-                rec_res = _sort_box_by_y_projection(layout_bbox, rec_res, 0.7)
+                rec_res = _sort_ocr_res_by_y_projection(
+                    label, layout_bbox, rec_res, 0.7
+                )
                 rec_res_first_bbox = rec_res["boxes"][0]
                 rec_res_end_bbox = rec_res["boxes"][-1]
-                if rec_res_first_bbox[0] - layout_bbox[0] < 20:
+                if rec_res_first_bbox[0] - layout_bbox[0] < 10:
                     seg_start_flag = False
-                if layout_bbox[2] - rec_res_end_bbox[2] < 20:
+                if layout_bbox[2] - rec_res_end_bbox[2] < 10:
                     seg_end_flag = False
                 if label == "formula":
                     rec_res["rec_texts"] = [
@@ -384,7 +435,7 @@ def get_structure_res(
                     ]
 
             if label in ["chart", "image"]:
-                structure_boxes.append(
+                single_block_layout_parsing_res.append(
                     {
                         "label": label,
                         f"{label}": {
@@ -399,7 +450,7 @@ def get_structure_res(
                     },
                 )
             else:
-                structure_boxes.append(
+                single_block_layout_parsing_res.append(
                     {
                         "label": label,
                         f"{label}": "".join(rec_res["rec_texts"]),
@@ -409,7 +460,7 @@ def get_structure_res(
                     },
                 )
 
-    return structure_boxes
+    return single_block_layout_parsing_res
 
 
 def _projection_by_bboxes(boxes: np.ndarray, axis: int) -> np.ndarray:
@@ -1145,7 +1196,7 @@ def get_layout_ordering(data, no_mask_labels=[], already_sorted=False):
     and assign an ordering index based on their positions.
 
     Modifies:
-        The 'parsing_result' list in 'layout_parsing_result' by adding an 'index' to each block.
+        The 'parsing_result' list in 'parsing_res_list' by adding an 'index' to each block.
 
     """
     if already_sorted: