pipeline_v2.py 31 KB

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