pipeline_v2.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683
  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. from typing import Optional, Union, Tuple, Iterator
  16. import numpy as np
  17. import re
  18. import copy
  19. from ....utils import logging
  20. from ...common.batch_sampler import ImageBatchSampler
  21. from ...common.reader import ReadImage
  22. from ...models.object_detection.result import DetResult
  23. from ...utils.pp_option import PaddlePredictorOption
  24. from ..base import BasePipeline
  25. from ..ocr.result import OCRResult
  26. from .result_v2 import LayoutParsingResultV2
  27. from .utils import get_single_block_parsing_res
  28. from .utils import get_sub_regions_ocr_res
  29. class LayoutParsingPipelineV2(BasePipeline):
  30. """Layout Parsing Pipeline V2"""
  31. entities = ["PP-StructureV3"]
  32. def __init__(
  33. self,
  34. config: dict,
  35. device: str = None,
  36. pp_option: PaddlePredictorOption = None,
  37. use_hpip: bool = False,
  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 high-performance inference (hpip) for prediction. Defaults to False.
  45. """
  46. super().__init__(
  47. device=device,
  48. pp_option=pp_option,
  49. use_hpip=use_hpip,
  50. )
  51. self.inintial_predictor(config)
  52. self.batch_sampler = ImageBatchSampler(batch_size=1)
  53. self.img_reader = ReadImage(format="BGR")
  54. def inintial_predictor(self, config: dict) -> None:
  55. """Initializes the predictor based on the provided configuration.
  56. Args:
  57. config (Dict): A dictionary containing the configuration for the predictor.
  58. Returns:
  59. None
  60. """
  61. self.use_doc_preprocessor = config.get("use_doc_preprocessor", True)
  62. self.use_general_ocr = config.get("use_general_ocr", True)
  63. self.use_table_recognition = config.get("use_table_recognition", True)
  64. self.use_seal_recognition = config.get("use_seal_recognition", True)
  65. self.use_formula_recognition = config.get(
  66. "use_formula_recognition",
  67. True,
  68. )
  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(
  105. general_ocr_config,
  106. )
  107. if self.use_seal_recognition:
  108. seal_recognition_config = config.get("SubPipelines", {}).get(
  109. "SealRecognition",
  110. {
  111. "pipeline_config_error": "config error for seal_recognition_pipeline!",
  112. },
  113. )
  114. self.seal_recognition_pipeline = self.create_pipeline(
  115. seal_recognition_config,
  116. )
  117. if self.use_table_recognition:
  118. table_recognition_config = config.get("SubPipelines", {}).get(
  119. "TableRecognition",
  120. {
  121. "pipeline_config_error": "config error for table_recognition_pipeline!",
  122. },
  123. )
  124. self.table_recognition_pipeline = self.create_pipeline(
  125. table_recognition_config,
  126. )
  127. if self.use_formula_recognition:
  128. formula_recognition_config = config.get("SubPipelines", {}).get(
  129. "FormulaRecognition",
  130. {
  131. "pipeline_config_error": "config error for formula_recognition_pipeline!",
  132. },
  133. )
  134. self.formula_recognition_pipeline = self.create_pipeline(
  135. formula_recognition_config,
  136. )
  137. return
  138. def get_text_paragraphs_ocr_res(
  139. self,
  140. overall_ocr_res: OCRResult,
  141. layout_det_res: DetResult,
  142. ) -> OCRResult:
  143. """
  144. Retrieves the OCR results for text paragraphs, excluding those of formulas, tables, and seals.
  145. Args:
  146. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  147. layout_det_res (DetResult): The detection result containing the layout information of the document.
  148. Returns:
  149. OCRResult: The OCR result for text paragraphs after excluding formulas, tables, and seals.
  150. """
  151. object_boxes = []
  152. for box_info in layout_det_res["boxes"]:
  153. if box_info["label"].lower() in ["formula", "table", "seal"]:
  154. object_boxes.append(box_info["coordinate"])
  155. object_boxes = np.array(object_boxes)
  156. sub_regions_ocr_res = get_sub_regions_ocr_res(
  157. overall_ocr_res, object_boxes, flag_within=False
  158. )
  159. return sub_regions_ocr_res
  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_layout_parsing_res(
  190. self,
  191. image: list,
  192. layout_det_res: DetResult,
  193. overall_ocr_res: OCRResult,
  194. table_res_list: list,
  195. seal_res_list: list,
  196. formula_res_list: list,
  197. text_det_limit_side_len: Optional[int] = None,
  198. text_det_limit_type: Optional[str] = None,
  199. text_det_thresh: Optional[float] = None,
  200. text_det_box_thresh: Optional[float] = None,
  201. text_det_unclip_ratio: Optional[float] = None,
  202. text_rec_score_thresh: Optional[float] = None,
  203. ) -> list:
  204. """
  205. Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
  206. Args:
  207. image (list): The input image.
  208. layout_det_res (DetResult): The detection result containing the layout information of the document.
  209. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  210. table_res_list (list): A list of table recognition results.
  211. seal_res_list (list): A list of seal recognition results.
  212. formula_res_list (list): A list of formula recognition results.
  213. text_det_limit_side_len (Optional[int], optional): The maximum side length of the text detection region. Defaults to None.
  214. text_det_limit_type (Optional[str], optional): The type of limit for the text detection region. Defaults to None.
  215. text_det_thresh (Optional[float], optional): The confidence threshold for text detection. Defaults to None.
  216. text_det_box_thresh (Optional[float], optional): The confidence threshold for text detection bounding boxes. Defaults to None
  217. text_det_unclip_ratio (Optional[float], optional): The unclip ratio for text detection. Defaults to None.
  218. text_rec_score_thresh (Optional[float], optional): The score threshold for text recognition. Defaults to None.
  219. Returns:
  220. list: A list of dictionaries representing the layout parsing result.
  221. """
  222. matched_ocr_dict = {}
  223. image = np.array(image)
  224. object_boxes = []
  225. for object_box_idx, box_info in enumerate(layout_det_res["boxes"]):
  226. box = box_info["coordinate"]
  227. label = box_info["label"].lower()
  228. object_boxes.append(box)
  229. if label not in ["formula", "table", "seal"]:
  230. _, matched_idxs = get_sub_regions_ocr_res(
  231. overall_ocr_res, [box], return_match_idx=True
  232. )
  233. for matched_idx in matched_idxs:
  234. if matched_ocr_dict.get(matched_idx, None) is None:
  235. matched_ocr_dict[matched_idx] = [object_box_idx]
  236. else:
  237. matched_ocr_dict[matched_idx].append(object_box_idx)
  238. already_processed = set()
  239. for matched_idx, layout_box_ids in matched_ocr_dict.items():
  240. if len(layout_box_ids) <= 1:
  241. continue
  242. # one ocr is matched to multiple layout boxes, split the text into multiple lines
  243. for idx in layout_box_ids:
  244. if idx in already_processed:
  245. continue
  246. already_processed.add(idx)
  247. wht_im = np.ones(image.shape, dtype=image.dtype) * 255
  248. box = object_boxes[idx]
  249. x1, y1, x2, y2 = [int(i) for i in box]
  250. wht_im[y1:y2, x1:x2, :] = image[y1:y2, x1:x2, :]
  251. sub_ocr_res = next(
  252. self.general_ocr_pipeline(
  253. wht_im,
  254. text_det_limit_side_len=text_det_limit_side_len,
  255. text_det_limit_type=text_det_limit_type,
  256. text_det_thresh=text_det_thresh,
  257. text_det_box_thresh=text_det_box_thresh,
  258. text_det_unclip_ratio=text_det_unclip_ratio,
  259. text_rec_score_thresh=text_rec_score_thresh,
  260. )
  261. )
  262. _, matched_idxs = get_sub_regions_ocr_res(
  263. overall_ocr_res, [box], return_match_idx=True
  264. )
  265. for matched_idx in sorted(matched_idxs, reverse=True):
  266. del overall_ocr_res["dt_polys"][matched_idx]
  267. del overall_ocr_res["rec_texts"][matched_idx]
  268. overall_ocr_res["rec_boxes"] = np.delete(
  269. overall_ocr_res["rec_boxes"], matched_idx, axis=0
  270. )
  271. del overall_ocr_res["rec_polys"][matched_idx]
  272. del overall_ocr_res["rec_scores"][matched_idx]
  273. if sub_ocr_res["rec_boxes"] is not []:
  274. overall_ocr_res["dt_polys"].extend(sub_ocr_res["dt_polys"])
  275. overall_ocr_res["rec_texts"].extend(sub_ocr_res["rec_texts"])
  276. overall_ocr_res["rec_boxes"] = np.concatenate(
  277. [overall_ocr_res["rec_boxes"], sub_ocr_res["rec_boxes"]], axis=0
  278. )
  279. overall_ocr_res["rec_polys"].extend(sub_ocr_res["rec_polys"])
  280. overall_ocr_res["rec_scores"].extend(sub_ocr_res["rec_scores"])
  281. for formula_res in formula_res_list:
  282. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  283. poly_points = [
  284. (x_min, y_min),
  285. (x_max, y_min),
  286. (x_max, y_max),
  287. (x_min, y_max),
  288. ]
  289. overall_ocr_res["dt_polys"].append(poly_points)
  290. overall_ocr_res["rec_texts"].append(f"${formula_res['rec_formula']}$")
  291. overall_ocr_res["rec_boxes"] = np.vstack(
  292. (overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
  293. )
  294. overall_ocr_res["rec_polys"].append(poly_points)
  295. overall_ocr_res["rec_scores"].append(1)
  296. parsing_res_list = get_single_block_parsing_res(
  297. overall_ocr_res=overall_ocr_res,
  298. layout_det_res=layout_det_res,
  299. table_res_list=table_res_list,
  300. seal_res_list=seal_res_list,
  301. )
  302. return parsing_res_list
  303. def get_model_settings(
  304. self,
  305. use_doc_orientation_classify: Union[bool, None],
  306. use_doc_unwarping: Union[bool, None],
  307. use_general_ocr: Union[bool, None],
  308. use_seal_recognition: Union[bool, None],
  309. use_table_recognition: Union[bool, None],
  310. use_formula_recognition: Union[bool, None],
  311. ) -> dict:
  312. """
  313. Get the model settings based on the provided parameters or default values.
  314. Args:
  315. use_doc_orientation_classify (Union[bool, None]): Enables document orientation classification if True. Defaults to system setting if None.
  316. use_doc_unwarping (Union[bool, None]): Enables document unwarping if True. Defaults to system setting if None.
  317. use_general_ocr (Union[bool, None]): Enables general OCR if True. Defaults to system setting if None.
  318. use_seal_recognition (Union[bool, None]): Enables seal recognition if True. Defaults to system setting if None.
  319. use_table_recognition (Union[bool, None]): Enables table recognition if True. Defaults to system setting if None.
  320. use_formula_recognition (Union[bool, None]): Enables formula recognition if True. Defaults to system setting if None.
  321. Returns:
  322. dict: A dictionary containing the model settings.
  323. """
  324. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  325. use_doc_preprocessor = self.use_doc_preprocessor
  326. else:
  327. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  328. use_doc_preprocessor = True
  329. else:
  330. use_doc_preprocessor = False
  331. if use_general_ocr is None:
  332. use_general_ocr = self.use_general_ocr
  333. if use_seal_recognition is None:
  334. use_seal_recognition = self.use_seal_recognition
  335. if use_table_recognition is None:
  336. use_table_recognition = self.use_table_recognition
  337. if use_formula_recognition is None:
  338. use_formula_recognition = self.use_formula_recognition
  339. return dict(
  340. use_doc_preprocessor=use_doc_preprocessor,
  341. use_general_ocr=use_general_ocr,
  342. use_seal_recognition=use_seal_recognition,
  343. use_table_recognition=use_table_recognition,
  344. use_formula_recognition=use_formula_recognition,
  345. )
  346. def predict(
  347. self,
  348. input: Union[str, list[str], np.ndarray, list[np.ndarray]],
  349. use_doc_orientation_classify: Union[bool, None] = None,
  350. use_doc_unwarping: Union[bool, None] = None,
  351. use_textline_orientation: Optional[bool] = None,
  352. use_general_ocr: Union[bool, None] = None,
  353. use_seal_recognition: Union[bool, None] = None,
  354. use_table_recognition: Union[bool, None] = None,
  355. use_formula_recognition: Union[bool, None] = None,
  356. layout_threshold: Optional[Union[float, dict]] = None,
  357. layout_nms: Optional[bool] = None,
  358. layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
  359. layout_merge_bboxes_mode: Optional[str] = None,
  360. text_det_limit_side_len: Union[int, None] = None,
  361. text_det_limit_type: Union[str, None] = None,
  362. text_det_thresh: Union[float, None] = None,
  363. text_det_box_thresh: Union[float, None] = None,
  364. text_det_unclip_ratio: Union[float, None] = None,
  365. text_rec_score_thresh: Union[float, None] = None,
  366. seal_det_limit_side_len: Union[int, None] = None,
  367. seal_det_limit_type: Union[str, None] = None,
  368. seal_det_thresh: Union[float, None] = None,
  369. seal_det_box_thresh: Union[float, None] = None,
  370. seal_det_unclip_ratio: Union[float, None] = None,
  371. seal_rec_score_thresh: Union[float, None] = None,
  372. **kwargs,
  373. ) -> LayoutParsingResultV2:
  374. """
  375. Predicts the layout parsing result for the given input.
  376. Args:
  377. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  378. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  379. use_textline_orientation (Optional[bool]): Whether to use textline orientation prediction.
  380. use_general_ocr (Optional[bool]): Whether to use general OCR.
  381. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  382. use_table_recognition (Optional[bool]): Whether to use table recognition.
  383. use_formula_recognition (Optional[bool]): Whether to use formula recognition.
  384. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
  385. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False.
  386. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box.
  387. Defaults to None.
  388. If it's a single number, then both width and height are used.
  389. If it's a tuple of two numbers, then they are used separately for width and height respectively.
  390. If it's None, then no unclipping will be performed.
  391. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None.
  392. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  393. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  394. text_det_thresh (Optional[float]): Threshold for text detection.
  395. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  396. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  397. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  398. seal_det_limit_side_len (Optional[int]): Maximum side length for seal detection.
  399. seal_det_limit_type (Optional[str]): Type of limit to apply for seal detection.
  400. seal_det_thresh (Optional[float]): Threshold for seal detection.
  401. seal_det_box_thresh (Optional[float]): Threshold for seal detection boxes.
  402. seal_det_unclip_ratio (Optional[float]): Ratio for unclipping seal detection boxes.
  403. seal_rec_score_thresh (Optional[float]): Score threshold for seal recognition.
  404. **kwargs (Any): Additional settings to extend functionality.
  405. Returns:
  406. LayoutParsingResultV2: The predicted layout parsing result.
  407. """
  408. model_settings = self.get_model_settings(
  409. use_doc_orientation_classify,
  410. use_doc_unwarping,
  411. use_general_ocr,
  412. use_seal_recognition,
  413. use_table_recognition,
  414. use_formula_recognition,
  415. )
  416. if not self.check_model_settings_valid(model_settings):
  417. yield {"error": "the input params for model settings are invalid!"}
  418. for img_id, batch_data in enumerate(self.batch_sampler(input)):
  419. image_array = self.img_reader(batch_data.instances)[0]
  420. if model_settings["use_doc_preprocessor"]:
  421. doc_preprocessor_res = next(
  422. self.doc_preprocessor_pipeline(
  423. image_array,
  424. use_doc_orientation_classify=use_doc_orientation_classify,
  425. use_doc_unwarping=use_doc_unwarping,
  426. ),
  427. )
  428. else:
  429. doc_preprocessor_res = {"output_img": image_array}
  430. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  431. layout_det_res = next(
  432. self.layout_det_model(
  433. doc_preprocessor_image,
  434. threshold=layout_threshold,
  435. layout_nms=layout_nms,
  436. layout_unclip_ratio=layout_unclip_ratio,
  437. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  438. )
  439. )
  440. if model_settings["use_formula_recognition"]:
  441. formula_res_all = next(
  442. self.formula_recognition_pipeline(
  443. doc_preprocessor_image,
  444. use_layout_detection=False,
  445. use_doc_orientation_classify=False,
  446. use_doc_unwarping=False,
  447. layout_det_res=layout_det_res,
  448. ),
  449. )
  450. formula_res_list = formula_res_all["formula_res_list"]
  451. else:
  452. formula_res_list = []
  453. for formula_res in formula_res_list:
  454. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  455. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = 255.0
  456. if (
  457. model_settings["use_general_ocr"]
  458. or model_settings["use_table_recognition"]
  459. ):
  460. overall_ocr_res = next(
  461. self.general_ocr_pipeline(
  462. doc_preprocessor_image,
  463. use_textline_orientation=use_textline_orientation,
  464. text_det_limit_side_len=text_det_limit_side_len,
  465. text_det_limit_type=text_det_limit_type,
  466. text_det_thresh=text_det_thresh,
  467. text_det_box_thresh=text_det_box_thresh,
  468. text_det_unclip_ratio=text_det_unclip_ratio,
  469. text_rec_score_thresh=text_rec_score_thresh,
  470. ),
  471. )
  472. else:
  473. overall_ocr_res = {}
  474. if model_settings["use_table_recognition"]:
  475. table_overall_ocr_res = copy.deepcopy(overall_ocr_res)
  476. for formula_res in formula_res_list:
  477. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  478. poly_points = [
  479. (x_min, y_min),
  480. (x_max, y_min),
  481. (x_max, y_max),
  482. (x_min, y_max),
  483. ]
  484. table_overall_ocr_res["dt_polys"].append(poly_points)
  485. table_overall_ocr_res["rec_texts"].append(
  486. f"${formula_res['rec_formula']}$"
  487. )
  488. table_overall_ocr_res["rec_boxes"] = np.vstack(
  489. (table_overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
  490. )
  491. table_overall_ocr_res["rec_polys"].append(poly_points)
  492. table_overall_ocr_res["rec_scores"].append(1)
  493. table_res_all = next(
  494. self.table_recognition_pipeline(
  495. doc_preprocessor_image,
  496. use_doc_orientation_classify=False,
  497. use_doc_unwarping=False,
  498. use_layout_detection=False,
  499. use_ocr_model=False,
  500. overall_ocr_res=table_overall_ocr_res,
  501. layout_det_res=layout_det_res,
  502. cell_sort_by_y_projection=True,
  503. ),
  504. )
  505. table_res_list = table_res_all["table_res_list"]
  506. else:
  507. table_res_list = []
  508. if model_settings["use_seal_recognition"]:
  509. seal_res_all = next(
  510. self.seal_recognition_pipeline(
  511. doc_preprocessor_image,
  512. use_doc_orientation_classify=False,
  513. use_doc_unwarping=False,
  514. use_layout_detection=False,
  515. layout_det_res=layout_det_res,
  516. seal_det_limit_side_len=seal_det_limit_side_len,
  517. seal_det_limit_type=seal_det_limit_type,
  518. seal_det_thresh=seal_det_thresh,
  519. seal_det_box_thresh=seal_det_box_thresh,
  520. seal_det_unclip_ratio=seal_det_unclip_ratio,
  521. seal_rec_score_thresh=seal_rec_score_thresh,
  522. ),
  523. )
  524. seal_res_list = seal_res_all["seal_res_list"]
  525. else:
  526. seal_res_list = []
  527. parsing_res_list = self.get_layout_parsing_res(
  528. doc_preprocessor_image,
  529. layout_det_res=layout_det_res,
  530. overall_ocr_res=overall_ocr_res,
  531. table_res_list=table_res_list,
  532. seal_res_list=seal_res_list,
  533. formula_res_list=formula_res_list,
  534. text_det_limit_side_len=text_det_limit_side_len,
  535. text_det_limit_type=text_det_limit_type,
  536. text_det_thresh=text_det_thresh,
  537. text_det_box_thresh=text_det_box_thresh,
  538. text_det_unclip_ratio=text_det_unclip_ratio,
  539. text_rec_score_thresh=text_rec_score_thresh,
  540. )
  541. for formula_res in formula_res_list:
  542. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  543. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
  544. "input_img"
  545. ]
  546. single_img_res = {
  547. "input_path": batch_data.input_paths[0],
  548. "page_index": batch_data.page_indexes[0],
  549. "doc_preprocessor_res": doc_preprocessor_res,
  550. "layout_det_res": layout_det_res,
  551. "overall_ocr_res": overall_ocr_res,
  552. "table_res_list": table_res_list,
  553. "seal_res_list": seal_res_list,
  554. "formula_res_list": formula_res_list,
  555. "parsing_res_list": parsing_res_list,
  556. "model_settings": model_settings,
  557. }
  558. yield LayoutParsingResultV2(single_img_res)
  559. def concatenate_markdown_pages(self, markdown_list: list) -> tuple:
  560. """
  561. Concatenate Markdown content from multiple pages into a single document.
  562. Args:
  563. markdown_list (list): A list containing Markdown data for each page.
  564. Returns:
  565. tuple: A tuple containing the processed Markdown text.
  566. """
  567. markdown_texts = ""
  568. previous_page_last_element_paragraph_end_flag = True
  569. for res in markdown_list:
  570. # Get the paragraph flags for the current page
  571. page_first_element_paragraph_start_flag: bool = res[
  572. "page_continuation_flags"
  573. ][0]
  574. page_last_element_paragraph_end_flag: bool = res["page_continuation_flags"][
  575. 1
  576. ]
  577. # Determine whether to add a space or a newline
  578. if (
  579. not page_first_element_paragraph_start_flag
  580. and not previous_page_last_element_paragraph_end_flag
  581. ):
  582. last_char_of_markdown = markdown_texts[-1] if markdown_texts else ""
  583. first_char_of_handler = (
  584. res["markdown_texts"][0] if res["markdown_texts"] else ""
  585. )
  586. # Check if the last character and the first character are Chinese characters
  587. last_is_chinese_char = (
  588. re.match(r"[\u4e00-\u9fff]", last_char_of_markdown)
  589. if last_char_of_markdown
  590. else False
  591. )
  592. first_is_chinese_char = (
  593. re.match(r"[\u4e00-\u9fff]", first_char_of_handler)
  594. if first_char_of_handler
  595. else False
  596. )
  597. if not (last_is_chinese_char or first_is_chinese_char):
  598. markdown_texts += " " + res["markdown_texts"]
  599. else:
  600. markdown_texts += res["markdown_texts"]
  601. else:
  602. markdown_texts += "\n\n" + res["markdown_texts"]
  603. previous_page_last_element_paragraph_end_flag = (
  604. page_last_element_paragraph_end_flag
  605. )
  606. return markdown_texts