common.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. import copy
  2. import json as json_parse
  3. import os
  4. import click
  5. from loguru import logger
  6. import magic_pdf.model as model_config
  7. from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_span_bbox,
  8. drow_model_bbox)
  9. from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
  10. from magic_pdf.pipe.OCRPipe import OCRPipe
  11. from magic_pdf.pipe.TXTPipe import TXTPipe
  12. from magic_pdf.pipe.UNIPipe import UNIPipe
  13. from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
  14. from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
  15. def prepare_env(output_dir, pdf_file_name, method):
  16. local_parent_dir = os.path.join(output_dir, pdf_file_name, method)
  17. local_image_dir = os.path.join(str(local_parent_dir), 'images')
  18. local_md_dir = local_parent_dir
  19. os.makedirs(local_image_dir, exist_ok=True)
  20. os.makedirs(local_md_dir, exist_ok=True)
  21. return local_image_dir, local_md_dir
  22. def do_parse(
  23. output_dir,
  24. pdf_file_name,
  25. pdf_bytes,
  26. model_list,
  27. parse_method,
  28. debug_able,
  29. f_draw_span_bbox=True,
  30. f_draw_layout_bbox=True,
  31. f_dump_md=True,
  32. f_dump_middle_json=True,
  33. f_dump_model_json=True,
  34. f_dump_orig_pdf=True,
  35. f_dump_content_list=False,
  36. f_make_md_mode=MakeMode.MM_MD,
  37. f_draw_model_bbox=False,
  38. ):
  39. if debug_able:
  40. logger.warning("debug mode is on")
  41. f_dump_content_list = True
  42. f_draw_model_bbox = True
  43. orig_model_list = copy.deepcopy(model_list)
  44. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name,
  45. parse_method)
  46. image_writer, md_writer = DiskReaderWriter(
  47. local_image_dir), DiskReaderWriter(local_md_dir)
  48. image_dir = str(os.path.basename(local_image_dir))
  49. if parse_method == 'auto':
  50. jso_useful_key = {'_pdf_type': '', 'model_list': model_list}
  51. pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True)
  52. elif parse_method == 'txt':
  53. pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True)
  54. elif parse_method == 'ocr':
  55. pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True)
  56. else:
  57. logger.error('unknown parse method')
  58. exit(1)
  59. pipe.pipe_classify()
  60. if len(model_list) == 0:
  61. if model_config.__use_inside_model__:
  62. pipe.pipe_analyze()
  63. orig_model_list = copy.deepcopy(pipe.model_list)
  64. else:
  65. logger.error('need model list input')
  66. exit(2)
  67. pipe.pipe_parse()
  68. pdf_info = pipe.pdf_mid_data['pdf_info']
  69. if f_draw_layout_bbox:
  70. draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
  71. if f_draw_span_bbox:
  72. draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
  73. if f_draw_model_bbox:
  74. drow_model_bbox(copy.deepcopy(orig_model_list), pdf_bytes, local_md_dir, pdf_file_name)
  75. md_content = pipe.pipe_mk_markdown(image_dir,
  76. drop_mode=DropMode.NONE,
  77. md_make_mode=f_make_md_mode)
  78. if f_dump_md:
  79. md_writer.write(
  80. content=md_content,
  81. path=f'{pdf_file_name}.md',
  82. mode=AbsReaderWriter.MODE_TXT,
  83. )
  84. if f_dump_middle_json:
  85. md_writer.write(
  86. content=json_parse.dumps(pipe.pdf_mid_data,
  87. ensure_ascii=False,
  88. indent=4),
  89. path=f'{pdf_file_name}_middle.json',
  90. mode=AbsReaderWriter.MODE_TXT,
  91. )
  92. if f_dump_model_json:
  93. md_writer.write(
  94. content=json_parse.dumps(orig_model_list,
  95. ensure_ascii=False,
  96. indent=4),
  97. path=f'{pdf_file_name}_model.json',
  98. mode=AbsReaderWriter.MODE_TXT,
  99. )
  100. if f_dump_orig_pdf:
  101. md_writer.write(
  102. content=pdf_bytes,
  103. path=f'{pdf_file_name}_origin.pdf',
  104. mode=AbsReaderWriter.MODE_BIN,
  105. )
  106. content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE)
  107. if f_dump_content_list:
  108. md_writer.write(
  109. content=json_parse.dumps(content_list,
  110. ensure_ascii=False,
  111. indent=4),
  112. path=f'{pdf_file_name}_content_list.json',
  113. mode=AbsReaderWriter.MODE_TXT,
  114. )
  115. logger.info(f'local output dir is {local_md_dir}')
  116. parse_pdf_methods = click.Choice(['ocr', 'txt', 'auto'])