para_pipeline.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298
  1. import os
  2. import sys
  3. import json
  4. from para.commons import *
  5. from para.raw_processor import RawBlockProcessor
  6. from para.layout_match_processor import LayoutFilterProcessor
  7. from para.stats import BlockStatisticsCalculator
  8. from para.stats import DocStatisticsCalculator
  9. from para.title_processor import TitleProcessor
  10. from para.block_termination_processor import BlockTerminationProcessor
  11. from para.block_continuation_processor import BlockContinuationProcessor
  12. from para.draw import DrawAnnos
  13. from para.exceptions import (
  14. DenseSingleLineBlockException,
  15. TitleDetectionException,
  16. TitleLevelException,
  17. ParaSplitException,
  18. ParaMergeException,
  19. DiscardByException,
  20. )
  21. if sys.version_info[0] >= 3:
  22. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  23. class ParaProcessPipeline:
  24. def __init__(self) -> None:
  25. pass
  26. def para_process_pipeline(self, pdf_info_dict, para_debug_mode=None, input_pdf_path=None, output_pdf_path=None):
  27. """
  28. This function processes the paragraphs, including:
  29. 1. Read raw input json file into pdf_dic
  30. 2. Detect and replace equations
  31. 3. Combine spans into a natural line
  32. 4. Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
  33. 5. Compute statistics for each block
  34. 6. Detect titles in the document
  35. 7. Detect paragraphs inside each block
  36. 8. Divide the level of the titles
  37. 9. Detect and combine paragraphs from different blocks into one paragraph
  38. 10. Check whether the final results after checking headings, dividing paragraphs within blocks, and merging paragraphs between blocks are plausible and reasonable.
  39. 11. Draw annotations on the pdf file
  40. Parameters
  41. ----------
  42. pdf_dic_json_fpath : str
  43. path to the pdf dictionary json file.
  44. Notice: data noises, including overlap blocks, header, footer, watermark, vertical margin note have been removed already.
  45. input_pdf_doc : str
  46. path to the input pdf file
  47. output_pdf_path : str
  48. path to the output pdf file
  49. Returns
  50. -------
  51. pdf_dict : dict
  52. result dictionary
  53. """
  54. error_info = None
  55. output_json_file = ""
  56. output_dir = ""
  57. if input_pdf_path is not None:
  58. input_pdf_path = os.path.abspath(input_pdf_path)
  59. # print_green_on_red(f">>>>>>>>>>>>>>>>>>> Process the paragraphs of {input_pdf_path}")
  60. if output_pdf_path is not None:
  61. output_dir = os.path.dirname(output_pdf_path)
  62. output_json_file = f"{output_dir}/pdf_dic.json"
  63. def __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode):
  64. """
  65. Save the pdf_dic to a json file
  66. """
  67. output_pdf_file_name = os.path.basename(output_pdf_path)
  68. # output_dir = os.path.dirname(output_pdf_path)
  69. output_dir = "\\tmp\\pdf_parse"
  70. output_pdf_file_name = output_pdf_file_name.replace(".pdf", f"_stage_{stage}.json")
  71. pdf_dic_json_fpath = os.path.join(output_dir, output_pdf_file_name)
  72. if not os.path.exists(output_dir):
  73. os.makedirs(output_dir)
  74. if para_debug_mode == "full":
  75. with open(pdf_dic_json_fpath, "w", encoding="utf-8") as f:
  76. json.dump(pdf_dic, f, indent=2, ensure_ascii=False)
  77. # Validate the output already exists
  78. if not os.path.exists(pdf_dic_json_fpath):
  79. print_red(f"Failed to save the pdf_dic to {pdf_dic_json_fpath}")
  80. return None
  81. else:
  82. print_green(f"Succeed to save the pdf_dic to {pdf_dic_json_fpath}")
  83. return pdf_dic_json_fpath
  84. """
  85. Preprocess the lines of block
  86. """
  87. # Find and replace the interline and inline equations, should be better done before the paragraph processing
  88. # Create "para_blocks" for each page.
  89. # equationProcessor = EquationsProcessor()
  90. # pdf_dic = equationProcessor.batch_process_blocks(pdf_info_dict)
  91. # Combine spans into a natural line
  92. rawBlockProcessor = RawBlockProcessor()
  93. pdf_dic = rawBlockProcessor.batch_process_blocks(pdf_info_dict)
  94. # print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
  95. # Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
  96. layoutFilter = LayoutFilterProcessor()
  97. pdf_dic = layoutFilter.batch_process_blocks(pdf_dic)
  98. # Compute statistics for each block
  99. blockStatisticsCalculator = BlockStatisticsCalculator()
  100. pdf_dic = blockStatisticsCalculator.batch_process_blocks(pdf_dic)
  101. # print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
  102. # Compute statistics for all blocks(namely this pdf document)
  103. docStatisticsCalculator = DocStatisticsCalculator()
  104. pdf_dic = docStatisticsCalculator.calc_stats_of_doc(pdf_dic)
  105. # print(f"pdf_dic['statistics']: {pdf_dic['statistics']}", end="\n\n")
  106. # Dump the first three stages of pdf_dic to a json file
  107. if para_debug_mode == "full":
  108. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode)
  109. """
  110. Detect titles in the document
  111. """
  112. doc_statistics = pdf_dic["statistics"]
  113. titleProcessor = TitleProcessor(doc_statistics)
  114. pdf_dic = titleProcessor.batch_process_blocks_detect_titles(pdf_dic)
  115. if para_debug_mode == "full":
  116. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="1", para_debug_mode=para_debug_mode)
  117. """
  118. Detect and divide the level of the titles
  119. """
  120. titleProcessor = TitleProcessor()
  121. pdf_dic = titleProcessor.batch_process_blocks_recog_title_level(pdf_dic)
  122. if para_debug_mode == "full":
  123. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="2", para_debug_mode=para_debug_mode)
  124. """
  125. Detect and split paragraphs inside each block
  126. """
  127. blockInnerParasProcessor = BlockTerminationProcessor()
  128. pdf_dic = blockInnerParasProcessor.batch_process_blocks(pdf_dic)
  129. if para_debug_mode == "full":
  130. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode=para_debug_mode)
  131. # pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode="full")
  132. # print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
  133. """
  134. Detect and combine paragraphs from different blocks into one paragraph
  135. """
  136. blockContinuationProcessor = BlockContinuationProcessor()
  137. pdf_dic = blockContinuationProcessor.batch_tag_paras(pdf_dic)
  138. pdf_dic = blockContinuationProcessor.batch_merge_paras(pdf_dic)
  139. if para_debug_mode == "full":
  140. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode=para_debug_mode)
  141. # pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode="full")
  142. # print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
  143. """
  144. Discard pdf files by checking exceptions and return the error info to the caller
  145. """
  146. discardByException = DiscardByException()
  147. is_discard_by_single_line_block = discardByException.discard_by_single_line_block(
  148. pdf_dic, exception=DenseSingleLineBlockException()
  149. )
  150. is_discard_by_title_detection = discardByException.discard_by_title_detection(
  151. pdf_dic, exception=TitleDetectionException()
  152. )
  153. is_discard_by_title_level = discardByException.discard_by_title_level(pdf_dic, exception=TitleLevelException())
  154. is_discard_by_split_para = discardByException.discard_by_split_para(pdf_dic, exception=ParaSplitException())
  155. is_discard_by_merge_para = discardByException.discard_by_merge_para(pdf_dic, exception=ParaMergeException())
  156. """
  157. if any(
  158. info is not None
  159. for info in [
  160. is_discard_by_single_line_block,
  161. is_discard_by_title_detection,
  162. is_discard_by_title_level,
  163. is_discard_by_split_para,
  164. is_discard_by_merge_para,
  165. ]
  166. ):
  167. error_info = next(
  168. (
  169. info
  170. for info in [
  171. is_discard_by_single_line_block,
  172. is_discard_by_title_detection,
  173. is_discard_by_title_level,
  174. is_discard_by_split_para,
  175. is_discard_by_merge_para,
  176. ]
  177. if info is not None
  178. ),
  179. None,
  180. )
  181. return pdf_dic, error_info
  182. if any(
  183. info is not None
  184. for info in [
  185. is_discard_by_single_line_block,
  186. is_discard_by_title_detection,
  187. is_discard_by_title_level,
  188. is_discard_by_split_para,
  189. is_discard_by_merge_para,
  190. ]
  191. ):
  192. error_info = next(
  193. (
  194. info
  195. for info in [
  196. is_discard_by_single_line_block,
  197. is_discard_by_title_detection,
  198. is_discard_by_title_level,
  199. is_discard_by_split_para,
  200. is_discard_by_merge_para,
  201. ]
  202. if info is not None
  203. ),
  204. None,
  205. )
  206. return pdf_dic, error_info
  207. """
  208. """
  209. Dump the final pdf_dic to a json file
  210. """
  211. if para_debug_mode is not None:
  212. with open(output_json_file, "w", encoding="utf-8") as f:
  213. json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
  214. """
  215. Draw the annotations
  216. """
  217. if is_discard_by_single_line_block is not None:
  218. error_info = is_discard_by_single_line_block
  219. elif is_discard_by_title_detection is not None:
  220. error_info = is_discard_by_title_detection
  221. elif is_discard_by_title_level is not None:
  222. error_info = is_discard_by_title_level
  223. elif is_discard_by_split_para is not None:
  224. error_info = is_discard_by_split_para
  225. elif is_discard_by_merge_para is not None:
  226. error_info = is_discard_by_merge_para
  227. if error_info is not None:
  228. return pdf_dic, error_info
  229. """
  230. Dump the final pdf_dic to a json file
  231. """
  232. if para_debug_mode is not None:
  233. with open(output_json_file, "w", encoding="utf-8") as f:
  234. json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
  235. """
  236. Draw the annotations
  237. """
  238. if para_debug_mode is not None:
  239. drawAnnos = DrawAnnos()
  240. drawAnnos.draw_annos(input_pdf_path, pdf_dic, output_pdf_path)
  241. """
  242. Remove the intermediate files which are generated in the process of paragraph processing if debug_mode is simple
  243. """
  244. if para_debug_mode is not None:
  245. for fpath in os.listdir(output_dir):
  246. if fpath.endswith(".json") and "stage" in fpath:
  247. os.remove(os.path.join(output_dir, fpath))
  248. return pdf_dic, error_info