pipeline_v2.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import annotations
  15. import os
  16. import sys
  17. from typing import Any, Dict, Optional, Union
  18. import cv2
  19. import numpy as np
  20. from ....utils import logging
  21. from ...common.batch_sampler import ImageBatchSampler
  22. from ...common.reader import ReadImage
  23. from ...models_new.object_detection.result import DetResult
  24. from ...utils.pp_option import PaddlePredictorOption
  25. from ..base import BasePipeline
  26. from ..components import convert_points_to_boxes
  27. from ..ocr.result import OCRResult
  28. from .result_v2 import LayoutParsingResultV2
  29. from .utils import get_structure_res
  30. from .utils import get_sub_regions_ocr_res
  31. # [TODO] 待更新models_new到models
  32. class LayoutParsingPipelineV2(BasePipeline):
  33. """Layout Parsing Pipeline V2"""
  34. entities = ["layout_parsing_v2"]
  35. def __init__(
  36. self,
  37. config: dict,
  38. device: str = None,
  39. pp_option: PaddlePredictorOption = None,
  40. use_hpip: bool = False,
  41. ) -> None:
  42. """Initializes the layout parsing pipeline.
  43. Args:
  44. config (Dict): Configuration dictionary containing various settings.
  45. device (str, optional): Device to run the predictions on. Defaults to None.
  46. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  47. use_hpip (bool, optional): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
  48. """
  49. super().__init__(
  50. device=device,
  51. pp_option=pp_option,
  52. use_hpip=use_hpip,
  53. )
  54. self.inintial_predictor(config)
  55. self.batch_sampler = ImageBatchSampler(batch_size=1)
  56. self.img_reader = ReadImage(format="BGR")
  57. def inintial_predictor(self, config: dict) -> None:
  58. """Initializes the predictor based on the provided configuration.
  59. Args:
  60. config (Dict): A dictionary containing the configuration for the predictor.
  61. Returns:
  62. None
  63. """
  64. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  65. self.use_general_ocr = config.get("use_general_ocr", True)
  66. self.use_table_recognition = config.get("use_table_recognition", True)
  67. self.use_seal_recognition = config.get("use_seal_recognition", True)
  68. self.use_formula_recognition = config.get(
  69. "use_formula_recognition",
  70. True,
  71. )
  72. if self.use_doc_preprocessor:
  73. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  74. "DocPreprocessor",
  75. {
  76. "pipeline_config_error": "config error for doc_preprocessor_pipeline!",
  77. },
  78. )
  79. self.doc_preprocessor_pipeline = self.create_pipeline(
  80. doc_preprocessor_config,
  81. )
  82. layout_det_config = config.get("SubModules", {}).get(
  83. "LayoutDetection",
  84. {"model_config_error": "config error for layout_det_model!"},
  85. )
  86. layout_kwargs = {}
  87. if (threshold := layout_det_config.get("threshold", None)) is not None:
  88. layout_kwargs["threshold"] = threshold
  89. if (layout_nms := layout_det_config.get("layout_nms", None)) is not None:
  90. layout_kwargs["layout_nms"] = layout_nms
  91. if (
  92. layout_unclip_ratio := layout_det_config.get("layout_unclip_ratio", None)
  93. ) is not None:
  94. layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio
  95. if (
  96. layout_merge_bboxes_mode := layout_det_config.get(
  97. "layout_merge_bboxes_mode", None
  98. )
  99. ) is not None:
  100. layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode
  101. self.layout_det_model = self.create_model(layout_det_config, **layout_kwargs)
  102. if self.use_general_ocr or self.use_table_recognition:
  103. general_ocr_config = config.get("SubPipelines", {}).get(
  104. "GeneralOCR",
  105. {"pipeline_config_error": "config error for general_ocr_pipeline!"},
  106. )
  107. self.general_ocr_pipeline = self.create_pipeline(
  108. general_ocr_config,
  109. )
  110. if self.use_seal_recognition:
  111. seal_recognition_config = config.get("SubPipelines", {}).get(
  112. "SealRecognition",
  113. {
  114. "pipeline_config_error": "config error for seal_recognition_pipeline!",
  115. },
  116. )
  117. self.seal_recognition_pipeline = self.create_pipeline(
  118. seal_recognition_config,
  119. )
  120. if self.use_table_recognition:
  121. table_recognition_config = config.get("SubPipelines", {}).get(
  122. "TableRecognition",
  123. {
  124. "pipeline_config_error": "config error for table_recognition_pipeline!",
  125. },
  126. )
  127. self.table_recognition_pipeline = self.create_pipeline(
  128. table_recognition_config,
  129. )
  130. if self.use_formula_recognition:
  131. formula_recognition_config = config.get("SubPipelines", {}).get(
  132. "FormulaRecognition",
  133. {
  134. "pipeline_config_error": "config error for formula_recognition_pipeline!",
  135. },
  136. )
  137. self.formula_recognition_pipeline = self.create_pipeline(
  138. formula_recognition_config,
  139. )
  140. return
  141. def get_text_paragraphs_ocr_res(
  142. self,
  143. overall_ocr_res: OCRResult,
  144. layout_det_res: DetResult,
  145. ) -> OCRResult:
  146. """
  147. Retrieves the OCR results for text paragraphs, excluding those of formulas, tables, and seals.
  148. Args:
  149. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  150. layout_det_res (DetResult): The detection result containing the layout information of the document.
  151. Returns:
  152. OCRResult: The OCR result for text paragraphs after excluding formulas, tables, and seals.
  153. """
  154. object_boxes = []
  155. for box_info in layout_det_res["boxes"]:
  156. if box_info["label"].lower() in ["formula", "table", "seal"]:
  157. object_boxes.append(box_info["coordinate"])
  158. object_boxes = np.array(object_boxes)
  159. return get_sub_regions_ocr_res(overall_ocr_res, object_boxes, flag_within=False)
  160. def check_model_settings_valid(self, input_params: dict) -> bool:
  161. """
  162. Check if the input parameters are valid based on the initialized models.
  163. Args:
  164. input_params (Dict): A dictionary containing input parameters.
  165. Returns:
  166. bool: True if all required models are initialized according to input parameters, False otherwise.
  167. """
  168. if input_params["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  169. logging.error(
  170. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized.",
  171. )
  172. return False
  173. if input_params["use_general_ocr"] and not self.use_general_ocr:
  174. logging.error(
  175. "Set use_general_ocr, but the models for general OCR are not initialized.",
  176. )
  177. return False
  178. if input_params["use_seal_recognition"] and not self.use_seal_recognition:
  179. logging.error(
  180. "Set use_seal_recognition, but the models for seal recognition are not initialized.",
  181. )
  182. return False
  183. if input_params["use_table_recognition"] and not self.use_table_recognition:
  184. logging.error(
  185. "Set use_table_recognition, but the models for table recognition are not initialized.",
  186. )
  187. return False
  188. return True
  189. def get_model_settings(
  190. self,
  191. use_doc_orientation_classify: bool | None,
  192. use_doc_unwarping: bool | None,
  193. use_general_ocr: bool | None,
  194. use_seal_recognition: bool | None,
  195. use_table_recognition: bool | None,
  196. use_formula_recognition: bool | None,
  197. ) -> dict:
  198. """
  199. Get the model settings based on the provided parameters or default values.
  200. Args:
  201. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  202. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  203. use_general_ocr (Optional[bool]): Whether to use general OCR.
  204. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  205. use_table_recognition (Optional[bool]): Whether to use table recognition.
  206. Returns:
  207. dict: A dictionary containing the model settings.
  208. """
  209. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  210. use_doc_preprocessor = self.use_doc_preprocessor
  211. else:
  212. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  213. use_doc_preprocessor = True
  214. else:
  215. use_doc_preprocessor = False
  216. if use_general_ocr is None:
  217. use_general_ocr = self.use_general_ocr
  218. if use_seal_recognition is None:
  219. use_seal_recognition = self.use_seal_recognition
  220. if use_table_recognition is None:
  221. use_table_recognition = self.use_table_recognition
  222. if use_formula_recognition is None:
  223. use_formula_recognition = self.use_formula_recognition
  224. return dict(
  225. use_doc_preprocessor=use_doc_preprocessor,
  226. use_general_ocr=use_general_ocr,
  227. use_seal_recognition=use_seal_recognition,
  228. use_table_recognition=use_table_recognition,
  229. use_formula_recognition=use_formula_recognition,
  230. )
  231. def predict(
  232. self,
  233. input: Union[str, list[str], np.ndarray, list[np.ndarray]],
  234. use_doc_orientation_classify: bool | None = None,
  235. use_doc_unwarping: bool | None = None,
  236. use_general_ocr: bool | None = None,
  237. use_seal_recognition: bool | None = None,
  238. use_table_recognition: bool | None = None,
  239. use_formula_recognition: bool | None = None,
  240. text_det_limit_side_len: int | None = None,
  241. text_det_limit_type: Union[str, None] = None,
  242. text_det_thresh: float | None = None,
  243. text_det_box_thresh: float | None = None,
  244. text_det_unclip_ratio: float | None = None,
  245. text_rec_score_thresh: float | None = None,
  246. seal_det_limit_side_len: int | None = None,
  247. seal_det_limit_type: Union[str, None] = None,
  248. seal_det_thresh: float | None = None,
  249. seal_det_box_thresh: float | None = None,
  250. seal_det_unclip_ratio: float | None = None,
  251. seal_rec_score_thresh: float | None = None,
  252. **kwargs,
  253. ) -> LayoutParsingResultV2:
  254. """
  255. This function predicts the layout parsing result for the given input.
  256. Args:
  257. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or pdf(s) to be processed.
  258. use_doc_orientation_classify (bool): Whether to use document orientation classification.
  259. use_doc_unwarping (bool): Whether to use document unwarping.
  260. use_general_ocr (bool): Whether to use general OCR.
  261. use_seal_recognition (bool): Whether to use seal recognition.
  262. use_table_recognition (bool): Whether to use table recognition.
  263. **kwargs: Additional keyword arguments.
  264. Returns:
  265. LayoutParsingResultV2: The predicted layout parsing result.
  266. """
  267. model_settings = self.get_model_settings(
  268. use_doc_orientation_classify,
  269. use_doc_unwarping,
  270. use_general_ocr,
  271. use_seal_recognition,
  272. use_table_recognition,
  273. use_formula_recognition,
  274. )
  275. if not self.check_model_settings_valid(model_settings):
  276. yield {"error": "the input params for model settings are invalid!"}
  277. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  278. image_array = self.img_reader(batch_data.instances)[0]
  279. if model_settings["use_doc_preprocessor"]:
  280. doc_preprocessor_res = next(
  281. self.doc_preprocessor_pipeline(
  282. image_array,
  283. use_doc_orientation_classify=use_doc_orientation_classify,
  284. use_doc_unwarping=use_doc_unwarping,
  285. ),
  286. )
  287. else:
  288. doc_preprocessor_res = {"output_img": image_array}
  289. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  290. layout_det_res = next(
  291. self.layout_det_model(doc_preprocessor_image),
  292. )
  293. if model_settings["use_formula_recognition"]:
  294. formula_res_all = next(
  295. self.formula_recognition_pipeline(
  296. doc_preprocessor_image,
  297. use_layout_detection=False,
  298. use_doc_orientation_classify=False,
  299. use_doc_unwarping=False,
  300. layout_det_res=layout_det_res,
  301. ),
  302. )
  303. formula_res_list = formula_res_all["formula_res_list"]
  304. else:
  305. formula_res_list = []
  306. for formula_res in formula_res_list:
  307. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  308. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = 255.0
  309. if (
  310. model_settings["use_general_ocr"]
  311. or model_settings["use_table_recognition"]
  312. ):
  313. overall_ocr_res = next(
  314. self.general_ocr_pipeline(
  315. doc_preprocessor_image,
  316. text_det_limit_side_len=text_det_limit_side_len,
  317. text_det_limit_type=text_det_limit_type,
  318. text_det_thresh=text_det_thresh,
  319. text_det_box_thresh=text_det_box_thresh,
  320. text_det_unclip_ratio=text_det_unclip_ratio,
  321. text_rec_score_thresh=text_rec_score_thresh,
  322. ),
  323. )
  324. for formula_res in formula_res_list:
  325. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  326. poly_points = [
  327. (x_min, y_min),
  328. (x_max, y_min),
  329. (x_max, y_max),
  330. (x_min, y_max),
  331. ]
  332. overall_ocr_res["dt_polys"].append(poly_points)
  333. overall_ocr_res["rec_texts"].append(
  334. f"${formula_res['rec_formula']}$"
  335. )
  336. overall_ocr_res["rec_boxes"] = np.vstack(
  337. (overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
  338. )
  339. overall_ocr_res["rec_polys"].append(poly_points)
  340. overall_ocr_res["rec_scores"].append(1)
  341. else:
  342. overall_ocr_res = {}
  343. if model_settings["use_general_ocr"]:
  344. text_paragraphs_ocr_res = self.get_text_paragraphs_ocr_res(
  345. overall_ocr_res,
  346. layout_det_res,
  347. )
  348. else:
  349. text_paragraphs_ocr_res = {}
  350. if model_settings["use_table_recognition"]:
  351. table_res_all = next(
  352. self.table_recognition_pipeline(
  353. doc_preprocessor_image,
  354. use_doc_orientation_classify=False,
  355. use_doc_unwarping=False,
  356. use_layout_detection=False,
  357. use_ocr_model=False,
  358. overall_ocr_res=overall_ocr_res,
  359. layout_det_res=layout_det_res,
  360. cell_sort_by_y_projection=True,
  361. ),
  362. )
  363. table_res_list = table_res_all["table_res_list"]
  364. else:
  365. table_res_list = []
  366. if model_settings["use_seal_recognition"]:
  367. seal_res_all = next(
  368. self.seal_recognition_pipeline(
  369. doc_preprocessor_image,
  370. use_doc_orientation_classify=False,
  371. use_doc_unwarping=False,
  372. use_layout_detection=False,
  373. layout_det_res=layout_det_res,
  374. seal_det_limit_side_len=seal_det_limit_side_len,
  375. seal_det_limit_type=seal_det_limit_type,
  376. seal_det_thresh=seal_det_thresh,
  377. seal_det_box_thresh=seal_det_box_thresh,
  378. seal_det_unclip_ratio=seal_det_unclip_ratio,
  379. seal_rec_score_thresh=seal_rec_score_thresh,
  380. ),
  381. )
  382. seal_res_list = seal_res_all["seal_res_list"]
  383. else:
  384. seal_res_list = []
  385. for formula_res in formula_res_list:
  386. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  387. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
  388. "input_img"
  389. ]
  390. structure_res = get_structure_res(
  391. overall_ocr_res,
  392. layout_det_res,
  393. table_res_list,
  394. )
  395. structure_res_list = [
  396. {
  397. "block_bbox": [0, 0, 2550, 2550],
  398. "block_size": [image_array.shape[1], image_array.shape[0]],
  399. "sub_blocks": structure_res,
  400. },
  401. ]
  402. single_img_res = {
  403. "input_path": batch_data.input_paths[0],
  404. "page_index": batch_data.page_indexes[0],
  405. "doc_preprocessor_res": doc_preprocessor_res,
  406. "layout_det_res": layout_det_res,
  407. "overall_ocr_res": overall_ocr_res,
  408. "text_paragraphs_ocr_res": text_paragraphs_ocr_res,
  409. "table_res_list": table_res_list,
  410. "seal_res_list": seal_res_list,
  411. "formula_res_list": formula_res_list,
  412. "layout_parsing_result": structure_res_list,
  413. "model_settings": model_settings,
  414. }
  415. yield LayoutParsingResultV2(single_img_res)