pipeline_ocr.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  1. import sys
  2. import time
  3. from loguru import logger
  4. from magic_pdf.dict2md.ocr_mkcontent import ocr_mk_mm_markdown, ocr_mk_nlp_markdown_with_para, \
  5. ocr_mk_mm_markdown_with_para_and_pagination, ocr_mk_mm_markdown_with_para, ocr_mk_mm_standard_format, \
  6. make_standard_format_with_para
  7. from magic_pdf.libs.commons import join_path, s3_image_save_path, formatted_time
  8. from magic_pdf.libs.json_compressor import JsonCompressor
  9. from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
  10. from magic_pdf.spark.base import get_data_source, exception_handler, get_pdf_bytes, get_bookname
  11. from magic_pdf.spark.s3 import get_s3_config
  12. def ocr_pdf_intermediate_dict_to_markdown(jso: dict, debug_mode=False) -> dict:
  13. if debug_mode:
  14. pass
  15. else: # 如果debug没开,则检测是否有needdrop字段
  16. if jso.get("need_drop", False):
  17. book_name = join_path(get_data_source(jso), jso["file_id"])
  18. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  19. jso["dropped"] = True
  20. return jso
  21. try:
  22. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  23. # 将 pdf_intermediate_dict 解压
  24. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  25. markdown_content = ocr_mk_mm_markdown(pdf_intermediate_dict)
  26. jso["content"] = markdown_content
  27. logger.info(
  28. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  29. file=sys.stderr,
  30. )
  31. # 把无用的信息清空
  32. jso["doc_layout_result"] = ""
  33. jso["pdf_intermediate_dict"] = ""
  34. jso["pdf_meta"] = ""
  35. except Exception as e:
  36. jso = exception_handler(jso, e)
  37. return jso
  38. def ocr_pdf_intermediate_dict_to_markdown_with_para(jso: dict, debug_mode=False) -> dict:
  39. if debug_mode:
  40. pass
  41. else: # 如果debug没开,则检测是否有needdrop字段
  42. if jso.get("need_drop", False):
  43. book_name = join_path(get_data_source(jso), jso["file_id"])
  44. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  45. jso["dropped"] = True
  46. return jso
  47. try:
  48. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  49. # 将 pdf_intermediate_dict 解压
  50. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  51. # markdown_content = ocr_mk_mm_markdown_with_para(pdf_intermediate_dict)
  52. markdown_content = ocr_mk_nlp_markdown_with_para(pdf_intermediate_dict)
  53. jso["content"] = markdown_content
  54. logger.info(
  55. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  56. file=sys.stderr,
  57. )
  58. # 把无用的信息清空
  59. jso["doc_layout_result"] = ""
  60. jso["pdf_intermediate_dict"] = ""
  61. jso["pdf_meta"] = ""
  62. except Exception as e:
  63. jso = exception_handler(jso, e)
  64. return jso
  65. def ocr_pdf_intermediate_dict_to_markdown_with_para_and_pagination(jso: dict, debug_mode=False) -> dict:
  66. if debug_mode:
  67. pass
  68. else: # 如果debug没开,则检测是否有needdrop字段
  69. if jso.get("need_drop", False):
  70. book_name = join_path(get_data_source(jso), jso["file_id"])
  71. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  72. jso["dropped"] = True
  73. return jso
  74. try:
  75. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  76. # 将 pdf_intermediate_dict 解压
  77. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  78. markdown_content = ocr_mk_mm_markdown_with_para_and_pagination(pdf_intermediate_dict)
  79. jso["content"] = markdown_content
  80. logger.info(
  81. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  82. file=sys.stderr,
  83. )
  84. # 把无用的信息清空
  85. # jso["doc_layout_result"] = ""
  86. jso["pdf_intermediate_dict"] = ""
  87. # jso["pdf_meta"] = ""
  88. except Exception as e:
  89. jso = exception_handler(jso, e)
  90. return jso
  91. def ocr_pdf_intermediate_dict_to_markdown_with_para_for_qa(
  92. jso: dict, debug_mode=False
  93. ) -> dict:
  94. if debug_mode:
  95. pass
  96. else: # 如果debug没开,则检测是否有needdrop字段
  97. if jso.get("need_drop", False):
  98. book_name = join_path(get_data_source(jso), jso["file_id"])
  99. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  100. jso["dropped"] = True
  101. return jso
  102. try:
  103. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  104. # 将 pdf_intermediate_dict 解压
  105. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  106. markdown_content = ocr_mk_mm_markdown_with_para(pdf_intermediate_dict)
  107. jso["content_ocr"] = markdown_content
  108. logger.info(
  109. f"book_name is:{get_data_source(jso)}/{jso['file_id']},markdown content length is {len(markdown_content)}",
  110. file=sys.stderr,
  111. )
  112. # 把无用的信息清空
  113. jso["doc_layout_result"] = ""
  114. jso["pdf_intermediate_dict"] = ""
  115. jso["mid_json_ocr"] = pdf_intermediate_dict
  116. jso["pdf_meta"] = ""
  117. except Exception as e:
  118. jso = exception_handler(jso, e)
  119. return jso
  120. def ocr_pdf_intermediate_dict_to_standard_format(jso: dict, debug_mode=False) -> dict:
  121. if debug_mode:
  122. pass
  123. else: # 如果debug没开,则检测是否有needdrop字段
  124. if jso.get("need_drop", False):
  125. book_name = join_path(get_data_source(jso), jso["file_id"])
  126. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  127. jso["dropped"] = True
  128. return jso
  129. try:
  130. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  131. # 将 pdf_intermediate_dict 解压
  132. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  133. standard_format = ocr_mk_mm_standard_format(pdf_intermediate_dict)
  134. jso["content_list"] = standard_format
  135. logger.info(
  136. f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",
  137. file=sys.stderr,
  138. )
  139. # 把无用的信息清空
  140. jso["doc_layout_result"] = ""
  141. jso["pdf_intermediate_dict"] = ""
  142. jso["pdf_meta"] = ""
  143. except Exception as e:
  144. jso = exception_handler(jso, e)
  145. return jso
  146. def ocr_pdf_intermediate_dict_to_standard_format_with_para(jso: dict, debug_mode=False) -> dict:
  147. if debug_mode:
  148. pass
  149. else: # 如果debug没开,则检测是否有needdrop字段
  150. if jso.get("need_drop", False):
  151. book_name = join_path(get_data_source(jso), jso["file_id"])
  152. logger.info(f"book_name is:{book_name} need drop", file=sys.stderr)
  153. jso["dropped"] = True
  154. return jso
  155. try:
  156. pdf_intermediate_dict = jso["pdf_intermediate_dict"]
  157. # 将 pdf_intermediate_dict 解压
  158. pdf_intermediate_dict = JsonCompressor.decompress_json(pdf_intermediate_dict)
  159. standard_format = make_standard_format_with_para(pdf_intermediate_dict)
  160. jso["content_list"] = standard_format
  161. logger.info(
  162. f"book_name is:{get_data_source(jso)}/{jso['file_id']},content_list length is {len(standard_format)}",
  163. file=sys.stderr,
  164. )
  165. # 把无用的信息清空
  166. jso["doc_layout_result"] = ""
  167. jso["pdf_intermediate_dict"] = ""
  168. jso["pdf_meta"] = ""
  169. except Exception as e:
  170. jso = exception_handler(jso, e)
  171. return jso
  172. def ocr_parse_pdf_core(pdf_bytes, model_output_json_list, book_name, start_page_id=0, debug_mode=False):
  173. save_path = s3_image_save_path
  174. image_s3_config = get_s3_config(save_path)
  175. start_time = time.time() # 记录开始时间
  176. # 先打印一下book_name和解析开始的时间
  177. logger.info(
  178. f"book_name is:{book_name},start_time is:{formatted_time(start_time)}",
  179. file=sys.stderr,
  180. )
  181. pdf_info_dict = parse_pdf_by_ocr(
  182. pdf_bytes,
  183. model_output_json_list,
  184. save_path,
  185. book_name,
  186. pdf_model_profile=None,
  187. image_s3_config=image_s3_config,
  188. start_page_id=start_page_id,
  189. debug_mode=debug_mode,
  190. )
  191. end_time = time.time() # 记录完成时间
  192. parse_time = int(end_time - start_time) # 计算执行时间
  193. # 解析完成后打印一下book_name和耗时
  194. logger.info(
  195. f"book_name is:{book_name},end_time is:{formatted_time(end_time)},cost_time is:{parse_time}",
  196. file=sys.stderr,
  197. )
  198. return pdf_info_dict, parse_time
  199. # 专门用来跑被drop的pdf,跑完之后需要把need_drop字段置为false
  200. def ocr_dropped_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  201. if not jso.get("need_drop", False):
  202. return jso
  203. else:
  204. try:
  205. pdf_bytes = get_pdf_bytes(jso)
  206. model_output_json_list = jso.get("doc_layout_result")
  207. book_name = get_bookname(jso)
  208. pdf_info_dict, parse_time = ocr_parse_pdf_core(
  209. pdf_bytes, model_output_json_list, book_name, start_page_id=start_page_id, debug_mode=debug_mode
  210. )
  211. jso["pdf_intermediate_dict"] = JsonCompressor.compress_json(pdf_info_dict)
  212. jso["parse_time"] = parse_time
  213. jso["need_drop"] = False
  214. except Exception as e:
  215. jso = exception_handler(jso, e)
  216. return jso
  217. def ocr_parse_pdf(jso: dict, start_page_id=0, debug_mode=False) -> dict:
  218. # 检测debug开关
  219. if debug_mode:
  220. pass
  221. else: # 如果debug没开,则检测是否有needdrop字段
  222. if jso.get("need_drop", False):
  223. return jso
  224. try:
  225. pdf_bytes = get_pdf_bytes(jso)
  226. model_output_json_list = jso.get("doc_layout_result")
  227. book_name = get_bookname(jso)
  228. pdf_info_dict, parse_time = ocr_parse_pdf_core(pdf_bytes, model_output_json_list, book_name,
  229. start_page_id=start_page_id, debug_mode=debug_mode)
  230. jso["pdf_intermediate_dict"] = JsonCompressor.compress_json(pdf_info_dict)
  231. jso["parse_time"] = parse_time
  232. except Exception as e:
  233. jso = exception_handler(jso, e)
  234. return jso