pipeline.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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 typing import Any, Dict, List, Optional, Tuple, Union
  15. import numpy as np
  16. from ....utils import logging
  17. from ....utils.deps import pipeline_requires_extra
  18. from ...common.batch_sampler import ImageBatchSampler
  19. from ...common.reader import ReadImage
  20. from ...models.object_detection.result import DetResult
  21. from ...utils.hpi import HPIConfig
  22. from ...utils.pp_option import PaddlePredictorOption
  23. from .._parallel import AutoParallelImageSimpleInferencePipeline
  24. from ..base import BasePipeline
  25. from ..components import CropByBoxes
  26. from ..ocr.result import OCRResult
  27. from .result import LayoutParsingResult
  28. from .utils import get_sub_regions_ocr_res, sorted_layout_boxes
  29. class _LayoutParsingPipeline(BasePipeline):
  30. """Layout Parsing Pipeline"""
  31. def __init__(
  32. self,
  33. config: Dict,
  34. device: str = None,
  35. pp_option: PaddlePredictorOption = None,
  36. use_hpip: bool = False,
  37. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  38. ) -> None:
  39. """Initializes the layout parsing pipeline.
  40. Args:
  41. config (Dict): Configuration dictionary containing various settings.
  42. device (str, optional): Device to run the predictions on. Defaults to None.
  43. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  44. use_hpip (bool, optional): Whether to use the high-performance
  45. inference plugin (HPIP) by default. Defaults to False.
  46. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  47. The default high-performance inference configuration dictionary.
  48. Defaults to None.
  49. """
  50. super().__init__(
  51. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  52. )
  53. self.inintial_predictor(config)
  54. self.batch_sampler = ImageBatchSampler(batch_size=1)
  55. self.img_reader = ReadImage(format="BGR")
  56. self._crop_by_boxes = CropByBoxes()
  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("use_formula_recognition", True)
  69. if self.use_doc_preprocessor:
  70. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  71. "DocPreprocessor",
  72. {
  73. "pipeline_config_error": "config error for doc_preprocessor_pipeline!"
  74. },
  75. )
  76. self.doc_preprocessor_pipeline = self.create_pipeline(
  77. doc_preprocessor_config
  78. )
  79. layout_det_config = config.get("SubModules", {}).get(
  80. "LayoutDetection",
  81. {"model_config_error": "config error for layout_det_model!"},
  82. )
  83. layout_kwargs = {}
  84. if (threshold := layout_det_config.get("threshold", None)) is not None:
  85. layout_kwargs["threshold"] = threshold
  86. if (layout_nms := layout_det_config.get("layout_nms", None)) is not None:
  87. layout_kwargs["layout_nms"] = layout_nms
  88. if (
  89. layout_unclip_ratio := layout_det_config.get("layout_unclip_ratio", None)
  90. ) is not None:
  91. layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio
  92. if (
  93. layout_merge_bboxes_mode := layout_det_config.get(
  94. "layout_merge_bboxes_mode", None
  95. )
  96. ) is not None:
  97. layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode
  98. self.layout_det_model = self.create_model(layout_det_config, **layout_kwargs)
  99. if self.use_general_ocr or self.use_table_recognition:
  100. general_ocr_config = config.get("SubPipelines", {}).get(
  101. "GeneralOCR",
  102. {"pipeline_config_error": "config error for general_ocr_pipeline!"},
  103. )
  104. self.general_ocr_pipeline = self.create_pipeline(general_ocr_config)
  105. if self.use_seal_recognition:
  106. seal_recognition_config = config.get("SubPipelines", {}).get(
  107. "SealRecognition",
  108. {
  109. "pipeline_config_error": "config error for seal_recognition_pipeline!"
  110. },
  111. )
  112. self.seal_recognition_pipeline = self.create_pipeline(
  113. seal_recognition_config
  114. )
  115. if self.use_table_recognition:
  116. table_recognition_config = config.get("SubPipelines", {}).get(
  117. "TableRecognition",
  118. {
  119. "pipeline_config_error": "config error for table_recognition_pipeline!"
  120. },
  121. )
  122. self.table_recognition_pipeline = self.create_pipeline(
  123. table_recognition_config
  124. )
  125. if self.use_formula_recognition:
  126. formula_recognition_config = config.get("SubPipelines", {}).get(
  127. "FormulaRecognition",
  128. {
  129. "pipeline_config_error": "config error for formula_recognition_pipeline!"
  130. },
  131. )
  132. self.formula_recognition_pipeline = self.create_pipeline(
  133. formula_recognition_config
  134. )
  135. return
  136. def get_layout_parsing_res(
  137. self,
  138. image: list,
  139. layout_det_res: DetResult,
  140. overall_ocr_res: OCRResult,
  141. table_res_list: list,
  142. seal_res_list: list,
  143. formula_res_list: list,
  144. text_det_limit_side_len: Optional[int] = None,
  145. text_det_limit_type: Optional[str] = None,
  146. text_det_thresh: Optional[float] = None,
  147. text_det_box_thresh: Optional[float] = None,
  148. text_det_unclip_ratio: Optional[float] = None,
  149. text_rec_score_thresh: Optional[float] = None,
  150. ) -> list:
  151. """
  152. Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
  153. Args:
  154. image (list): The input image.
  155. layout_det_res (DetResult): The detection result containing the layout information of the document.
  156. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  157. table_res_list (list): A list of table recognition results.
  158. seal_res_list (list): A list of seal recognition results.
  159. formula_res_list (list): A list of formula recognition results.
  160. text_det_limit_side_len (Optional[int], optional): The maximum side length of the text detection region. Defaults to None.
  161. text_det_limit_type (Optional[str], optional): The type of limit for the text detection region. Defaults to None.
  162. text_det_thresh (Optional[float], optional): The confidence threshold for text detection. Defaults to None.
  163. text_det_box_thresh (Optional[float], optional): The confidence threshold for text detection bounding boxes. Defaults to None
  164. text_det_unclip_ratio (Optional[float], optional): The unclip ratio for text detection. Defaults to None.
  165. text_rec_score_thresh (Optional[float], optional): The score threshold for text recognition. Defaults to None.
  166. Returns:
  167. list: A list of dictionaries representing the layout parsing result.
  168. """
  169. layout_parsing_res = []
  170. matched_ocr_dict = {}
  171. formula_index = 0
  172. table_index = 0
  173. seal_index = 0
  174. image = np.array(image)
  175. object_boxes = []
  176. for object_box_idx, box_info in enumerate(layout_det_res["boxes"]):
  177. single_box_res = {}
  178. box = box_info["coordinate"]
  179. label = box_info["label"].lower()
  180. single_box_res["block_bbox"] = box
  181. single_box_res["block_label"] = label
  182. single_box_res["block_content"] = ""
  183. object_boxes.append(box)
  184. if label == "formula":
  185. if len(formula_res_list) > 0:
  186. assert (
  187. len(formula_res_list) > formula_index
  188. ), f"The number of \
  189. formula regions of layout parsing pipeline \
  190. and formula recognition pipeline are different!"
  191. single_box_res["block_content"] = formula_res_list[formula_index][
  192. "rec_formula"
  193. ]
  194. formula_index += 1
  195. elif label == "table":
  196. if len(table_res_list) > 0:
  197. assert (
  198. len(table_res_list) > table_index
  199. ), f"The number of \
  200. table regions of layout parsing pipeline \
  201. and table recognition pipeline are different!"
  202. single_box_res["block_content"] = table_res_list[table_index][
  203. "pred_html"
  204. ]
  205. table_index += 1
  206. elif label == "seal":
  207. if len(seal_res_list) > 0:
  208. assert (
  209. len(seal_res_list) > seal_index
  210. ), f"The number of \
  211. seal regions of layout parsing pipeline \
  212. and seal recognition pipeline are different!"
  213. single_box_res["block_content"] = ", ".join(
  214. seal_res_list[seal_index]["rec_texts"]
  215. )
  216. seal_index += 1
  217. else:
  218. ocr_res_in_box, matched_idxes = get_sub_regions_ocr_res(
  219. overall_ocr_res, [box], return_match_idx=True
  220. )
  221. for matched_idx in matched_idxes:
  222. if matched_ocr_dict.get(matched_idx, None) is None:
  223. matched_ocr_dict[matched_idx] = [object_box_idx]
  224. else:
  225. matched_ocr_dict[matched_idx].append(object_box_idx)
  226. single_box_res["block_content"] = "\n".join(ocr_res_in_box["rec_texts"])
  227. layout_parsing_res.append(single_box_res)
  228. for layout_box_ids in matched_ocr_dict.values():
  229. # one ocr is matched to multiple layout boxes, split the text into multiple lines
  230. if len(layout_box_ids) > 1:
  231. for idx in layout_box_ids:
  232. wht_im = np.ones(image.shape, dtype=image.dtype) * 255
  233. box = layout_parsing_res[idx]["block_bbox"]
  234. x1, y1, x2, y2 = [int(i) for i in box]
  235. wht_im[y1:y2, x1:x2, :] = image[y1:y2, x1:x2, :]
  236. sub_ocr_res = next(
  237. self.general_ocr_pipeline(
  238. wht_im,
  239. text_det_limit_side_len=text_det_limit_side_len,
  240. text_det_limit_type=text_det_limit_type,
  241. text_det_thresh=text_det_thresh,
  242. text_det_box_thresh=text_det_box_thresh,
  243. text_det_unclip_ratio=text_det_unclip_ratio,
  244. text_rec_score_thresh=text_rec_score_thresh,
  245. )
  246. )
  247. layout_parsing_res[idx]["block_content"] = "\n".join(
  248. sub_ocr_res["rec_texts"]
  249. )
  250. ocr_without_layout_boxes = get_sub_regions_ocr_res(
  251. overall_ocr_res, object_boxes, flag_within=False
  252. )
  253. for ocr_rec_box, ocr_rec_text in zip(
  254. ocr_without_layout_boxes["rec_boxes"], ocr_without_layout_boxes["rec_texts"]
  255. ):
  256. single_box_res = {}
  257. single_box_res["block_bbox"] = ocr_rec_box
  258. single_box_res["block_label"] = "other_text"
  259. single_box_res["block_content"] = ocr_rec_text
  260. layout_parsing_res.append(single_box_res)
  261. layout_parsing_res = sorted_layout_boxes(layout_parsing_res, w=image.shape[1])
  262. return layout_parsing_res
  263. def check_model_settings_valid(self, input_params: Dict) -> bool:
  264. """
  265. Check if the input parameters are valid based on the initialized models.
  266. Args:
  267. input_params (Dict): A dictionary containing input parameters.
  268. Returns:
  269. bool: True if all required models are initialized according to input parameters, False otherwise.
  270. """
  271. if input_params["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  272. logging.error(
  273. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized."
  274. )
  275. return False
  276. if input_params["use_general_ocr"] and not self.use_general_ocr:
  277. logging.error(
  278. "Set use_general_ocr, but the models for general OCR are not initialized."
  279. )
  280. return False
  281. if input_params["use_seal_recognition"] and not self.use_seal_recognition:
  282. logging.error(
  283. "Set use_seal_recognition, but the models for seal recognition are not initialized."
  284. )
  285. return False
  286. if input_params["use_table_recognition"] and not self.use_table_recognition:
  287. logging.error(
  288. "Set use_table_recognition, but the models for table recognition are not initialized."
  289. )
  290. return False
  291. return True
  292. def get_model_settings(
  293. self,
  294. use_doc_orientation_classify: Optional[bool],
  295. use_doc_unwarping: Optional[bool],
  296. use_general_ocr: Optional[bool],
  297. use_seal_recognition: Optional[bool],
  298. use_table_recognition: Optional[bool],
  299. use_formula_recognition: Optional[bool],
  300. ) -> dict:
  301. """
  302. Get the model settings based on the provided parameters or default values.
  303. Args:
  304. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  305. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  306. use_general_ocr (Optional[bool]): Whether to use general OCR.
  307. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  308. use_table_recognition (Optional[bool]): Whether to use table recognition.
  309. Returns:
  310. dict: A dictionary containing the model settings.
  311. """
  312. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  313. use_doc_preprocessor = self.use_doc_preprocessor
  314. else:
  315. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  316. use_doc_preprocessor = True
  317. else:
  318. use_doc_preprocessor = False
  319. if use_general_ocr is None:
  320. use_general_ocr = self.use_general_ocr
  321. if use_seal_recognition is None:
  322. use_seal_recognition = self.use_seal_recognition
  323. if use_table_recognition is None:
  324. use_table_recognition = self.use_table_recognition
  325. if use_formula_recognition is None:
  326. use_formula_recognition = self.use_formula_recognition
  327. return dict(
  328. use_doc_preprocessor=use_doc_preprocessor,
  329. use_general_ocr=use_general_ocr,
  330. use_seal_recognition=use_seal_recognition,
  331. use_table_recognition=use_table_recognition,
  332. use_formula_recognition=use_formula_recognition,
  333. )
  334. def predict(
  335. self,
  336. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  337. use_doc_orientation_classify: Optional[bool] = None,
  338. use_doc_unwarping: Optional[bool] = None,
  339. use_textline_orientation: Optional[bool] = None,
  340. use_general_ocr: Optional[bool] = None,
  341. use_seal_recognition: Optional[bool] = None,
  342. use_table_recognition: Optional[bool] = None,
  343. use_formula_recognition: Optional[bool] = None,
  344. layout_threshold: Optional[Union[float, dict]] = None,
  345. layout_nms: Optional[bool] = None,
  346. layout_unclip_ratio: Optional[Union[float, Tuple[float, float], dict]] = None,
  347. layout_merge_bboxes_mode: Optional[str] = None,
  348. text_det_limit_side_len: Optional[int] = None,
  349. text_det_limit_type: Optional[str] = None,
  350. text_det_thresh: Optional[float] = None,
  351. text_det_box_thresh: Optional[float] = None,
  352. text_det_unclip_ratio: Optional[float] = None,
  353. text_rec_score_thresh: Optional[float] = None,
  354. seal_det_limit_side_len: Optional[int] = None,
  355. seal_det_limit_type: Optional[str] = None,
  356. seal_det_thresh: Optional[float] = None,
  357. seal_det_box_thresh: Optional[float] = None,
  358. seal_det_unclip_ratio: Optional[float] = None,
  359. seal_rec_score_thresh: Optional[float] = None,
  360. **kwargs,
  361. ) -> LayoutParsingResult:
  362. """
  363. This function predicts the layout parsing result for the given input.
  364. Args:
  365. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): The input image(s) or pdf(s) to be processed.
  366. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  367. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  368. use_textline_orientation (Optional[bool]): Whether to use textline orientation prediction.
  369. use_general_ocr (Optional[bool]): Whether to use general OCR.
  370. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  371. use_table_recognition (Optional[bool]): Whether to use table recognition.
  372. use_formula_recognition (Optional[bool]): Whether to use formula recognition.
  373. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
  374. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False.
  375. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box.
  376. Defaults to None.
  377. If it's a single number, then both width and height are used.
  378. If it's a tuple of two numbers, then they are used separately for width and height respectively.
  379. If it's None, then no unclipping will be performed.
  380. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None.
  381. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  382. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  383. text_det_thresh (Optional[float]): Threshold for text detection.
  384. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  385. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  386. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  387. seal_det_limit_side_len (Optional[int]): Maximum side length for seal detection.
  388. seal_det_limit_type (Optional[str]): Type of limit to apply for seal detection.
  389. seal_det_thresh (Optional[float]): Threshold for seal detection.
  390. seal_det_box_thresh (Optional[float]): Threshold for seal detection boxes.
  391. seal_det_unclip_ratio (Optional[float]): Ratio for unclipping seal detection boxes.
  392. seal_rec_score_thresh (Optional[float]): Score threshold for seal recognition.
  393. **kwargs: Additional keyword arguments.
  394. Returns:
  395. LayoutParsingResult: The predicted layout parsing result.
  396. """
  397. model_settings = self.get_model_settings(
  398. use_doc_orientation_classify,
  399. use_doc_unwarping,
  400. use_general_ocr,
  401. use_seal_recognition,
  402. use_table_recognition,
  403. use_formula_recognition,
  404. )
  405. if not self.check_model_settings_valid(model_settings):
  406. yield {"error": "the input params for model settings are invalid!"}
  407. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  408. image_array = self.img_reader(batch_data.instances)[0]
  409. if model_settings["use_doc_preprocessor"]:
  410. doc_preprocessor_res = next(
  411. self.doc_preprocessor_pipeline(
  412. image_array,
  413. use_doc_orientation_classify=use_doc_orientation_classify,
  414. use_doc_unwarping=use_doc_unwarping,
  415. )
  416. )
  417. else:
  418. doc_preprocessor_res = {"output_img": image_array}
  419. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  420. layout_det_res = next(
  421. self.layout_det_model(
  422. doc_preprocessor_image,
  423. threshold=layout_threshold,
  424. layout_nms=layout_nms,
  425. layout_unclip_ratio=layout_unclip_ratio,
  426. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  427. )
  428. )
  429. if (
  430. model_settings["use_general_ocr"]
  431. or model_settings["use_table_recognition"]
  432. ):
  433. overall_ocr_res = next(
  434. self.general_ocr_pipeline(
  435. doc_preprocessor_image,
  436. use_textline_orientation=use_textline_orientation,
  437. text_det_limit_side_len=text_det_limit_side_len,
  438. text_det_limit_type=text_det_limit_type,
  439. text_det_thresh=text_det_thresh,
  440. text_det_box_thresh=text_det_box_thresh,
  441. text_det_unclip_ratio=text_det_unclip_ratio,
  442. text_rec_score_thresh=text_rec_score_thresh,
  443. )
  444. )
  445. else:
  446. overall_ocr_res = {}
  447. if model_settings["use_table_recognition"]:
  448. table_res_all = next(
  449. self.table_recognition_pipeline(
  450. doc_preprocessor_image,
  451. use_doc_orientation_classify=False,
  452. use_doc_unwarping=False,
  453. use_layout_detection=False,
  454. use_ocr_model=False,
  455. overall_ocr_res=overall_ocr_res,
  456. layout_det_res=layout_det_res,
  457. )
  458. )
  459. table_res_list = table_res_all["table_res_list"]
  460. else:
  461. table_res_list = []
  462. if model_settings["use_seal_recognition"]:
  463. seal_res_all = next(
  464. self.seal_recognition_pipeline(
  465. doc_preprocessor_image,
  466. use_doc_orientation_classify=False,
  467. use_doc_unwarping=False,
  468. use_layout_detection=False,
  469. layout_det_res=layout_det_res,
  470. seal_det_limit_side_len=seal_det_limit_side_len,
  471. seal_det_limit_type=seal_det_limit_type,
  472. seal_det_thresh=seal_det_thresh,
  473. seal_det_box_thresh=seal_det_box_thresh,
  474. seal_det_unclip_ratio=seal_det_unclip_ratio,
  475. seal_rec_score_thresh=seal_rec_score_thresh,
  476. )
  477. )
  478. seal_res_list = seal_res_all["seal_res_list"]
  479. else:
  480. seal_res_list = []
  481. if model_settings["use_formula_recognition"]:
  482. formula_res_all = next(
  483. self.formula_recognition_pipeline(
  484. doc_preprocessor_image,
  485. use_layout_detection=False,
  486. use_doc_orientation_classify=False,
  487. use_doc_unwarping=False,
  488. layout_det_res=layout_det_res,
  489. )
  490. )
  491. formula_res_list = formula_res_all["formula_res_list"]
  492. else:
  493. formula_res_list = []
  494. parsing_res_list = self.get_layout_parsing_res(
  495. doc_preprocessor_image,
  496. layout_det_res=layout_det_res,
  497. overall_ocr_res=overall_ocr_res,
  498. table_res_list=table_res_list,
  499. seal_res_list=seal_res_list,
  500. formula_res_list=formula_res_list,
  501. text_det_limit_side_len=text_det_limit_side_len,
  502. text_det_limit_type=text_det_limit_type,
  503. text_det_thresh=text_det_thresh,
  504. text_det_box_thresh=text_det_box_thresh,
  505. text_det_unclip_ratio=text_det_unclip_ratio,
  506. text_rec_score_thresh=text_rec_score_thresh,
  507. )
  508. single_img_res = {
  509. "input_path": batch_data.input_paths[0],
  510. "page_index": batch_data.page_indexes[0],
  511. "doc_preprocessor_res": doc_preprocessor_res,
  512. "layout_det_res": layout_det_res,
  513. "overall_ocr_res": overall_ocr_res,
  514. "table_res_list": table_res_list,
  515. "seal_res_list": seal_res_list,
  516. "formula_res_list": formula_res_list,
  517. "parsing_res_list": parsing_res_list,
  518. "model_settings": model_settings,
  519. }
  520. yield LayoutParsingResult(single_img_res)
  521. @pipeline_requires_extra("ocr")
  522. class LayoutParsingPipeline(AutoParallelImageSimpleInferencePipeline):
  523. entities = ["layout_parsing"]
  524. @property
  525. def _pipeline_cls(self):
  526. return _LayoutParsingPipeline
  527. def _get_batch_size(self, config):
  528. return 1