demo.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241
  1. # Copyright (c) Opendatalab. All rights reserved.
  2. import copy
  3. import json
  4. import os
  5. from pathlib import Path
  6. from loguru import logger
  7. from mineru.cli.common import prepare_env, read_fn
  8. from mineru.utils.pdf_page_tools import convert_pdf_bytes_to_bytes_by_pypdfium2
  9. from mineru.data.data_reader_writer import FileBasedDataWriter
  10. from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox
  11. from mineru.utils.enum_class import MakeMode
  12. from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
  13. from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
  14. from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
  15. from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
  16. from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
  17. from mineru.utils.guess_suffix_or_lang import guess_suffix_by_path
  18. def do_parse(
  19. output_dir, # Output directory for storing parsing results
  20. pdf_file_names: list[str], # List of PDF file names to be parsed
  21. pdf_bytes_list: list[bytes], # List of PDF bytes to be parsed
  22. p_lang_list: list[str], # List of languages for each PDF, default is 'ch' (Chinese)
  23. backend="pipeline", # The backend for parsing PDF, default is 'pipeline'
  24. parse_method="auto", # The method for parsing PDF, default is 'auto'
  25. formula_enable=True, # Enable formula parsing
  26. table_enable=True, # Enable table parsing
  27. server_url=None, # Server URL for vlm-http-client backend
  28. f_draw_layout_bbox=True, # Whether to draw layout bounding boxes
  29. f_draw_span_bbox=True, # Whether to draw span bounding boxes
  30. f_dump_md=True, # Whether to dump markdown files
  31. f_dump_middle_json=True, # Whether to dump middle JSON files
  32. f_dump_model_output=True, # Whether to dump model output files
  33. f_dump_orig_pdf=True, # Whether to dump original PDF files
  34. f_dump_content_list=True, # Whether to dump content list files
  35. f_make_md_mode=MakeMode.MM_MD, # The mode for making markdown content, default is MM_MD
  36. start_page_id=0, # Start page ID for parsing, default is 0
  37. end_page_id=None, # End page ID for parsing, default is None (parse all pages until the end of the document)
  38. ):
  39. if backend == "pipeline":
  40. for idx, pdf_bytes in enumerate(pdf_bytes_list):
  41. new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
  42. pdf_bytes_list[idx] = new_pdf_bytes
  43. infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = pipeline_doc_analyze(pdf_bytes_list, p_lang_list, parse_method=parse_method, formula_enable=formula_enable,table_enable=table_enable)
  44. for idx, model_list in enumerate(infer_results):
  45. model_json = copy.deepcopy(model_list)
  46. pdf_file_name = pdf_file_names[idx]
  47. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
  48. image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
  49. images_list = all_image_lists[idx]
  50. pdf_doc = all_pdf_docs[idx]
  51. _lang = lang_list[idx]
  52. _ocr_enable = ocr_enabled_list[idx]
  53. middle_json = pipeline_result_to_middle_json(model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, formula_enable)
  54. pdf_info = middle_json["pdf_info"]
  55. pdf_bytes = pdf_bytes_list[idx]
  56. _process_output(
  57. pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
  58. md_writer, f_draw_layout_bbox, f_draw_span_bbox, f_dump_orig_pdf,
  59. f_dump_md, f_dump_content_list, f_dump_middle_json, f_dump_model_output,
  60. f_make_md_mode, middle_json, model_json, is_pipeline=True
  61. )
  62. else:
  63. if backend.startswith("vlm-"):
  64. backend = backend[4:]
  65. f_draw_span_bbox = False
  66. parse_method = "vlm"
  67. for idx, pdf_bytes in enumerate(pdf_bytes_list):
  68. pdf_file_name = pdf_file_names[idx]
  69. pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
  70. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
  71. image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
  72. middle_json, infer_result = vlm_doc_analyze(pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url)
  73. pdf_info = middle_json["pdf_info"]
  74. _process_output(
  75. pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
  76. md_writer, f_draw_layout_bbox, f_draw_span_bbox, f_dump_orig_pdf,
  77. f_dump_md, f_dump_content_list, f_dump_middle_json, f_dump_model_output,
  78. f_make_md_mode, middle_json, infer_result, is_pipeline=False
  79. )
  80. def _process_output(
  81. pdf_info,
  82. pdf_bytes,
  83. pdf_file_name,
  84. local_md_dir,
  85. local_image_dir,
  86. md_writer,
  87. f_draw_layout_bbox,
  88. f_draw_span_bbox,
  89. f_dump_orig_pdf,
  90. f_dump_md,
  91. f_dump_content_list,
  92. f_dump_middle_json,
  93. f_dump_model_output,
  94. f_make_md_mode,
  95. middle_json,
  96. model_output=None,
  97. is_pipeline=True
  98. ):
  99. """处理输出文件"""
  100. if f_draw_layout_bbox:
  101. draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_layout.pdf")
  102. if f_draw_span_bbox:
  103. draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_span.pdf")
  104. if f_dump_orig_pdf:
  105. md_writer.write(
  106. f"{pdf_file_name}_origin.pdf",
  107. pdf_bytes,
  108. )
  109. image_dir = str(os.path.basename(local_image_dir))
  110. if f_dump_md:
  111. make_func = pipeline_union_make if is_pipeline else vlm_union_make
  112. md_content_str = make_func(pdf_info, f_make_md_mode, image_dir)
  113. md_writer.write_string(
  114. f"{pdf_file_name}.md",
  115. md_content_str,
  116. )
  117. if f_dump_content_list:
  118. make_func = pipeline_union_make if is_pipeline else vlm_union_make
  119. content_list = make_func(pdf_info, MakeMode.CONTENT_LIST, image_dir)
  120. md_writer.write_string(
  121. f"{pdf_file_name}_content_list.json",
  122. json.dumps(content_list, ensure_ascii=False, indent=4),
  123. )
  124. if f_dump_middle_json:
  125. md_writer.write_string(
  126. f"{pdf_file_name}_middle.json",
  127. json.dumps(middle_json, ensure_ascii=False, indent=4),
  128. )
  129. if f_dump_model_output:
  130. md_writer.write_string(
  131. f"{pdf_file_name}_model.json",
  132. json.dumps(model_output, ensure_ascii=False, indent=4),
  133. )
  134. logger.info(f"local output dir is {local_md_dir}")
  135. def parse_doc(
  136. path_list: list[Path],
  137. output_dir,
  138. lang="ch",
  139. backend="pipeline",
  140. method="auto",
  141. server_url=None,
  142. start_page_id=0,
  143. end_page_id=None
  144. ):
  145. """
  146. Parameter description:
  147. path_list: List of document paths to be parsed, can be PDF or image files.
  148. output_dir: Output directory for storing parsing results.
  149. lang: Language option, default is 'ch', optional values include['ch', 'ch_server', 'ch_lite', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']。
  150. Input the languages in the pdf (if known) to improve OCR accuracy. Optional.
  151. Adapted only for the case where the backend is set to "pipeline"
  152. backend: the backend for parsing pdf:
  153. pipeline: More general.
  154. vlm-transformers: More general.
  155. vlm-vllm-engine: Faster(engine).
  156. vlm-http-client: Faster(client).
  157. without method specified, pipeline will be used by default.
  158. method: the method for parsing pdf:
  159. auto: Automatically determine the method based on the file type.
  160. txt: Use text extraction method.
  161. ocr: Use OCR method for image-based PDFs.
  162. Without method specified, 'auto' will be used by default.
  163. Adapted only for the case where the backend is set to "pipeline".
  164. server_url: When the backend is `http-client`, you need to specify the server_url, for example:`http://127.0.0.1:30000`
  165. start_page_id: Start page ID for parsing, default is 0
  166. end_page_id: End page ID for parsing, default is None (parse all pages until the end of the document)
  167. """
  168. try:
  169. file_name_list = []
  170. pdf_bytes_list = []
  171. lang_list = []
  172. for path in path_list:
  173. file_name = str(Path(path).stem)
  174. pdf_bytes = read_fn(path)
  175. file_name_list.append(file_name)
  176. pdf_bytes_list.append(pdf_bytes)
  177. lang_list.append(lang)
  178. do_parse(
  179. output_dir=output_dir,
  180. pdf_file_names=file_name_list,
  181. pdf_bytes_list=pdf_bytes_list,
  182. p_lang_list=lang_list,
  183. backend=backend,
  184. parse_method=method,
  185. server_url=server_url,
  186. start_page_id=start_page_id,
  187. end_page_id=end_page_id
  188. )
  189. except Exception as e:
  190. logger.exception(e)
  191. if __name__ == '__main__':
  192. # args
  193. __dir__ = os.path.dirname(os.path.abspath(__file__))
  194. pdf_files_dir = os.path.join(__dir__, "pdfs")
  195. output_dir = os.path.join(__dir__, "output")
  196. pdf_suffixes = ["pdf"]
  197. image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg"]
  198. doc_path_list = []
  199. for doc_path in Path(pdf_files_dir).glob('*'):
  200. if guess_suffix_by_path(doc_path) in pdf_suffixes + image_suffixes:
  201. doc_path_list.append(doc_path)
  202. """如果您由于网络问题无法下载模型,可以设置环境变量MINERU_MODEL_SOURCE为modelscope使用免代理仓库下载模型"""
  203. # os.environ['MINERU_MODEL_SOURCE'] = "modelscope"
  204. """Use pipeline mode if your environment does not support VLM"""
  205. parse_doc(doc_path_list, output_dir, backend="pipeline")
  206. """To enable VLM mode, change the backend to 'vlm-xxx'"""
  207. # parse_doc(doc_path_list, output_dir, backend="vlm-transformers") # more general.
  208. # parse_doc(doc_path_list, output_dir, backend="vlm-mlx-engine") # faster than transformers in macOS 13.5+.
  209. # parse_doc(doc_path_list, output_dir, backend="vlm-vllm-engine") # faster(engine).
  210. # parse_doc(doc_path_list, output_dir, backend="vlm-http-client", server_url="http://127.0.0.1:30000") # faster(client).