pipeline_v2.py 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710
  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"].size > 0:
  273. sub_ocr_res["rec_labels"] = ["text"] * len(sub_ocr_res["rec_texts"])
  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. overall_ocr_res["rec_labels"].extend(sub_ocr_res["rec_labels"])
  282. for formula_res in formula_res_list:
  283. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  284. poly_points = [
  285. (x_min, y_min),
  286. (x_max, y_min),
  287. (x_max, y_max),
  288. (x_min, y_max),
  289. ]
  290. overall_ocr_res["dt_polys"].append(poly_points)
  291. overall_ocr_res["rec_texts"].append(f"${formula_res['rec_formula']}$")
  292. overall_ocr_res["rec_boxes"] = np.vstack(
  293. (overall_ocr_res["rec_boxes"], [formula_res["dt_polys"]])
  294. )
  295. overall_ocr_res["rec_labels"].append("formula")
  296. overall_ocr_res["rec_polys"].append(poly_points)
  297. overall_ocr_res["rec_scores"].append(1)
  298. parsing_res_list = get_single_block_parsing_res(
  299. self.general_ocr_pipeline,
  300. overall_ocr_res=overall_ocr_res,
  301. layout_det_res=layout_det_res,
  302. table_res_list=table_res_list,
  303. seal_res_list=seal_res_list,
  304. )
  305. return parsing_res_list
  306. def get_model_settings(
  307. self,
  308. use_doc_orientation_classify: Union[bool, None],
  309. use_doc_unwarping: Union[bool, None],
  310. use_general_ocr: Union[bool, None],
  311. use_seal_recognition: Union[bool, None],
  312. use_table_recognition: Union[bool, None],
  313. use_formula_recognition: Union[bool, None],
  314. ) -> dict:
  315. """
  316. Get the model settings based on the provided parameters or default values.
  317. Args:
  318. use_doc_orientation_classify (Union[bool, None]): Enables document orientation classification if True. Defaults to system setting if None.
  319. use_doc_unwarping (Union[bool, None]): Enables document unwarping if True. Defaults to system setting if None.
  320. use_general_ocr (Union[bool, None]): Enables general OCR if True. Defaults to system setting if None.
  321. use_seal_recognition (Union[bool, None]): Enables seal recognition if True. Defaults to system setting if None.
  322. use_table_recognition (Union[bool, None]): Enables table recognition if True. Defaults to system setting if None.
  323. use_formula_recognition (Union[bool, None]): Enables formula recognition if True. Defaults to system setting if None.
  324. Returns:
  325. dict: A dictionary containing the model settings.
  326. """
  327. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  328. use_doc_preprocessor = self.use_doc_preprocessor
  329. else:
  330. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  331. use_doc_preprocessor = True
  332. else:
  333. use_doc_preprocessor = False
  334. if use_general_ocr is None:
  335. use_general_ocr = self.use_general_ocr
  336. if use_seal_recognition is None:
  337. use_seal_recognition = self.use_seal_recognition
  338. if use_table_recognition is None:
  339. use_table_recognition = self.use_table_recognition
  340. if use_formula_recognition is None:
  341. use_formula_recognition = self.use_formula_recognition
  342. return dict(
  343. use_doc_preprocessor=use_doc_preprocessor,
  344. use_general_ocr=use_general_ocr,
  345. use_seal_recognition=use_seal_recognition,
  346. use_table_recognition=use_table_recognition,
  347. use_formula_recognition=use_formula_recognition,
  348. )
  349. def predict(
  350. self,
  351. input: Union[str, list[str], np.ndarray, list[np.ndarray]],
  352. use_doc_orientation_classify: Union[bool, None] = None,
  353. use_doc_unwarping: Union[bool, None] = None,
  354. use_textline_orientation: Optional[bool] = None,
  355. use_general_ocr: Union[bool, None] = None,
  356. use_seal_recognition: Union[bool, None] = None,
  357. use_table_recognition: Union[bool, None] = None,
  358. use_formula_recognition: Union[bool, None] = None,
  359. layout_threshold: Optional[Union[float, dict]] = None,
  360. layout_nms: Optional[bool] = None,
  361. layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
  362. layout_merge_bboxes_mode: Optional[str] = None,
  363. text_det_limit_side_len: Union[int, None] = None,
  364. text_det_limit_type: Union[str, None] = None,
  365. text_det_thresh: Union[float, None] = None,
  366. text_det_box_thresh: Union[float, None] = None,
  367. text_det_unclip_ratio: Union[float, None] = None,
  368. text_rec_score_thresh: Union[float, None] = None,
  369. seal_det_limit_side_len: Union[int, None] = None,
  370. seal_det_limit_type: Union[str, None] = None,
  371. seal_det_thresh: Union[float, None] = None,
  372. seal_det_box_thresh: Union[float, None] = None,
  373. seal_det_unclip_ratio: Union[float, None] = None,
  374. seal_rec_score_thresh: Union[float, None] = None,
  375. **kwargs,
  376. ) -> LayoutParsingResultV2:
  377. """
  378. Predicts the layout parsing result for the given input.
  379. Args:
  380. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  381. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  382. use_textline_orientation (Optional[bool]): Whether to use textline orientation prediction.
  383. use_general_ocr (Optional[bool]): Whether to use general OCR.
  384. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  385. use_table_recognition (Optional[bool]): Whether to use table recognition.
  386. use_formula_recognition (Optional[bool]): Whether to use formula recognition.
  387. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
  388. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False.
  389. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box.
  390. Defaults to None.
  391. If it's a single number, then both width and height are used.
  392. If it's a tuple of two numbers, then they are used separately for width and height respectively.
  393. If it's None, then no unclipping will be performed.
  394. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None.
  395. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  396. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  397. text_det_thresh (Optional[float]): Threshold for text detection.
  398. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  399. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  400. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  401. seal_det_limit_side_len (Optional[int]): Maximum side length for seal detection.
  402. seal_det_limit_type (Optional[str]): Type of limit to apply for seal detection.
  403. seal_det_thresh (Optional[float]): Threshold for seal detection.
  404. seal_det_box_thresh (Optional[float]): Threshold for seal detection boxes.
  405. seal_det_unclip_ratio (Optional[float]): Ratio for unclipping seal detection boxes.
  406. seal_rec_score_thresh (Optional[float]): Score threshold for seal recognition.
  407. **kwargs (Any): Additional settings to extend functionality.
  408. Returns:
  409. LayoutParsingResultV2: The predicted layout parsing result.
  410. """
  411. model_settings = self.get_model_settings(
  412. use_doc_orientation_classify,
  413. use_doc_unwarping,
  414. use_general_ocr,
  415. use_seal_recognition,
  416. use_table_recognition,
  417. use_formula_recognition,
  418. )
  419. if not self.check_model_settings_valid(model_settings):
  420. yield {"error": "the input params for model settings are invalid!"}
  421. for batch_data in self.batch_sampler(input):
  422. image_array = self.img_reader(batch_data.instances)[0]
  423. if model_settings["use_doc_preprocessor"]:
  424. doc_preprocessor_res = next(
  425. self.doc_preprocessor_pipeline(
  426. image_array,
  427. use_doc_orientation_classify=use_doc_orientation_classify,
  428. use_doc_unwarping=use_doc_unwarping,
  429. ),
  430. )
  431. else:
  432. doc_preprocessor_res = {"output_img": image_array}
  433. doc_preprocessor_image = doc_preprocessor_res["output_img"]
  434. layout_det_res = next(
  435. self.layout_det_model(
  436. doc_preprocessor_image,
  437. threshold=layout_threshold,
  438. layout_nms=layout_nms,
  439. layout_unclip_ratio=layout_unclip_ratio,
  440. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  441. )
  442. )
  443. imgs_in_doc = gather_imgs(doc_preprocessor_image, layout_det_res["boxes"])
  444. if model_settings["use_formula_recognition"]:
  445. formula_res_all = next(
  446. self.formula_recognition_pipeline(
  447. doc_preprocessor_image,
  448. use_layout_detection=False,
  449. use_doc_orientation_classify=False,
  450. use_doc_unwarping=False,
  451. layout_det_res=layout_det_res,
  452. ),
  453. )
  454. formula_res_list = formula_res_all["formula_res_list"]
  455. else:
  456. formula_res_list = []
  457. for formula_res in formula_res_list:
  458. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  459. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = 255.0
  460. if (
  461. model_settings["use_general_ocr"]
  462. or model_settings["use_table_recognition"]
  463. ):
  464. overall_ocr_res = next(
  465. self.general_ocr_pipeline(
  466. doc_preprocessor_image,
  467. use_textline_orientation=use_textline_orientation,
  468. text_det_limit_side_len=text_det_limit_side_len,
  469. text_det_limit_type=text_det_limit_type,
  470. text_det_thresh=text_det_thresh,
  471. text_det_box_thresh=text_det_box_thresh,
  472. text_det_unclip_ratio=text_det_unclip_ratio,
  473. text_rec_score_thresh=text_rec_score_thresh,
  474. ),
  475. )
  476. else:
  477. overall_ocr_res = {}
  478. overall_ocr_res["rec_labels"] = ["text"] * len(overall_ocr_res["rec_texts"])
  479. if model_settings["use_table_recognition"]:
  480. table_contents = copy.deepcopy(overall_ocr_res)
  481. for formula_res in formula_res_list:
  482. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  483. poly_points = [
  484. (x_min, y_min),
  485. (x_max, y_min),
  486. (x_max, y_max),
  487. (x_min, y_max),
  488. ]
  489. table_contents["dt_polys"].append(poly_points)
  490. table_contents["rec_texts"].append(
  491. f"${formula_res['rec_formula']}$"
  492. )
  493. table_contents["rec_boxes"] = np.vstack(
  494. (table_contents["rec_boxes"], [formula_res["dt_polys"]])
  495. )
  496. table_contents["rec_polys"].append(poly_points)
  497. table_contents["rec_scores"].append(1)
  498. for img in imgs_in_doc:
  499. img_path = img["path"]
  500. x_min, y_min, x_max, y_max = img["coordinate"]
  501. poly_points = [
  502. (x_min, y_min),
  503. (x_max, y_min),
  504. (x_max, y_max),
  505. (x_min, y_max),
  506. ]
  507. table_contents["dt_polys"].append(poly_points)
  508. table_contents["rec_texts"].append(
  509. f'<div style="text-align: center;"><img src="{img_path}" alt="Image" /></div>'
  510. )
  511. table_contents["rec_boxes"] = np.vstack(
  512. (table_contents["rec_boxes"], img["coordinate"])
  513. )
  514. table_contents["rec_polys"].append(poly_points)
  515. table_contents["rec_scores"].append(img["score"])
  516. table_res_all = next(
  517. self.table_recognition_pipeline(
  518. doc_preprocessor_image,
  519. use_doc_orientation_classify=False,
  520. use_doc_unwarping=False,
  521. use_layout_detection=False,
  522. use_ocr_model=False,
  523. overall_ocr_res=table_contents,
  524. layout_det_res=layout_det_res,
  525. cell_sort_by_y_projection=True,
  526. ),
  527. )
  528. table_res_list = table_res_all["table_res_list"]
  529. else:
  530. table_res_list = []
  531. if model_settings["use_seal_recognition"]:
  532. seal_res_all = next(
  533. self.seal_recognition_pipeline(
  534. doc_preprocessor_image,
  535. use_doc_orientation_classify=False,
  536. use_doc_unwarping=False,
  537. use_layout_detection=False,
  538. layout_det_res=layout_det_res,
  539. seal_det_limit_side_len=seal_det_limit_side_len,
  540. seal_det_limit_type=seal_det_limit_type,
  541. seal_det_thresh=seal_det_thresh,
  542. seal_det_box_thresh=seal_det_box_thresh,
  543. seal_det_unclip_ratio=seal_det_unclip_ratio,
  544. seal_rec_score_thresh=seal_rec_score_thresh,
  545. ),
  546. )
  547. seal_res_list = seal_res_all["seal_res_list"]
  548. else:
  549. seal_res_list = []
  550. parsing_res_list = self.get_layout_parsing_res(
  551. doc_preprocessor_image,
  552. layout_det_res=layout_det_res,
  553. overall_ocr_res=overall_ocr_res,
  554. table_res_list=table_res_list,
  555. seal_res_list=seal_res_list,
  556. formula_res_list=formula_res_list,
  557. text_det_limit_side_len=text_det_limit_side_len,
  558. text_det_limit_type=text_det_limit_type,
  559. text_det_thresh=text_det_thresh,
  560. text_det_box_thresh=text_det_box_thresh,
  561. text_det_unclip_ratio=text_det_unclip_ratio,
  562. text_rec_score_thresh=text_rec_score_thresh,
  563. )
  564. for formula_res in formula_res_list:
  565. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  566. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
  567. "input_img"
  568. ]
  569. single_img_res = {
  570. "input_path": batch_data.input_paths[0],
  571. "page_index": batch_data.page_indexes[0],
  572. "doc_preprocessor_res": doc_preprocessor_res,
  573. "layout_det_res": layout_det_res,
  574. "overall_ocr_res": overall_ocr_res,
  575. "table_res_list": table_res_list,
  576. "seal_res_list": seal_res_list,
  577. "formula_res_list": formula_res_list,
  578. "parsing_res_list": parsing_res_list,
  579. "imgs_in_doc": imgs_in_doc,
  580. "model_settings": model_settings,
  581. }
  582. yield LayoutParsingResultV2(single_img_res)
  583. def concatenate_markdown_pages(self, markdown_list: list) -> tuple:
  584. """
  585. Concatenate Markdown content from multiple pages into a single document.
  586. Args:
  587. markdown_list (list): A list containing Markdown data for each page.
  588. Returns:
  589. tuple: A tuple containing the processed Markdown text.
  590. """
  591. markdown_texts = ""
  592. previous_page_last_element_paragraph_end_flag = True
  593. for res in markdown_list:
  594. # Get the paragraph flags for the current page
  595. page_first_element_paragraph_start_flag: bool = res[
  596. "page_continuation_flags"
  597. ][0]
  598. page_last_element_paragraph_end_flag: bool = res["page_continuation_flags"][
  599. 1
  600. ]
  601. # Determine whether to add a space or a newline
  602. if (
  603. not page_first_element_paragraph_start_flag
  604. and not previous_page_last_element_paragraph_end_flag
  605. ):
  606. last_char_of_markdown = markdown_texts[-1] if markdown_texts else ""
  607. first_char_of_handler = (
  608. res["markdown_texts"][0] if res["markdown_texts"] else ""
  609. )
  610. # Check if the last character and the first character are Chinese characters
  611. last_is_chinese_char = (
  612. re.match(r"[\u4e00-\u9fff]", last_char_of_markdown)
  613. if last_char_of_markdown
  614. else False
  615. )
  616. first_is_chinese_char = (
  617. re.match(r"[\u4e00-\u9fff]", first_char_of_handler)
  618. if first_char_of_handler
  619. else False
  620. )
  621. if not (last_is_chinese_char or first_is_chinese_char):
  622. markdown_texts += " " + res["markdown_texts"]
  623. else:
  624. markdown_texts += res["markdown_texts"]
  625. else:
  626. markdown_texts += "\n\n" + res["markdown_texts"]
  627. previous_page_last_element_paragraph_end_flag = (
  628. page_last_element_paragraph_end_flag
  629. )
  630. return markdown_texts