pipeline_v2.py 32 KB

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