浏览代码

fix layout_parsing_v2 ocr (#3267)

shuai.liu 9 月之前
父节点
当前提交
2ec417a6f0
共有 1 个文件被更改,包括 113 次插入31 次删除
  1. 113 31
      paddlex/inference/pipelines/layout_parsing/pipeline_v2.py

+ 113 - 31
paddlex/inference/pipelines/layout_parsing/pipeline_v2.py

@@ -225,20 +225,113 @@ class LayoutParsingPipelineV2(BasePipeline):
         overall_ocr_res: OCRResult,
         overall_ocr_res: OCRResult,
         table_res_list: list,
         table_res_list: list,
         seal_res_list: list,
         seal_res_list: list,
+        formula_res_list: list,
+        text_det_limit_side_len: Optional[int] = None,
+        text_det_limit_type: Optional[str] = None,
+        text_det_thresh: Optional[float] = None,
+        text_det_box_thresh: Optional[float] = None,
+        text_det_unclip_ratio: Optional[float] = None,
+        text_rec_score_thresh: Optional[float] = None,
     ) -> list:
     ) -> list:
         """
         """
-        Get the layout parsing result based on the layout detection result, OCR result, and other recognition results.
-
+        Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
         Args:
         Args:
             image (list): The input image.
             image (list): The input image.
-            layout_det_res (DetResult): The layout detection results.
-            overall_ocr_res (OCRResult): The overall OCR results.
-            table_res_list (list): A list of table detection results.
-            seal_res_list (list): A list of seal detection results.
-
+            layout_det_res (DetResult): The detection result containing the layout information of the document.
+            overall_ocr_res (OCRResult): The overall OCR result containing text information.
+            table_res_list (list): A list of table recognition results.
+            seal_res_list (list): A list of seal recognition results.
+            formula_res_list (list): A list of formula recognition results.
+            text_det_limit_side_len (Optional[int], optional): The maximum side length of the text detection region. Defaults to None.
+            text_det_limit_type (Optional[str], optional): The type of limit for the text detection region. Defaults to None.
+            text_det_thresh (Optional[float], optional): The confidence threshold for text detection. Defaults to None.
+            text_det_box_thresh (Optional[float], optional): The confidence threshold for text detection bounding boxes. Defaults to None
+            text_det_unclip_ratio (Optional[float], optional): The unclip ratio for text detection. Defaults to None.
+            text_rec_score_thresh (Optional[float], optional): The score threshold for text recognition. Defaults to None.
         Returns:
         Returns:
             list: A list of dictionaries representing the layout parsing result.
             list: A list of dictionaries representing the layout parsing result.
         """
         """
+        matched_ocr_dict = {}
+        image = np.array(image)
+        object_boxes = []
+
+        for object_box_idx, box_info in enumerate(layout_det_res["boxes"]):
+            box = box_info["coordinate"]
+            label = box_info["label"].lower()
+            object_boxes.append(box)
+
+            if label not in ["formula", "table", "seal"]:
+                _, matched_idxs = get_sub_regions_ocr_res(
+                    overall_ocr_res, [box], return_match_idx=True
+                )
+                for matched_idx in matched_idxs:
+                    if matched_ocr_dict.get(matched_idx, None) is None:
+                        matched_ocr_dict[matched_idx] = [object_box_idx]
+                    else:
+                        matched_ocr_dict[matched_idx].append(object_box_idx)
+
+        already_processed = set()
+        for matched_idx, layout_box_ids in matched_ocr_dict.items():
+            if len(layout_box_ids) <= 1:
+                continue
+
+            # one ocr is matched to multiple layout boxes, split the text into multiple lines
+            for idx in layout_box_ids:
+                if idx in already_processed:
+                    continue
+
+                already_processed.add(idx)
+                wht_im = np.ones(image.shape, dtype=image.dtype) * 255
+                box = object_boxes[idx]
+                x1, y1, x2, y2 = [int(i) for i in box]
+                wht_im[y1:y2, x1:x2, :] = image[y1:y2, x1:x2, :]
+                sub_ocr_res = next(
+                    self.general_ocr_pipeline(
+                        wht_im,
+                        text_det_limit_side_len=text_det_limit_side_len,
+                        text_det_limit_type=text_det_limit_type,
+                        text_det_thresh=text_det_thresh,
+                        text_det_box_thresh=text_det_box_thresh,
+                        text_det_unclip_ratio=text_det_unclip_ratio,
+                        text_rec_score_thresh=text_rec_score_thresh,
+                    )
+                )
+                _, matched_idxs = get_sub_regions_ocr_res(
+                    overall_ocr_res, [box], return_match_idx=True
+                )
+                for matched_idx in sorted(matched_idxs, reverse=True):
+                    del overall_ocr_res["dt_polys"][matched_idx]
+                    del overall_ocr_res["rec_texts"][matched_idx]
+                    overall_ocr_res["rec_boxes"] = np.delete(
+                        overall_ocr_res["rec_boxes"], matched_idx, axis=0
+                    )
+                    del overall_ocr_res["rec_polys"][matched_idx]
+                    del overall_ocr_res["rec_scores"][matched_idx]
+
+                overall_ocr_res["dt_polys"].extend(sub_ocr_res["dt_polys"])
+                overall_ocr_res["rec_texts"].extend(sub_ocr_res["rec_texts"])
+                overall_ocr_res["rec_boxes"] = np.concatenate(
+                    [overall_ocr_res["rec_boxes"], sub_ocr_res["rec_boxes"]], axis=0
+                )
+                overall_ocr_res["rec_polys"].extend(sub_ocr_res["rec_polys"])
+                overall_ocr_res["rec_scores"].extend(sub_ocr_res["rec_scores"])
+
+        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),
+            ]
+            overall_ocr_res["dt_polys"].append(poly_points)
+            overall_ocr_res["rec_texts"].append(f"${formula_res['rec_formula']}$")
+            overall_ocr_res["rec_boxes"] = np.vstack(
+                (overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
+            )
+            overall_ocr_res["rec_polys"].append(poly_points)
+            overall_ocr_res["rec_scores"].append(1)
+
         layout_parsing_res = get_single_block_parsing_res(
         layout_parsing_res = get_single_block_parsing_res(
             overall_ocr_res=overall_ocr_res,
             overall_ocr_res=overall_ocr_res,
             layout_det_res=layout_det_res,
             layout_det_res=layout_det_res,
@@ -445,24 +538,6 @@ class LayoutParsingPipelineV2(BasePipeline):
                         text_rec_score_thresh=text_rec_score_thresh,
                         text_rec_score_thresh=text_rec_score_thresh,
                     ),
                     ),
                 )
                 )
-
-                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),
-                    ]
-                    overall_ocr_res["dt_polys"].append(poly_points)
-                    overall_ocr_res["rec_texts"].append(
-                        f"${formula_res['rec_formula']}$"
-                    )
-                    overall_ocr_res["rec_boxes"] = np.vstack(
-                        (overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
-                    )
-                    overall_ocr_res["rec_polys"].append(poly_points)
-                    overall_ocr_res["rec_scores"].append(1)
             else:
             else:
                 overall_ocr_res = {}
                 overall_ocr_res = {}
 
 
@@ -511,20 +586,27 @@ class LayoutParsingPipelineV2(BasePipeline):
             else:
             else:
                 seal_res_list = []
                 seal_res_list = []
 
 
-            for formula_res in formula_res_list:
-                x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
-                doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
-                    "input_img"
-                ]
-
             parsing_res_list = self.get_layout_parsing_res(
             parsing_res_list = self.get_layout_parsing_res(
                 doc_preprocessor_image,
                 doc_preprocessor_image,
                 layout_det_res=layout_det_res,
                 layout_det_res=layout_det_res,
                 overall_ocr_res=overall_ocr_res,
                 overall_ocr_res=overall_ocr_res,
                 table_res_list=table_res_list,
                 table_res_list=table_res_list,
                 seal_res_list=seal_res_list,
                 seal_res_list=seal_res_list,
+                formula_res_list=formula_res_list,
+                text_det_limit_side_len=text_det_limit_side_len,
+                text_det_limit_type=text_det_limit_type,
+                text_det_thresh=text_det_thresh,
+                text_det_box_thresh=text_det_box_thresh,
+                text_det_unclip_ratio=text_det_unclip_ratio,
+                text_rec_score_thresh=text_rec_score_thresh,
             )
             )
 
 
+            for formula_res in formula_res_list:
+                x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
+                doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
+                    "input_img"
+                ]
+
             single_img_res = {
             single_img_res = {
                 "input_path": batch_data.input_paths[0],
                 "input_path": batch_data.input_paths[0],
                 "page_index": batch_data.page_indexes[0],
                 "page_index": batch_data.page_indexes[0],