pipeline_v2.py 61 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345
  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, List, Optional, Tuple, Union
  18. import numpy as np
  19. from PIL import Image
  20. from ....utils import logging
  21. from ....utils.deps import pipeline_requires_extra
  22. from ...common.batch_sampler import ImageBatchSampler
  23. from ...common.reader import ReadImage
  24. from ...models.object_detection.result import DetResult
  25. from ...utils.benchmark import benchmark
  26. from ...utils.hpi import HPIConfig
  27. from ...utils.pp_option import PaddlePredictorOption
  28. from .._parallel import AutoParallelImageSimpleInferencePipeline
  29. from ..base import BasePipeline
  30. from ..ocr.result import OCRResult
  31. from .layout_objects import LayoutBlock, LayoutRegion
  32. from .result_v2 import LayoutParsingResultV2
  33. from .setting import BLOCK_LABEL_MAP, BLOCK_SETTINGS, REGION_SETTINGS
  34. from .utils import (
  35. calculate_bbox_area,
  36. calculate_minimum_enclosing_bbox,
  37. calculate_overlap_ratio,
  38. convert_formula_res_to_ocr_format,
  39. gather_imgs,
  40. get_bbox_intersection,
  41. get_sub_regions_ocr_res,
  42. remove_overlap_blocks,
  43. shrink_supplement_region_bbox,
  44. update_region_box,
  45. )
  46. from .xycut_enhanced import xycut_enhanced
  47. @benchmark.time_methods
  48. class _LayoutParsingPipelineV2(BasePipeline):
  49. """Layout Parsing Pipeline V2"""
  50. def __init__(
  51. self,
  52. config: dict,
  53. device: str = None,
  54. pp_option: PaddlePredictorOption = None,
  55. use_hpip: bool = False,
  56. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  57. ) -> None:
  58. """Initializes the layout parsing pipeline.
  59. Args:
  60. config (Dict): Configuration dictionary containing various settings.
  61. device (str, optional): Device to run the predictions on. Defaults to None.
  62. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  63. use_hpip (bool, optional): Whether to use the high-performance
  64. inference plugin (HPIP) by default. Defaults to False.
  65. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  66. The default high-performance inference configuration dictionary.
  67. Defaults to None.
  68. """
  69. super().__init__(
  70. device=device,
  71. pp_option=pp_option,
  72. use_hpip=use_hpip,
  73. hpi_config=hpi_config,
  74. )
  75. self.inintial_predictor(config)
  76. self.batch_sampler = ImageBatchSampler(batch_size=config.get("batch_size", 1))
  77. self.img_reader = ReadImage(format="BGR")
  78. def close(self):
  79. if getattr(self, "chart_recognition_model"):
  80. self.chart_recognition_model.close()
  81. def inintial_predictor(self, config: dict) -> None:
  82. """Initializes the predictor based on the provided configuration.
  83. Args:
  84. config (Dict): A dictionary containing the configuration for the predictor.
  85. Returns:
  86. None
  87. """
  88. if (
  89. config.get("use_doc_preprocessor", True)
  90. or config.get("use_doc_orientation_classify", True)
  91. or config.get("use_doc_unwarping", True)
  92. ):
  93. self.use_doc_preprocessor = True
  94. else:
  95. self.use_doc_preprocessor = False
  96. self.use_table_recognition = config.get("use_table_recognition", True)
  97. self.use_seal_recognition = config.get("use_seal_recognition", True)
  98. self.format_block_content = config.get("format_block_content", False)
  99. self.use_region_detection = config.get(
  100. "use_region_detection",
  101. True,
  102. )
  103. self.use_formula_recognition = config.get(
  104. "use_formula_recognition",
  105. True,
  106. )
  107. self.use_chart_recognition = config.get(
  108. "use_chart_recognition",
  109. False,
  110. )
  111. if self.use_doc_preprocessor:
  112. doc_preprocessor_config = config.get("SubPipelines", {}).get(
  113. "DocPreprocessor",
  114. {
  115. "pipeline_config_error": "config error for doc_preprocessor_pipeline!",
  116. },
  117. )
  118. self.doc_preprocessor_pipeline = self.create_pipeline(
  119. doc_preprocessor_config,
  120. )
  121. if self.use_region_detection:
  122. region_detection_config = config.get("SubModules", {}).get(
  123. "RegionDetection",
  124. {
  125. "model_config_error": "config error for block_region_detection_model!"
  126. },
  127. )
  128. self.region_detection_model = self.create_model(
  129. region_detection_config,
  130. )
  131. layout_det_config = config.get("SubModules", {}).get(
  132. "LayoutDetection",
  133. {"model_config_error": "config error for layout_det_model!"},
  134. )
  135. layout_kwargs = {}
  136. if (threshold := layout_det_config.get("threshold", None)) is not None:
  137. layout_kwargs["threshold"] = threshold
  138. if (layout_nms := layout_det_config.get("layout_nms", None)) is not None:
  139. layout_kwargs["layout_nms"] = layout_nms
  140. if (
  141. layout_unclip_ratio := layout_det_config.get("layout_unclip_ratio", None)
  142. ) is not None:
  143. layout_kwargs["layout_unclip_ratio"] = layout_unclip_ratio
  144. if (
  145. layout_merge_bboxes_mode := layout_det_config.get(
  146. "layout_merge_bboxes_mode", None
  147. )
  148. ) is not None:
  149. layout_kwargs["layout_merge_bboxes_mode"] = layout_merge_bboxes_mode
  150. self.layout_det_model = self.create_model(layout_det_config, **layout_kwargs)
  151. general_ocr_config = config.get("SubPipelines", {}).get(
  152. "GeneralOCR",
  153. {"pipeline_config_error": "config error for general_ocr_pipeline!"},
  154. )
  155. self.general_ocr_pipeline = self.create_pipeline(
  156. general_ocr_config,
  157. )
  158. if self.use_seal_recognition:
  159. seal_recognition_config = config.get("SubPipelines", {}).get(
  160. "SealRecognition",
  161. {
  162. "pipeline_config_error": "config error for seal_recognition_pipeline!",
  163. },
  164. )
  165. self.seal_recognition_pipeline = self.create_pipeline(
  166. seal_recognition_config,
  167. )
  168. if self.use_table_recognition:
  169. table_recognition_config = config.get("SubPipelines", {}).get(
  170. "TableRecognition",
  171. {
  172. "pipeline_config_error": "config error for table_recognition_pipeline!",
  173. },
  174. )
  175. self.table_recognition_pipeline = self.create_pipeline(
  176. table_recognition_config,
  177. )
  178. if self.use_formula_recognition:
  179. formula_recognition_config = config.get("SubPipelines", {}).get(
  180. "FormulaRecognition",
  181. {
  182. "pipeline_config_error": "config error for formula_recognition_pipeline!",
  183. },
  184. )
  185. self.formula_recognition_pipeline = self.create_pipeline(
  186. formula_recognition_config,
  187. )
  188. # TODO(gaotingquan): init the model at any time
  189. chart_recognition_config = config.get("SubModules", {}).get(
  190. "ChartRecognition",
  191. {"model_config_error": "config error for block_region_detection_model!"},
  192. )
  193. self.chart_recognition_model = self.create_model(
  194. chart_recognition_config,
  195. )
  196. return
  197. def get_text_paragraphs_ocr_res(
  198. self,
  199. overall_ocr_res: OCRResult,
  200. layout_det_res: DetResult,
  201. ) -> OCRResult:
  202. """
  203. Retrieves the OCR results for text paragraphs, excluding those of formulas, tables, and seals.
  204. Args:
  205. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  206. layout_det_res (DetResult): The detection result containing the layout information of the document.
  207. Returns:
  208. OCRResult: The OCR result for text paragraphs after excluding formulas, tables, and seals.
  209. """
  210. object_boxes = []
  211. for box_info in layout_det_res["boxes"]:
  212. if box_info["label"].lower() in ["formula", "table", "seal"]:
  213. object_boxes.append(box_info["coordinate"])
  214. object_boxes = np.array(object_boxes)
  215. sub_regions_ocr_res = get_sub_regions_ocr_res(
  216. overall_ocr_res, object_boxes, flag_within=False
  217. )
  218. return sub_regions_ocr_res
  219. def check_model_settings_valid(self, input_params: dict) -> bool:
  220. """
  221. Check if the input parameters are valid based on the initialized models.
  222. Args:
  223. input_params (Dict): A dictionary containing input parameters.
  224. Returns:
  225. bool: True if all required models are initialized according to input parameters, False otherwise.
  226. """
  227. if input_params["use_doc_preprocessor"] and not self.use_doc_preprocessor:
  228. logging.error(
  229. "Set use_doc_preprocessor, but the models for doc preprocessor are not initialized.",
  230. )
  231. return False
  232. if input_params["use_seal_recognition"] and not self.use_seal_recognition:
  233. logging.error(
  234. "Set use_seal_recognition, but the models for seal recognition are not initialized.",
  235. )
  236. return False
  237. if input_params["use_table_recognition"] and not self.use_table_recognition:
  238. logging.error(
  239. "Set use_table_recognition, but the models for table recognition are not initialized.",
  240. )
  241. return False
  242. return True
  243. def standardized_data(
  244. self,
  245. image: list,
  246. region_det_res: DetResult,
  247. layout_det_res: DetResult,
  248. overall_ocr_res: OCRResult,
  249. formula_res_list: list,
  250. text_rec_model: Any,
  251. text_rec_score_thresh: Union[float, None] = None,
  252. ) -> list:
  253. """
  254. Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
  255. Args:
  256. image (list): The input image.
  257. overall_ocr_res (OCRResult): An object containing the overall OCR results, including detected text boxes and recognized text. The structure is expected to have:
  258. - "input_img": The image on which OCR was performed.
  259. - "dt_boxes": A list of detected text box coordinates.
  260. - "rec_texts": A list of recognized text corresponding to the detected boxes.
  261. layout_det_res (DetResult): An object containing the layout detection results, including detected layout boxes and their labels. The structure is expected to have:
  262. - "boxes": A list of dictionaries with keys "coordinate" for box coordinates and "block_label" for the type of content.
  263. table_res_list (list): A list of table detection results, where each item is a dictionary containing:
  264. - "block_bbox": The bounding box of the table layout.
  265. - "pred_html": The predicted HTML representation of the table.
  266. formula_res_list (list): A list of formula recognition results.
  267. text_rec_model (Any): The text recognition model.
  268. text_rec_score_thresh (Optional[float], optional): The score threshold for text recognition. Defaults to None.
  269. Returns:
  270. list: A list of dictionaries representing the layout parsing result.
  271. """
  272. matched_ocr_dict = {}
  273. region_to_block_map = {}
  274. block_to_ocr_map = {}
  275. object_boxes = []
  276. footnote_list = []
  277. paragraph_title_list = []
  278. bottom_text_y_max = 0
  279. max_block_area = 0.0
  280. doc_title_num = 0
  281. base_region_bbox = [65535, 65535, 0, 0]
  282. layout_det_res = remove_overlap_blocks(
  283. layout_det_res,
  284. threshold=0.5,
  285. smaller=True,
  286. )
  287. # convert formula_res_list to OCRResult format
  288. convert_formula_res_to_ocr_format(formula_res_list, overall_ocr_res)
  289. # match layout boxes and ocr boxes and get some information for layout_order_config
  290. for box_idx, box_info in enumerate(layout_det_res["boxes"]):
  291. box = box_info["coordinate"]
  292. label = box_info["label"].lower()
  293. object_boxes.append(box)
  294. _, _, _, y2 = box
  295. # update the region box and max_block_area according to the layout boxes
  296. base_region_bbox = update_region_box(box, base_region_bbox)
  297. max_block_area = max(max_block_area, calculate_bbox_area(box))
  298. # update_layout_order_config_block_index(layout_order_config, label, box_idx)
  299. # set the label of footnote to text, when it is above the text boxes
  300. if label == "footnote":
  301. footnote_list.append(box_idx)
  302. elif label == "paragraph_title":
  303. paragraph_title_list.append(box_idx)
  304. if label == "text":
  305. bottom_text_y_max = max(y2, bottom_text_y_max)
  306. if label == "doc_title":
  307. doc_title_num += 1
  308. if label not in ["formula", "table", "seal"]:
  309. _, matched_idxes = get_sub_regions_ocr_res(
  310. overall_ocr_res, [box], return_match_idx=True
  311. )
  312. block_to_ocr_map[box_idx] = matched_idxes
  313. for matched_idx in matched_idxes:
  314. if matched_ocr_dict.get(matched_idx, None) is None:
  315. matched_ocr_dict[matched_idx] = [box_idx]
  316. else:
  317. matched_ocr_dict[matched_idx].append(box_idx)
  318. # fix the footnote label
  319. for footnote_idx in footnote_list:
  320. if (
  321. layout_det_res["boxes"][footnote_idx]["coordinate"][3]
  322. < bottom_text_y_max
  323. ):
  324. layout_det_res["boxes"][footnote_idx]["label"] = "text"
  325. # check if there is only one paragraph title and without doc_title
  326. only_one_paragraph_title = len(paragraph_title_list) == 1 and doc_title_num == 0
  327. if only_one_paragraph_title:
  328. paragraph_title_block_area = calculate_bbox_area(
  329. layout_det_res["boxes"][paragraph_title_list[0]]["coordinate"]
  330. )
  331. title_area_max_block_threshold = BLOCK_SETTINGS.get(
  332. "title_conversion_area_ratio_threshold", 0.3
  333. )
  334. if (
  335. paragraph_title_block_area
  336. > max_block_area * title_area_max_block_threshold
  337. ):
  338. layout_det_res["boxes"][paragraph_title_list[0]]["label"] = "doc_title"
  339. # Replace the OCR information of the hurdles.
  340. for overall_ocr_idx, layout_box_ids in matched_ocr_dict.items():
  341. if len(layout_box_ids) > 1:
  342. matched_no = 0
  343. overall_ocr_box = copy.deepcopy(
  344. overall_ocr_res["rec_boxes"][overall_ocr_idx]
  345. )
  346. overall_ocr_dt_poly = copy.deepcopy(
  347. overall_ocr_res["dt_polys"][overall_ocr_idx]
  348. )
  349. for box_idx in layout_box_ids:
  350. layout_box = layout_det_res["boxes"][box_idx]["coordinate"]
  351. crop_box = get_bbox_intersection(overall_ocr_box, layout_box)
  352. for ocr_idx in block_to_ocr_map[box_idx]:
  353. ocr_box = overall_ocr_res["rec_boxes"][ocr_idx]
  354. iou = calculate_overlap_ratio(ocr_box, crop_box, "small")
  355. if iou > 0.8:
  356. overall_ocr_res["rec_texts"][ocr_idx] = ""
  357. x1, y1, x2, y2 = [int(i) for i in crop_box]
  358. crop_img = np.array(image)[y1:y2, x1:x2]
  359. crop_img_rec_res = list(text_rec_model([crop_img]))[0]
  360. crop_img_dt_poly = get_bbox_intersection(
  361. overall_ocr_dt_poly, layout_box, return_format="poly"
  362. )
  363. crop_img_rec_score = crop_img_rec_res["rec_score"]
  364. crop_img_rec_text = crop_img_rec_res["rec_text"]
  365. text_rec_score_thresh = (
  366. text_rec_score_thresh
  367. if text_rec_score_thresh is not None
  368. else (self.general_ocr_pipeline.text_rec_score_thresh)
  369. )
  370. if crop_img_rec_score >= text_rec_score_thresh:
  371. matched_no += 1
  372. if matched_no == 1:
  373. # the first matched ocr be replaced by the first matched layout box
  374. overall_ocr_res["dt_polys"][
  375. overall_ocr_idx
  376. ] = crop_img_dt_poly
  377. overall_ocr_res["rec_boxes"][overall_ocr_idx] = crop_box
  378. overall_ocr_res["rec_polys"][
  379. overall_ocr_idx
  380. ] = crop_img_dt_poly
  381. overall_ocr_res["rec_scores"][
  382. overall_ocr_idx
  383. ] = crop_img_rec_score
  384. overall_ocr_res["rec_texts"][
  385. overall_ocr_idx
  386. ] = crop_img_rec_text
  387. else:
  388. # the other matched ocr be appended to the overall ocr result
  389. overall_ocr_res["dt_polys"].append(crop_img_dt_poly)
  390. if len(overall_ocr_res["rec_boxes"]) == 0:
  391. overall_ocr_res["rec_boxes"] = np.array([crop_box])
  392. else:
  393. overall_ocr_res["rec_boxes"] = np.vstack(
  394. (overall_ocr_res["rec_boxes"], crop_box)
  395. )
  396. overall_ocr_res["rec_polys"].append(crop_img_dt_poly)
  397. overall_ocr_res["rec_scores"].append(crop_img_rec_score)
  398. overall_ocr_res["rec_texts"].append(crop_img_rec_text)
  399. overall_ocr_res["rec_labels"].append("text")
  400. block_to_ocr_map[box_idx].remove(overall_ocr_idx)
  401. block_to_ocr_map[box_idx].append(
  402. len(overall_ocr_res["rec_texts"]) - 1
  403. )
  404. # use layout bbox to do ocr recognition when there is no matched ocr
  405. for layout_box_idx, overall_ocr_idxes in block_to_ocr_map.items():
  406. has_text = False
  407. for idx in overall_ocr_idxes:
  408. if overall_ocr_res["rec_texts"][idx] != "":
  409. has_text = True
  410. break
  411. if not has_text and layout_det_res["boxes"][layout_box_idx][
  412. "label"
  413. ] not in BLOCK_LABEL_MAP.get("vision_labels", []):
  414. crop_box = layout_det_res["boxes"][layout_box_idx]["coordinate"]
  415. x1, y1, x2, y2 = [int(i) for i in crop_box]
  416. crop_img = np.array(image)[y1:y2, x1:x2]
  417. crop_img_rec_res = list(text_rec_model([crop_img]))[0]
  418. crop_img_dt_poly = get_bbox_intersection(
  419. crop_box, crop_box, return_format="poly"
  420. )
  421. crop_img_rec_score = crop_img_rec_res["rec_score"]
  422. crop_img_rec_text = crop_img_rec_res["rec_text"]
  423. text_rec_score_thresh = (
  424. text_rec_score_thresh
  425. if text_rec_score_thresh is not None
  426. else (self.general_ocr_pipeline.text_rec_score_thresh)
  427. )
  428. if crop_img_rec_score >= text_rec_score_thresh:
  429. if len(overall_ocr_res["rec_boxes"]) == 0:
  430. overall_ocr_res["rec_boxes"] = np.array([crop_box])
  431. else:
  432. overall_ocr_res["rec_boxes"] = np.vstack(
  433. (overall_ocr_res["rec_boxes"], crop_box)
  434. )
  435. overall_ocr_res["rec_polys"].append(crop_img_dt_poly)
  436. overall_ocr_res["rec_scores"].append(crop_img_rec_score)
  437. overall_ocr_res["rec_texts"].append(crop_img_rec_text)
  438. overall_ocr_res["rec_labels"].append("text")
  439. block_to_ocr_map[layout_box_idx].append(
  440. len(overall_ocr_res["rec_texts"]) - 1
  441. )
  442. # when there is no layout detection result but there is ocr result, convert ocr detection result to layout detection result
  443. if len(layout_det_res["boxes"]) == 0 and len(overall_ocr_res["rec_boxes"]) > 0:
  444. for idx, ocr_rec_box in enumerate(overall_ocr_res["rec_boxes"]):
  445. base_region_bbox = update_region_box(ocr_rec_box, base_region_bbox)
  446. layout_det_res["boxes"].append(
  447. {
  448. "label": "text",
  449. "coordinate": ocr_rec_box,
  450. "score": overall_ocr_res["rec_scores"][idx],
  451. }
  452. )
  453. block_to_ocr_map[idx] = [idx]
  454. mask_labels = (
  455. BLOCK_LABEL_MAP.get("unordered_labels", [])
  456. + BLOCK_LABEL_MAP.get("header_labels", [])
  457. + BLOCK_LABEL_MAP.get("footer_labels", [])
  458. )
  459. block_bboxes = [box["coordinate"] for box in layout_det_res["boxes"]]
  460. region_det_res["boxes"] = sorted(
  461. region_det_res["boxes"],
  462. key=lambda item: calculate_bbox_area(item["coordinate"]),
  463. )
  464. if len(region_det_res["boxes"]) == 0:
  465. region_det_res["boxes"] = [
  466. {
  467. "coordinate": base_region_bbox,
  468. "label": "SupplementaryRegion",
  469. "score": 1,
  470. }
  471. ]
  472. region_to_block_map[0] = range(len(block_bboxes))
  473. else:
  474. block_idxes_set = set(range(len(block_bboxes)))
  475. # match block to region
  476. for region_idx, region_info in enumerate(region_det_res["boxes"]):
  477. matched_idxes = []
  478. region_to_block_map[region_idx] = []
  479. region_bbox = region_info["coordinate"]
  480. for block_idx in block_idxes_set:
  481. if layout_det_res["boxes"][block_idx]["label"] in mask_labels:
  482. continue
  483. overlap_ratio = calculate_overlap_ratio(
  484. region_bbox, block_bboxes[block_idx], mode="small"
  485. )
  486. if overlap_ratio > REGION_SETTINGS.get(
  487. "match_block_overlap_ratio_threshold", 0.8
  488. ):
  489. matched_idxes.append(block_idx)
  490. old_region_bbox_matched_idxes = []
  491. if len(matched_idxes) > 0:
  492. while len(old_region_bbox_matched_idxes) != len(matched_idxes):
  493. old_region_bbox_matched_idxes = copy.deepcopy(matched_idxes)
  494. matched_idxes = []
  495. matched_bboxes = [
  496. block_bboxes[idx] for idx in old_region_bbox_matched_idxes
  497. ]
  498. new_region_bbox = calculate_minimum_enclosing_bbox(
  499. matched_bboxes
  500. )
  501. for block_idx in block_idxes_set:
  502. if (
  503. layout_det_res["boxes"][block_idx]["label"]
  504. in mask_labels
  505. ):
  506. continue
  507. overlap_ratio = calculate_overlap_ratio(
  508. new_region_bbox, block_bboxes[block_idx], mode="small"
  509. )
  510. if overlap_ratio > REGION_SETTINGS.get(
  511. "match_block_overlap_ratio_threshold", 0.8
  512. ):
  513. matched_idxes.append(block_idx)
  514. for block_idx in matched_idxes:
  515. block_idxes_set.remove(block_idx)
  516. region_to_block_map[region_idx] = matched_idxes
  517. region_det_res["boxes"][region_idx]["coordinate"] = new_region_bbox
  518. # Supplement region when there is no matched block
  519. while len(block_idxes_set) > 0:
  520. unmatched_bboxes = [block_bboxes[idx] for idx in block_idxes_set]
  521. if len(unmatched_bboxes) == 0:
  522. break
  523. supplement_region_bbox = calculate_minimum_enclosing_bbox(
  524. unmatched_bboxes
  525. )
  526. matched_idxes = []
  527. # check if the new region bbox is overlapped with other region bbox, if have, then shrink the new region bbox
  528. for region_idx, region_info in enumerate(region_det_res["boxes"]):
  529. if len(region_to_block_map[region_idx]) == 0:
  530. continue
  531. region_bbox = region_info["coordinate"]
  532. overlap_ratio = calculate_overlap_ratio(
  533. supplement_region_bbox, region_bbox
  534. )
  535. if overlap_ratio > 0:
  536. supplement_region_bbox, matched_idxes = (
  537. shrink_supplement_region_bbox(
  538. supplement_region_bbox,
  539. region_bbox,
  540. image.shape[1],
  541. image.shape[0],
  542. block_idxes_set,
  543. block_bboxes,
  544. )
  545. )
  546. matched_idxes = [
  547. idx
  548. for idx in matched_idxes
  549. if layout_det_res["boxes"][idx]["label"] not in mask_labels
  550. ]
  551. if len(matched_idxes) == 0:
  552. matched_idxes = [
  553. idx
  554. for idx in block_idxes_set
  555. if layout_det_res["boxes"][idx]["label"] not in mask_labels
  556. ]
  557. if len(matched_idxes) == 0:
  558. break
  559. matched_bboxes = [block_bboxes[idx] for idx in matched_idxes]
  560. supplement_region_bbox = calculate_minimum_enclosing_bbox(
  561. matched_bboxes
  562. )
  563. region_idx = len(region_det_res["boxes"])
  564. region_to_block_map[region_idx] = list(matched_idxes)
  565. for block_idx in matched_idxes:
  566. block_idxes_set.remove(block_idx)
  567. region_det_res["boxes"].append(
  568. {
  569. "coordinate": supplement_region_bbox,
  570. "label": "SupplementaryRegion",
  571. "score": 1,
  572. }
  573. )
  574. mask_idxes = [
  575. idx
  576. for idx in range(len(layout_det_res["boxes"]))
  577. if layout_det_res["boxes"][idx]["label"] in mask_labels
  578. ]
  579. for idx in mask_idxes:
  580. bbox = layout_det_res["boxes"][idx]["coordinate"]
  581. region_idx = len(region_det_res["boxes"])
  582. region_to_block_map[region_idx] = [idx]
  583. region_det_res["boxes"].append(
  584. {
  585. "coordinate": bbox,
  586. "label": "SupplementaryRegion",
  587. "score": 1,
  588. }
  589. )
  590. region_block_ocr_idx_map = dict(
  591. region_to_block_map=region_to_block_map,
  592. block_to_ocr_map=block_to_ocr_map,
  593. )
  594. return region_block_ocr_idx_map, region_det_res, layout_det_res
  595. def get_layout_parsing_objects(
  596. self,
  597. image: list,
  598. region_block_ocr_idx_map: dict,
  599. region_det_res: DetResult,
  600. overall_ocr_res: OCRResult,
  601. layout_det_res: DetResult,
  602. table_res_list: list,
  603. seal_res_list: list,
  604. chart_res_list: list,
  605. text_rec_model: Any,
  606. text_rec_score_thresh: Union[float, None] = None,
  607. ) -> list:
  608. """
  609. Extract structured information from OCR and layout detection results.
  610. Args:
  611. image (list): The input image.
  612. overall_ocr_res (OCRResult): An object containing the overall OCR results, including detected text boxes and recognized text. The structure is expected to have:
  613. - "input_img": The image on which OCR was performed.
  614. - "dt_boxes": A list of detected text box coordinates.
  615. - "rec_texts": A list of recognized text corresponding to the detected boxes.
  616. layout_det_res (DetResult): An object containing the layout detection results, including detected layout boxes and their labels. The structure is expected to have:
  617. - "boxes": A list of dictionaries with keys "coordinate" for box coordinates and "block_label" for the type of content.
  618. table_res_list (list): A list of table detection results, where each item is a dictionary containing:
  619. - "block_bbox": The bounding box of the table layout.
  620. - "pred_html": The predicted HTML representation of the table.
  621. seal_res_list (List): A list of seal detection results. The details of each item depend on the specific application context.
  622. text_rec_model (Any): A model for text recognition.
  623. text_rec_score_thresh (Union[float, None]): The minimum score required for a recognized character to be considered valid. If None, use the default value specified during initialization. Default is None.
  624. Returns:
  625. list: A list of structured boxes where each item is a dictionary containing:
  626. - "block_label": The label of the content (e.g., 'table', 'chart', 'image').
  627. - The label as a key with either table HTML or image data and text.
  628. - "block_bbox": The coordinates of the layout box.
  629. """
  630. table_index = 0
  631. seal_index = 0
  632. chart_index = 0
  633. layout_parsing_blocks: List[LayoutBlock] = []
  634. for box_idx, box_info in enumerate(layout_det_res["boxes"]):
  635. label = box_info["label"]
  636. block_bbox = box_info["coordinate"]
  637. rec_res = {"boxes": [], "rec_texts": [], "rec_labels": []}
  638. block = LayoutBlock(label=label, bbox=block_bbox)
  639. if label == "table" and len(table_res_list) > 0:
  640. block.content = table_res_list[table_index]["pred_html"]
  641. table_index += 1
  642. elif label == "seal" and len(seal_res_list) > 0:
  643. block.content = "\n".join(seal_res_list[seal_index]["rec_texts"])
  644. seal_index += 1
  645. elif label == "chart" and len(chart_res_list) > 0:
  646. block.content = chart_res_list[chart_index]
  647. chart_index += 1
  648. else:
  649. if label == "formula":
  650. _, ocr_idx_list = get_sub_regions_ocr_res(
  651. overall_ocr_res, [block_bbox], return_match_idx=True
  652. )
  653. region_block_ocr_idx_map["block_to_ocr_map"][box_idx] = ocr_idx_list
  654. else:
  655. ocr_idx_list = region_block_ocr_idx_map["block_to_ocr_map"].get(
  656. box_idx, []
  657. )
  658. for box_no in ocr_idx_list:
  659. rec_res["boxes"].append(overall_ocr_res["rec_boxes"][box_no])
  660. rec_res["rec_texts"].append(
  661. overall_ocr_res["rec_texts"][box_no],
  662. )
  663. rec_res["rec_labels"].append(
  664. overall_ocr_res["rec_labels"][box_no],
  665. )
  666. block.update_text_content(
  667. image=image,
  668. ocr_rec_res=rec_res,
  669. text_rec_model=text_rec_model,
  670. text_rec_score_thresh=text_rec_score_thresh,
  671. )
  672. if (
  673. label
  674. in ["seal", "table", "formula", "chart"]
  675. + BLOCK_LABEL_MAP["image_labels"]
  676. ):
  677. x_min, y_min, x_max, y_max = list(map(int, block_bbox))
  678. img_path = (
  679. f"imgs/img_in_{block.label}_box_{x_min}_{y_min}_{x_max}_{y_max}.jpg"
  680. )
  681. img = Image.fromarray(image[y_min:y_max, x_min:x_max, ::-1])
  682. block.image = {"path": img_path, "img": img}
  683. layout_parsing_blocks.append(block)
  684. page_region_bbox = [65535, 65535, 0, 0]
  685. layout_parsing_regions: List[LayoutRegion] = []
  686. for region_idx, region_info in enumerate(region_det_res["boxes"]):
  687. region_bbox = np.array(region_info["coordinate"]).astype("int")
  688. region_blocks = [
  689. layout_parsing_blocks[idx]
  690. for idx in region_block_ocr_idx_map["region_to_block_map"][region_idx]
  691. ]
  692. if region_blocks:
  693. page_region_bbox = update_region_box(region_bbox, page_region_bbox)
  694. region = LayoutRegion(bbox=region_bbox, blocks=region_blocks)
  695. layout_parsing_regions.append(region)
  696. layout_parsing_page = LayoutRegion(
  697. bbox=np.array(page_region_bbox).astype("int"), blocks=layout_parsing_regions
  698. )
  699. return layout_parsing_page
  700. def sort_layout_parsing_blocks(
  701. self, layout_parsing_page: LayoutRegion
  702. ) -> List[LayoutBlock]:
  703. layout_parsing_regions = xycut_enhanced(layout_parsing_page)
  704. parsing_res_list = []
  705. for region in layout_parsing_regions:
  706. layout_parsing_blocks = xycut_enhanced(region)
  707. parsing_res_list.extend(layout_parsing_blocks)
  708. return parsing_res_list
  709. def get_layout_parsing_res(
  710. self,
  711. image: list,
  712. region_det_res: DetResult,
  713. layout_det_res: DetResult,
  714. overall_ocr_res: OCRResult,
  715. table_res_list: list,
  716. seal_res_list: list,
  717. chart_res_list: list,
  718. formula_res_list: list,
  719. text_rec_score_thresh: Union[float, None] = None,
  720. ) -> list:
  721. """
  722. Retrieves the layout parsing result based on the layout detection result, OCR result, and other recognition results.
  723. Args:
  724. image (list): The input image.
  725. layout_det_res (DetResult): The detection result containing the layout information of the document.
  726. overall_ocr_res (OCRResult): The overall OCR result containing text information.
  727. table_res_list (list): A list of table recognition results.
  728. seal_res_list (list): A list of seal recognition results.
  729. formula_res_list (list): A list of formula recognition results.
  730. text_rec_score_thresh (Optional[float], optional): The score threshold for text recognition. Defaults to None.
  731. Returns:
  732. list: A list of dictionaries representing the layout parsing result.
  733. """
  734. # Standardize data
  735. region_block_ocr_idx_map, region_det_res, layout_det_res = (
  736. self.standardized_data(
  737. image=image,
  738. region_det_res=region_det_res,
  739. layout_det_res=layout_det_res,
  740. overall_ocr_res=overall_ocr_res,
  741. formula_res_list=formula_res_list,
  742. text_rec_model=self.general_ocr_pipeline.text_rec_model,
  743. text_rec_score_thresh=text_rec_score_thresh,
  744. )
  745. )
  746. # Format layout parsing block
  747. layout_parsing_page = self.get_layout_parsing_objects(
  748. image=image,
  749. region_block_ocr_idx_map=region_block_ocr_idx_map,
  750. region_det_res=region_det_res,
  751. overall_ocr_res=overall_ocr_res,
  752. layout_det_res=layout_det_res,
  753. table_res_list=table_res_list,
  754. seal_res_list=seal_res_list,
  755. chart_res_list=chart_res_list,
  756. text_rec_model=self.general_ocr_pipeline.text_rec_model,
  757. text_rec_score_thresh=self.general_ocr_pipeline.text_rec_score_thresh,
  758. )
  759. parsing_res_list = self.sort_layout_parsing_blocks(layout_parsing_page)
  760. order_index = 1
  761. for index, block in enumerate(parsing_res_list):
  762. block.index = index
  763. if block.label in BLOCK_LABEL_MAP["visualize_index_labels"]:
  764. block.order_index = order_index
  765. order_index += 1
  766. return parsing_res_list
  767. def get_model_settings(
  768. self,
  769. use_doc_orientation_classify: Union[bool, None],
  770. use_doc_unwarping: Union[bool, None],
  771. use_seal_recognition: Union[bool, None],
  772. use_table_recognition: Union[bool, None],
  773. use_formula_recognition: Union[bool, None],
  774. use_chart_recognition: Union[bool, None],
  775. use_region_detection: Union[bool, None],
  776. format_block_content: Union[bool, None],
  777. ) -> dict:
  778. """
  779. Get the model settings based on the provided parameters or default values.
  780. Args:
  781. use_doc_orientation_classify (Union[bool, None]): Enables document orientation classification if True. Defaults to system setting if None.
  782. use_doc_unwarping (Union[bool, None]): Enables document unwarping if True. Defaults to system setting if None.
  783. use_seal_recognition (Union[bool, None]): Enables seal recognition if True. Defaults to system setting if None.
  784. use_table_recognition (Union[bool, None]): Enables table recognition if True. Defaults to system setting if None.
  785. use_formula_recognition (Union[bool, None]): Enables formula recognition if True. Defaults to system setting if None.
  786. format_block_content (Union[bool, None]): Enables block content formatting if True. Defaults to system setting if None.
  787. Returns:
  788. dict: A dictionary containing the model settings.
  789. """
  790. if use_doc_orientation_classify is None and use_doc_unwarping is None:
  791. use_doc_preprocessor = self.use_doc_preprocessor
  792. else:
  793. if use_doc_orientation_classify is True or use_doc_unwarping is True:
  794. use_doc_preprocessor = True
  795. else:
  796. use_doc_preprocessor = False
  797. if use_seal_recognition is None:
  798. use_seal_recognition = self.use_seal_recognition
  799. if use_table_recognition is None:
  800. use_table_recognition = self.use_table_recognition
  801. if use_formula_recognition is None:
  802. use_formula_recognition = self.use_formula_recognition
  803. if use_region_detection is None:
  804. use_region_detection = self.use_region_detection
  805. if use_chart_recognition is None:
  806. use_chart_recognition = self.use_chart_recognition
  807. if format_block_content is None:
  808. format_block_content = self.format_block_content
  809. return dict(
  810. use_doc_preprocessor=use_doc_preprocessor,
  811. use_seal_recognition=use_seal_recognition,
  812. use_table_recognition=use_table_recognition,
  813. use_formula_recognition=use_formula_recognition,
  814. use_chart_recognition=use_chart_recognition,
  815. use_region_detection=use_region_detection,
  816. format_block_content=format_block_content,
  817. )
  818. def predict(
  819. self,
  820. input: Union[str, list[str], np.ndarray, list[np.ndarray]],
  821. use_doc_orientation_classify: Union[bool, None] = None,
  822. use_doc_unwarping: Union[bool, None] = None,
  823. use_textline_orientation: Optional[bool] = None,
  824. use_seal_recognition: Union[bool, None] = None,
  825. use_table_recognition: Union[bool, None] = None,
  826. use_formula_recognition: Union[bool, None] = None,
  827. use_chart_recognition: Union[bool, None] = None,
  828. use_region_detection: Union[bool, None] = None,
  829. format_block_content: Union[bool, None] = None,
  830. layout_threshold: Optional[Union[float, dict]] = None,
  831. layout_nms: Optional[bool] = None,
  832. layout_unclip_ratio: Optional[Union[float, Tuple[float, float], dict]] = None,
  833. layout_merge_bboxes_mode: Optional[str] = None,
  834. text_det_limit_side_len: Union[int, None] = None,
  835. text_det_limit_type: Union[str, None] = None,
  836. text_det_thresh: Union[float, None] = None,
  837. text_det_box_thresh: Union[float, None] = None,
  838. text_det_unclip_ratio: Union[float, None] = None,
  839. text_rec_score_thresh: Union[float, None] = None,
  840. seal_det_limit_side_len: Union[int, None] = None,
  841. seal_det_limit_type: Union[str, None] = None,
  842. seal_det_thresh: Union[float, None] = None,
  843. seal_det_box_thresh: Union[float, None] = None,
  844. seal_det_unclip_ratio: Union[float, None] = None,
  845. seal_rec_score_thresh: Union[float, None] = None,
  846. use_wired_table_cells_trans_to_html: bool = False,
  847. use_wireless_table_cells_trans_to_html: bool = False,
  848. use_table_orientation_classify: bool = True,
  849. use_ocr_results_with_table_cells: bool = True,
  850. use_e2e_wired_table_rec_model: bool = False,
  851. use_e2e_wireless_table_rec_model: bool = True,
  852. **kwargs,
  853. ) -> LayoutParsingResultV2:
  854. """
  855. Predicts the layout parsing result for the given input.
  856. Args:
  857. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): Input image path, list of image paths,
  858. numpy array of an image, or list of numpy arrays.
  859. use_doc_orientation_classify (Optional[bool]): Whether to use document orientation classification.
  860. use_doc_unwarping (Optional[bool]): Whether to use document unwarping.
  861. use_textline_orientation (Optional[bool]): Whether to use textline orientation prediction.
  862. use_seal_recognition (Optional[bool]): Whether to use seal recognition.
  863. use_table_recognition (Optional[bool]): Whether to use table recognition.
  864. use_formula_recognition (Optional[bool]): Whether to use formula recognition.
  865. use_region_detection (Optional[bool]): Whether to use region detection.
  866. format_block_content (Optional[bool]): Whether to format block content.
  867. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
  868. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False.
  869. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box.
  870. Defaults to None.
  871. If it's a single number, then both width and height are used.
  872. If it's a tuple of two numbers, then they are used separately for width and height respectively.
  873. If it's None, then no unclipping will be performed.
  874. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None.
  875. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  876. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  877. text_det_thresh (Optional[float]): Threshold for text detection.
  878. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  879. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  880. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  881. seal_det_limit_side_len (Optional[int]): Maximum side length for seal detection.
  882. seal_det_limit_type (Optional[str]): Type of limit to apply for seal detection.
  883. seal_det_thresh (Optional[float]): Threshold for seal detection.
  884. seal_det_box_thresh (Optional[float]): Threshold for seal detection boxes.
  885. seal_det_unclip_ratio (Optional[float]): Ratio for unclipping seal detection boxes.
  886. seal_rec_score_thresh (Optional[float]): Score threshold for seal recognition.
  887. use_wired_table_cells_trans_to_html (bool): Whether to use wired table cells trans to HTML.
  888. use_wireless_table_cells_trans_to_html (bool): Whether to use wireless table cells trans to HTML.
  889. use_table_orientation_classify (bool): Whether to use table orientation classification.
  890. use_ocr_results_with_table_cells (bool): Whether to use OCR results processed by table cells.
  891. use_e2e_wired_table_rec_model (bool): Whether to use end-to-end wired table recognition model.
  892. use_e2e_wireless_table_rec_model (bool): Whether to use end-to-end wireless table recognition model.
  893. **kwargs (Any): Additional settings to extend functionality.
  894. Returns:
  895. LayoutParsingResultV2: The predicted layout parsing result.
  896. """
  897. model_settings = self.get_model_settings(
  898. use_doc_orientation_classify,
  899. use_doc_unwarping,
  900. use_seal_recognition,
  901. use_table_recognition,
  902. use_formula_recognition,
  903. use_chart_recognition,
  904. use_region_detection,
  905. format_block_content,
  906. )
  907. if not self.check_model_settings_valid(model_settings):
  908. yield {"error": "the input params for model settings are invalid!"}
  909. for batch_data in self.batch_sampler(input):
  910. image_arrays = self.img_reader(batch_data.instances)
  911. if model_settings["use_doc_preprocessor"]:
  912. doc_preprocessor_results = list(
  913. self.doc_preprocessor_pipeline(
  914. image_arrays,
  915. use_doc_orientation_classify=use_doc_orientation_classify,
  916. use_doc_unwarping=use_doc_unwarping,
  917. )
  918. )
  919. else:
  920. doc_preprocessor_results = [{"output_img": arr} for arr in image_arrays]
  921. doc_preprocessor_images = [
  922. item["output_img"] for item in doc_preprocessor_results
  923. ]
  924. layout_det_results = list(
  925. self.layout_det_model(
  926. doc_preprocessor_images,
  927. threshold=layout_threshold,
  928. layout_nms=layout_nms,
  929. layout_unclip_ratio=layout_unclip_ratio,
  930. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  931. )
  932. )
  933. imgs_in_doc = [
  934. gather_imgs(img, res["boxes"])
  935. for img, res in zip(doc_preprocessor_images, layout_det_results)
  936. ]
  937. if model_settings["use_region_detection"]:
  938. region_det_results = list(
  939. self.region_detection_model(
  940. doc_preprocessor_images,
  941. layout_nms=True,
  942. layout_merge_bboxes_mode="small",
  943. ),
  944. )
  945. else:
  946. region_det_results = [{"boxes": []} for _ in doc_preprocessor_images]
  947. if model_settings["use_formula_recognition"]:
  948. formula_res_all = list(
  949. self.formula_recognition_pipeline(
  950. doc_preprocessor_images,
  951. use_layout_detection=False,
  952. use_doc_orientation_classify=False,
  953. use_doc_unwarping=False,
  954. layout_det_res=layout_det_results,
  955. ),
  956. )
  957. formula_res_lists = [
  958. item["formula_res_list"] for item in formula_res_all
  959. ]
  960. else:
  961. formula_res_lists = [[] for _ in doc_preprocessor_images]
  962. for doc_preprocessor_image, formula_res_list in zip(
  963. doc_preprocessor_images, formula_res_lists
  964. ):
  965. for formula_res in formula_res_list:
  966. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  967. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = 255.0
  968. overall_ocr_results = list(
  969. self.general_ocr_pipeline(
  970. doc_preprocessor_images,
  971. use_textline_orientation=use_textline_orientation,
  972. text_det_limit_side_len=text_det_limit_side_len,
  973. text_det_limit_type=text_det_limit_type,
  974. text_det_thresh=text_det_thresh,
  975. text_det_box_thresh=text_det_box_thresh,
  976. text_det_unclip_ratio=text_det_unclip_ratio,
  977. text_rec_score_thresh=text_rec_score_thresh,
  978. ),
  979. )
  980. for overall_ocr_res in overall_ocr_results:
  981. overall_ocr_res["rec_labels"] = ["text"] * len(
  982. overall_ocr_res["rec_texts"]
  983. )
  984. if model_settings["use_table_recognition"]:
  985. table_res_lists = []
  986. for (
  987. layout_det_res,
  988. doc_preprocessor_image,
  989. overall_ocr_res,
  990. formula_res_list,
  991. imgs_in_doc_for_img,
  992. ) in zip(
  993. layout_det_results,
  994. doc_preprocessor_images,
  995. overall_ocr_results,
  996. formula_res_lists,
  997. imgs_in_doc,
  998. ):
  999. table_contents_for_img = copy.deepcopy(overall_ocr_res)
  1000. for formula_res in formula_res_list:
  1001. x_min, y_min, x_max, y_max = list(
  1002. map(int, formula_res["dt_polys"])
  1003. )
  1004. poly_points = [
  1005. (x_min, y_min),
  1006. (x_max, y_min),
  1007. (x_max, y_max),
  1008. (x_min, y_max),
  1009. ]
  1010. table_contents_for_img["dt_polys"].append(poly_points)
  1011. rec_formula = formula_res["rec_formula"]
  1012. if not rec_formula.startswith("$") or not rec_formula.endswith(
  1013. "$"
  1014. ):
  1015. rec_formula = f"${rec_formula}$"
  1016. table_contents_for_img["rec_texts"].append(f"{rec_formula}")
  1017. if table_contents_for_img["rec_boxes"].size == 0:
  1018. table_contents_for_img["rec_boxes"] = np.array(
  1019. [formula_res["dt_polys"]]
  1020. )
  1021. else:
  1022. table_contents_for_img["rec_boxes"] = np.vstack(
  1023. (
  1024. table_contents_for_img["rec_boxes"],
  1025. [formula_res["dt_polys"]],
  1026. )
  1027. )
  1028. table_contents_for_img["rec_polys"].append(poly_points)
  1029. table_contents_for_img["rec_scores"].append(1)
  1030. for img in imgs_in_doc_for_img:
  1031. img_path = img["path"]
  1032. x_min, y_min, x_max, y_max = img["coordinate"]
  1033. poly_points = [
  1034. (x_min, y_min),
  1035. (x_max, y_min),
  1036. (x_max, y_max),
  1037. (x_min, y_max),
  1038. ]
  1039. table_contents_for_img["dt_polys"].append(poly_points)
  1040. table_contents_for_img["rec_texts"].append(
  1041. f'<div style="text-align: center;"><img src="{img_path}" alt="Image" /></div>'
  1042. )
  1043. if table_contents_for_img["rec_boxes"].size == 0:
  1044. table_contents_for_img["rec_boxes"] = np.array(
  1045. [img["coordinate"]]
  1046. )
  1047. else:
  1048. table_contents_for_img["rec_boxes"] = np.vstack(
  1049. (table_contents_for_img["rec_boxes"], img["coordinate"])
  1050. )
  1051. table_contents_for_img["rec_polys"].append(poly_points)
  1052. table_contents_for_img["rec_scores"].append(img["score"])
  1053. table_res_all = list(
  1054. self.table_recognition_pipeline(
  1055. doc_preprocessor_image,
  1056. use_doc_orientation_classify=False,
  1057. use_doc_unwarping=False,
  1058. use_layout_detection=False,
  1059. use_ocr_model=False,
  1060. overall_ocr_res=table_contents_for_img,
  1061. layout_det_res=layout_det_res,
  1062. cell_sort_by_y_projection=True,
  1063. use_wired_table_cells_trans_to_html=use_wired_table_cells_trans_to_html,
  1064. use_wireless_table_cells_trans_to_html=use_wireless_table_cells_trans_to_html,
  1065. use_table_orientation_classify=use_table_orientation_classify,
  1066. use_ocr_results_with_table_cells=use_ocr_results_with_table_cells,
  1067. use_e2e_wired_table_rec_model=use_e2e_wired_table_rec_model,
  1068. use_e2e_wireless_table_rec_model=use_e2e_wireless_table_rec_model,
  1069. ),
  1070. )
  1071. single_table_res_lists = [
  1072. item["table_res_list"] for item in table_res_all
  1073. ]
  1074. table_res_lists.extend(single_table_res_lists)
  1075. else:
  1076. table_res_lists = [[] for _ in doc_preprocessor_images]
  1077. if model_settings["use_seal_recognition"]:
  1078. seal_res_all = list(
  1079. self.seal_recognition_pipeline(
  1080. doc_preprocessor_images,
  1081. use_doc_orientation_classify=False,
  1082. use_doc_unwarping=False,
  1083. use_layout_detection=False,
  1084. layout_det_res=layout_det_results,
  1085. seal_det_limit_side_len=seal_det_limit_side_len,
  1086. seal_det_limit_type=seal_det_limit_type,
  1087. seal_det_thresh=seal_det_thresh,
  1088. seal_det_box_thresh=seal_det_box_thresh,
  1089. seal_det_unclip_ratio=seal_det_unclip_ratio,
  1090. seal_rec_score_thresh=seal_rec_score_thresh,
  1091. ),
  1092. )
  1093. seal_res_lists = [item["seal_res_list"] for item in seal_res_all]
  1094. else:
  1095. seal_res_lists = [[] for _ in doc_preprocessor_images]
  1096. for (
  1097. input_path,
  1098. page_index,
  1099. doc_preprocessor_image,
  1100. doc_preprocessor_res,
  1101. layout_det_res,
  1102. region_det_res,
  1103. overall_ocr_res,
  1104. table_res_list,
  1105. seal_res_list,
  1106. formula_res_list,
  1107. imgs_in_doc_for_img,
  1108. ) in zip(
  1109. batch_data.input_paths,
  1110. batch_data.page_indexes,
  1111. doc_preprocessor_images,
  1112. doc_preprocessor_results,
  1113. layout_det_results,
  1114. region_det_results,
  1115. overall_ocr_results,
  1116. table_res_lists,
  1117. seal_res_lists,
  1118. formula_res_lists,
  1119. imgs_in_doc,
  1120. ):
  1121. chart_res_list = []
  1122. if model_settings["use_chart_recognition"]:
  1123. chart_imgs_list = []
  1124. for bbox in layout_det_res["boxes"]:
  1125. if bbox["label"] == "chart":
  1126. x_min, y_min, x_max, y_max = bbox["coordinate"]
  1127. chart_img = doc_preprocessor_image[
  1128. int(y_min) : int(y_max), int(x_min) : int(x_max), :
  1129. ]
  1130. chart_imgs_list.append({"image": chart_img})
  1131. for chart_res_batch in self.chart_recognition_model(
  1132. input=chart_imgs_list
  1133. ):
  1134. chart_res_list.append(chart_res_batch["result"])
  1135. parsing_res_list = self.get_layout_parsing_res(
  1136. doc_preprocessor_image,
  1137. region_det_res=region_det_res,
  1138. layout_det_res=layout_det_res,
  1139. overall_ocr_res=overall_ocr_res,
  1140. table_res_list=table_res_list,
  1141. seal_res_list=seal_res_list,
  1142. chart_res_list=chart_res_list,
  1143. formula_res_list=formula_res_list,
  1144. text_rec_score_thresh=text_rec_score_thresh,
  1145. )
  1146. for formula_res in formula_res_list:
  1147. x_min, y_min, x_max, y_max = list(map(int, formula_res["dt_polys"]))
  1148. doc_preprocessor_image[y_min:y_max, x_min:x_max, :] = formula_res[
  1149. "input_img"
  1150. ]
  1151. single_img_res = {
  1152. "input_path": input_path,
  1153. "page_index": page_index,
  1154. "doc_preprocessor_res": doc_preprocessor_res,
  1155. "layout_det_res": layout_det_res,
  1156. "region_det_res": region_det_res,
  1157. "overall_ocr_res": overall_ocr_res,
  1158. "table_res_list": table_res_list,
  1159. "seal_res_list": seal_res_list,
  1160. "chart_res_list": chart_res_list,
  1161. "formula_res_list": formula_res_list,
  1162. "parsing_res_list": parsing_res_list,
  1163. "imgs_in_doc": imgs_in_doc_for_img,
  1164. "model_settings": model_settings,
  1165. }
  1166. yield LayoutParsingResultV2(single_img_res)
  1167. def concatenate_markdown_pages(self, markdown_list: list) -> tuple:
  1168. """
  1169. Concatenate Markdown content from multiple pages into a single document.
  1170. Args:
  1171. markdown_list (list): A list containing Markdown data for each page.
  1172. Returns:
  1173. tuple: A tuple containing the processed Markdown text.
  1174. """
  1175. markdown_texts = ""
  1176. previous_page_last_element_paragraph_end_flag = True
  1177. for res in markdown_list:
  1178. # Get the paragraph flags for the current page
  1179. page_first_element_paragraph_start_flag: bool = res[
  1180. "page_continuation_flags"
  1181. ][0]
  1182. page_last_element_paragraph_end_flag: bool = res["page_continuation_flags"][
  1183. 1
  1184. ]
  1185. # Determine whether to add a space or a newline
  1186. if (
  1187. not page_first_element_paragraph_start_flag
  1188. and not previous_page_last_element_paragraph_end_flag
  1189. ):
  1190. last_char_of_markdown = markdown_texts[-1] if markdown_texts else ""
  1191. first_char_of_handler = (
  1192. res["markdown_texts"][0] if res["markdown_texts"] else ""
  1193. )
  1194. # Check if the last character and the first character are Chinese characters
  1195. last_is_chinese_char = (
  1196. re.match(r"[\u4e00-\u9fff]", last_char_of_markdown)
  1197. if last_char_of_markdown
  1198. else False
  1199. )
  1200. first_is_chinese_char = (
  1201. re.match(r"[\u4e00-\u9fff]", first_char_of_handler)
  1202. if first_char_of_handler
  1203. else False
  1204. )
  1205. if not (last_is_chinese_char or first_is_chinese_char):
  1206. markdown_texts += " " + res["markdown_texts"]
  1207. else:
  1208. markdown_texts += res["markdown_texts"]
  1209. else:
  1210. markdown_texts += "\n\n" + res["markdown_texts"]
  1211. previous_page_last_element_paragraph_end_flag = (
  1212. page_last_element_paragraph_end_flag
  1213. )
  1214. return markdown_texts
  1215. @pipeline_requires_extra("ocr")
  1216. class LayoutParsingPipelineV2(AutoParallelImageSimpleInferencePipeline):
  1217. entities = ["PP-StructureV3"]
  1218. @property
  1219. def _pipeline_cls(self):
  1220. return _LayoutParsingPipelineV2
  1221. def _get_batch_size(self, config):
  1222. return config.get("batch_size", 1)