pipeline_v2.py 33 KB

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