|
@@ -13,11 +13,7 @@
|
|
|
# limitations under the License.
|
|
# limitations under the License.
|
|
|
from __future__ import annotations
|
|
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
|
|
import numpy as np
|
|
|
|
|
|
|
|
from ....utils import logging
|
|
from ....utils import logging
|
|
@@ -26,10 +22,9 @@ from ...common.reader import ReadImage
|
|
|
from ...models_new.object_detection.result import DetResult
|
|
from ...models_new.object_detection.result import DetResult
|
|
|
from ...utils.pp_option import PaddlePredictorOption
|
|
from ...utils.pp_option import PaddlePredictorOption
|
|
|
from ..base import BasePipeline
|
|
from ..base import BasePipeline
|
|
|
-from ..components import convert_points_to_boxes
|
|
|
|
|
from ..ocr.result import OCRResult
|
|
from ..ocr.result import OCRResult
|
|
|
from .result_v2 import LayoutParsingResultV2
|
|
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
|
|
from .utils import get_sub_regions_ocr_res
|
|
|
|
|
|
|
|
# [TODO] 待更新models_new到models
|
|
# [TODO] 待更新models_new到models
|
|
@@ -183,7 +178,10 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
if box_info["label"].lower() in ["formula", "table", "seal"]:
|
|
if box_info["label"].lower() in ["formula", "table", "seal"]:
|
|
|
object_boxes.append(box_info["coordinate"])
|
|
object_boxes.append(box_info["coordinate"])
|
|
|
object_boxes = np.array(object_boxes)
|
|
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:
|
|
def check_model_settings_valid(self, input_params: dict) -> bool:
|
|
|
"""
|
|
"""
|
|
@@ -222,6 +220,49 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
|
|
|
|
|
return True
|
|
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(
|
|
def get_model_settings(
|
|
|
self,
|
|
self,
|
|
|
use_doc_orientation_classify: bool | None,
|
|
use_doc_orientation_classify: bool | None,
|
|
@@ -293,6 +334,10 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
seal_det_box_thresh: float | None = None,
|
|
seal_det_box_thresh: float | None = None,
|
|
|
seal_det_unclip_ratio: float | None = None,
|
|
seal_det_unclip_ratio: float | None = None,
|
|
|
seal_rec_score_thresh: 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,
|
|
**kwargs,
|
|
|
) -> LayoutParsingResultV2:
|
|
) -> LayoutParsingResultV2:
|
|
|
"""
|
|
"""
|
|
@@ -340,7 +385,13 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
doc_preprocessor_image = doc_preprocessor_res["output_img"]
|
|
doc_preprocessor_image = doc_preprocessor_res["output_img"]
|
|
|
|
|
|
|
|
layout_det_res = next(
|
|
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"]:
|
|
if model_settings["use_formula_recognition"]:
|
|
@@ -448,18 +499,13 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
"input_img"
|
|
"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 = {
|
|
single_img_res = {
|
|
|
"input_path": batch_data.input_paths[0],
|
|
"input_path": batch_data.input_paths[0],
|
|
@@ -471,7 +517,7 @@ class LayoutParsingPipelineV2(BasePipeline):
|
|
|
"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,
|
|
"formula_res_list": formula_res_list,
|
|
|
- "layout_parsing_result": structure_res_list,
|
|
|
|
|
|
|
+ "parsing_res_list": parsing_res_list,
|
|
|
"model_settings": model_settings,
|
|
"model_settings": model_settings,
|
|
|
}
|
|
}
|
|
|
yield LayoutParsingResultV2(single_img_res)
|
|
yield LayoutParsingResultV2(single_img_res)
|