common.py 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  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. f_draw_span_bbox=True,
  29. f_draw_layout_bbox=True,
  30. f_dump_md=True,
  31. f_dump_middle_json=True,
  32. f_dump_model_json=True,
  33. f_dump_orig_pdf=True,
  34. f_dump_content_list=False,
  35. f_make_md_mode=MakeMode.MM_MD,
  36. f_draw_model_bbox=False,
  37. ):
  38. orig_model_list = copy.deepcopy(model_list)
  39. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name,
  40. parse_method)
  41. image_writer, md_writer = DiskReaderWriter(
  42. local_image_dir), DiskReaderWriter(local_md_dir)
  43. image_dir = str(os.path.basename(local_image_dir))
  44. if parse_method == 'auto':
  45. jso_useful_key = {'_pdf_type': '', 'model_list': model_list}
  46. pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True)
  47. elif parse_method == 'txt':
  48. pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True)
  49. elif parse_method == 'ocr':
  50. pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True)
  51. else:
  52. logger.error('unknown parse method')
  53. exit(1)
  54. pipe.pipe_classify()
  55. if len(model_list) == 0:
  56. if model_config.__use_inside_model__:
  57. pipe.pipe_analyze()
  58. orig_model_list = copy.deepcopy(pipe.model_list)
  59. else:
  60. logger.error('need model list input')
  61. exit(2)
  62. pipe.pipe_parse()
  63. pdf_info = pipe.pdf_mid_data['pdf_info']
  64. if f_draw_layout_bbox:
  65. draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
  66. if f_draw_span_bbox:
  67. draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
  68. if f_draw_model_bbox:
  69. drow_model_bbox(orig_model_list, pdf_bytes, local_md_dir,
  70. pdf_file_name)
  71. md_content = pipe.pipe_mk_markdown(image_dir,
  72. drop_mode=DropMode.NONE,
  73. md_make_mode=f_make_md_mode)
  74. if f_dump_md:
  75. md_writer.write(
  76. content=md_content,
  77. path=f'{pdf_file_name}.md',
  78. mode=AbsReaderWriter.MODE_TXT,
  79. )
  80. if f_dump_middle_json:
  81. md_writer.write(
  82. content=json_parse.dumps(pipe.pdf_mid_data,
  83. ensure_ascii=False,
  84. indent=4),
  85. path=f'{pdf_file_name}_middle.json',
  86. mode=AbsReaderWriter.MODE_TXT,
  87. )
  88. if f_dump_model_json:
  89. md_writer.write(
  90. content=json_parse.dumps(orig_model_list,
  91. ensure_ascii=False,
  92. indent=4),
  93. path=f'{pdf_file_name}_model.json',
  94. mode=AbsReaderWriter.MODE_TXT,
  95. )
  96. if f_dump_orig_pdf:
  97. md_writer.write(
  98. content=pdf_bytes,
  99. path=f'{pdf_file_name}_origin.pdf',
  100. mode=AbsReaderWriter.MODE_BIN,
  101. )
  102. content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE)
  103. if f_dump_content_list:
  104. md_writer.write(
  105. content=json_parse.dumps(content_list,
  106. ensure_ascii=False,
  107. indent=4),
  108. path=f'{pdf_file_name}_content_list.json',
  109. mode=AbsReaderWriter.MODE_TXT,
  110. )
  111. logger.info(f'local output dir is {local_md_dir}')
  112. parse_pdf_methods = click.Choice(['ocr', 'txt', 'auto'])