pipeline_v2.py 33 KB

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