common.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375
  1. # Copyright (c) Opendatalab. All rights reserved.
  2. import json
  3. import os
  4. import copy
  5. from pathlib import Path
  6. from loguru import logger
  7. from mineru.data.data_reader_writer import FileBasedDataWriter
  8. from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox, draw_line_sort_bbox
  9. from mineru.utils.enum_class import MakeMode
  10. from mineru.utils.guess_suffix_or_lang import guess_suffix_by_bytes
  11. from mineru.utils.pdf_image_tools import images_bytes_to_pdf_bytes
  12. from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
  13. from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
  14. from mineru.backend.vlm.vlm_analyze import aio_doc_analyze as aio_vlm_doc_analyze
  15. from mineru.utils.pdf_page_tools import convert_pdf_bytes_to_bytes_by_pypdfium2
  16. pdf_suffixes = ["pdf"]
  17. image_suffixes = ["png", "jpeg", "jp2", "webp", "gif", "bmp", "jpg", "tiff"]
  18. os.environ["TOKENIZERS_PARALLELISM"] = "false"
  19. def read_fn(path):
  20. if not isinstance(path, Path):
  21. path = Path(path)
  22. with open(str(path), "rb") as input_file:
  23. file_bytes = input_file.read()
  24. file_suffix = guess_suffix_by_bytes(file_bytes, path)
  25. if file_suffix in image_suffixes:
  26. return images_bytes_to_pdf_bytes(file_bytes)
  27. elif file_suffix in pdf_suffixes:
  28. return file_bytes
  29. else:
  30. raise Exception(f"Unknown file suffix: {file_suffix}")
  31. def prepare_env(output_dir, pdf_file_name, parse_method):
  32. local_md_dir = str(os.path.join(output_dir, pdf_file_name, parse_method))
  33. local_image_dir = os.path.join(str(local_md_dir), "images")
  34. os.makedirs(local_image_dir, exist_ok=True)
  35. os.makedirs(local_md_dir, exist_ok=True)
  36. return local_image_dir, local_md_dir
  37. def _prepare_pdf_bytes(pdf_bytes_list, start_page_id, end_page_id):
  38. """准备处理PDF字节数据"""
  39. result = []
  40. for pdf_bytes in pdf_bytes_list:
  41. new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
  42. result.append(new_pdf_bytes)
  43. return result
  44. def _process_output(
  45. pdf_info,
  46. pdf_bytes,
  47. pdf_file_name,
  48. local_md_dir,
  49. local_image_dir,
  50. md_writer,
  51. f_draw_layout_bbox,
  52. f_draw_span_bbox,
  53. f_dump_orig_pdf,
  54. f_dump_md,
  55. f_dump_content_list,
  56. f_dump_middle_json,
  57. f_dump_model_output,
  58. f_make_md_mode,
  59. middle_json,
  60. model_output=None,
  61. is_pipeline=True
  62. ):
  63. f_draw_line_sort_bbox = False
  64. from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
  65. """处理输出文件"""
  66. if f_draw_layout_bbox:
  67. draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_layout.pdf")
  68. if f_draw_span_bbox:
  69. draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_span.pdf")
  70. if f_dump_orig_pdf:
  71. md_writer.write(
  72. f"{pdf_file_name}_origin.pdf",
  73. pdf_bytes,
  74. )
  75. if f_draw_line_sort_bbox:
  76. draw_line_sort_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_line_sort.pdf")
  77. image_dir = str(os.path.basename(local_image_dir))
  78. if f_dump_md:
  79. make_func = pipeline_union_make if is_pipeline else vlm_union_make
  80. md_content_str = make_func(pdf_info, f_make_md_mode, image_dir)
  81. md_writer.write_string(
  82. f"{pdf_file_name}.md",
  83. md_content_str,
  84. )
  85. if f_dump_content_list:
  86. make_func = pipeline_union_make if is_pipeline else vlm_union_make
  87. content_list = make_func(pdf_info, MakeMode.CONTENT_LIST, image_dir)
  88. md_writer.write_string(
  89. f"{pdf_file_name}_content_list.json",
  90. json.dumps(content_list, ensure_ascii=False, indent=4),
  91. )
  92. if f_dump_middle_json:
  93. md_writer.write_string(
  94. f"{pdf_file_name}_middle.json",
  95. json.dumps(middle_json, ensure_ascii=False, indent=4),
  96. )
  97. if f_dump_model_output:
  98. md_writer.write_string(
  99. f"{pdf_file_name}_model.json",
  100. json.dumps(model_output, ensure_ascii=False, indent=4),
  101. )
  102. logger.info(f"local output dir is {local_md_dir}")
  103. def _process_pipeline(
  104. output_dir,
  105. pdf_file_names,
  106. pdf_bytes_list,
  107. p_lang_list,
  108. parse_method,
  109. p_formula_enable,
  110. p_table_enable,
  111. f_draw_layout_bbox,
  112. f_draw_span_bbox,
  113. f_dump_md,
  114. f_dump_middle_json,
  115. f_dump_model_output,
  116. f_dump_orig_pdf,
  117. f_dump_content_list,
  118. f_make_md_mode,
  119. ):
  120. """处理pipeline后端逻辑"""
  121. from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
  122. from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
  123. infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = (
  124. pipeline_doc_analyze(
  125. pdf_bytes_list, p_lang_list, parse_method=parse_method,
  126. formula_enable=p_formula_enable, table_enable=p_table_enable
  127. )
  128. )
  129. for idx, model_list in enumerate(infer_results):
  130. model_json = copy.deepcopy(model_list)
  131. pdf_file_name = pdf_file_names[idx]
  132. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
  133. image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
  134. images_list = all_image_lists[idx]
  135. pdf_doc = all_pdf_docs[idx]
  136. _lang = lang_list[idx]
  137. _ocr_enable = ocr_enabled_list[idx]
  138. middle_json = pipeline_result_to_middle_json(
  139. model_list, images_list, pdf_doc, image_writer,
  140. _lang, _ocr_enable, p_formula_enable
  141. )
  142. pdf_info = middle_json["pdf_info"]
  143. pdf_bytes = pdf_bytes_list[idx]
  144. _process_output(
  145. pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
  146. md_writer, f_draw_layout_bbox, f_draw_span_bbox, f_dump_orig_pdf,
  147. f_dump_md, f_dump_content_list, f_dump_middle_json, f_dump_model_output,
  148. f_make_md_mode, middle_json, model_json, is_pipeline=True
  149. )
  150. async def _async_process_vlm(
  151. output_dir,
  152. pdf_file_names,
  153. pdf_bytes_list,
  154. backend,
  155. f_draw_layout_bbox,
  156. f_draw_span_bbox,
  157. f_dump_md,
  158. f_dump_middle_json,
  159. f_dump_model_output,
  160. f_dump_orig_pdf,
  161. f_dump_content_list,
  162. f_make_md_mode,
  163. server_url=None,
  164. **kwargs,
  165. ):
  166. """异步处理VLM后端逻辑"""
  167. parse_method = "vlm"
  168. f_draw_span_bbox = False
  169. if not backend.endswith("client"):
  170. server_url = None
  171. for idx, pdf_bytes in enumerate(pdf_bytes_list):
  172. pdf_file_name = pdf_file_names[idx]
  173. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
  174. image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
  175. middle_json, infer_result = await aio_vlm_doc_analyze(
  176. pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url, **kwargs,
  177. )
  178. pdf_info = middle_json["pdf_info"]
  179. _process_output(
  180. pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
  181. md_writer, f_draw_layout_bbox, f_draw_span_bbox, f_dump_orig_pdf,
  182. f_dump_md, f_dump_content_list, f_dump_middle_json, f_dump_model_output,
  183. f_make_md_mode, middle_json, infer_result, is_pipeline=False
  184. )
  185. def _process_vlm(
  186. output_dir,
  187. pdf_file_names,
  188. pdf_bytes_list,
  189. backend,
  190. f_draw_layout_bbox,
  191. f_draw_span_bbox,
  192. f_dump_md,
  193. f_dump_middle_json,
  194. f_dump_model_output,
  195. f_dump_orig_pdf,
  196. f_dump_content_list,
  197. f_make_md_mode,
  198. server_url=None,
  199. **kwargs,
  200. ):
  201. """同步处理VLM后端逻辑"""
  202. parse_method = "vlm"
  203. f_draw_span_bbox = False
  204. if not backend.endswith("client"):
  205. server_url = None
  206. for idx, pdf_bytes in enumerate(pdf_bytes_list):
  207. pdf_file_name = pdf_file_names[idx]
  208. local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
  209. image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
  210. middle_json, infer_result = vlm_doc_analyze(
  211. pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url, **kwargs,
  212. )
  213. pdf_info = middle_json["pdf_info"]
  214. _process_output(
  215. pdf_info, pdf_bytes, pdf_file_name, local_md_dir, local_image_dir,
  216. md_writer, f_draw_layout_bbox, f_draw_span_bbox, f_dump_orig_pdf,
  217. f_dump_md, f_dump_content_list, f_dump_middle_json, f_dump_model_output,
  218. f_make_md_mode, middle_json, infer_result, is_pipeline=False
  219. )
  220. def do_parse(
  221. output_dir,
  222. pdf_file_names: list[str],
  223. pdf_bytes_list: list[bytes],
  224. p_lang_list: list[str],
  225. backend="pipeline",
  226. parse_method="auto",
  227. formula_enable=True,
  228. table_enable=True,
  229. server_url=None,
  230. f_draw_layout_bbox=True,
  231. f_draw_span_bbox=True,
  232. f_dump_md=True,
  233. f_dump_middle_json=True,
  234. f_dump_model_output=True,
  235. f_dump_orig_pdf=True,
  236. f_dump_content_list=True,
  237. f_make_md_mode=MakeMode.MM_MD,
  238. start_page_id=0,
  239. end_page_id=None,
  240. **kwargs,
  241. ):
  242. # 预处理PDF字节数据
  243. pdf_bytes_list = _prepare_pdf_bytes(pdf_bytes_list, start_page_id, end_page_id)
  244. if backend == "pipeline":
  245. _process_pipeline(
  246. output_dir, pdf_file_names, pdf_bytes_list, p_lang_list,
  247. parse_method, formula_enable, table_enable,
  248. f_draw_layout_bbox, f_draw_span_bbox, f_dump_md, f_dump_middle_json,
  249. f_dump_model_output, f_dump_orig_pdf, f_dump_content_list, f_make_md_mode
  250. )
  251. else:
  252. if backend.startswith("vlm-"):
  253. backend = backend[4:]
  254. if backend == "vllm-async-engine":
  255. raise Exception("vlm-vllm-async-engine backend is not supported in sync mode, please use vlm-vllm-engine backend")
  256. os.environ['MINERU_VLM_FORMULA_ENABLE'] = str(formula_enable)
  257. os.environ['MINERU_VLM_TABLE_ENABLE'] = str(table_enable)
  258. _process_vlm(
  259. output_dir, pdf_file_names, pdf_bytes_list, backend,
  260. f_draw_layout_bbox, f_draw_span_bbox, f_dump_md, f_dump_middle_json,
  261. f_dump_model_output, f_dump_orig_pdf, f_dump_content_list, f_make_md_mode,
  262. server_url, **kwargs,
  263. )
  264. async def aio_do_parse(
  265. output_dir,
  266. pdf_file_names: list[str],
  267. pdf_bytes_list: list[bytes],
  268. p_lang_list: list[str],
  269. backend="pipeline",
  270. parse_method="auto",
  271. formula_enable=True,
  272. table_enable=True,
  273. server_url=None,
  274. f_draw_layout_bbox=True,
  275. f_draw_span_bbox=True,
  276. f_dump_md=True,
  277. f_dump_middle_json=True,
  278. f_dump_model_output=True,
  279. f_dump_orig_pdf=True,
  280. f_dump_content_list=True,
  281. f_make_md_mode=MakeMode.MM_MD,
  282. start_page_id=0,
  283. end_page_id=None,
  284. **kwargs,
  285. ):
  286. # 预处理PDF字节数据
  287. pdf_bytes_list = _prepare_pdf_bytes(pdf_bytes_list, start_page_id, end_page_id)
  288. if backend == "pipeline":
  289. # pipeline模式暂不支持异步,使用同步处理方式
  290. _process_pipeline(
  291. output_dir, pdf_file_names, pdf_bytes_list, p_lang_list,
  292. parse_method, formula_enable, table_enable,
  293. f_draw_layout_bbox, f_draw_span_bbox, f_dump_md, f_dump_middle_json,
  294. f_dump_model_output, f_dump_orig_pdf, f_dump_content_list, f_make_md_mode
  295. )
  296. else:
  297. if backend.startswith("vlm-"):
  298. backend = backend[4:]
  299. if backend == "vllm-engine":
  300. raise Exception("vlm-vllm-engine backend is not supported in async mode, please use vlm-vllm-async-engine backend")
  301. os.environ['MINERU_VLM_FORMULA_ENABLE'] = str(formula_enable)
  302. os.environ['MINERU_VLM_TABLE_ENABLE'] = str(table_enable)
  303. await _async_process_vlm(
  304. output_dir, pdf_file_names, pdf_bytes_list, backend,
  305. f_draw_layout_bbox, f_draw_span_bbox, f_dump_md, f_dump_middle_json,
  306. f_dump_model_output, f_dump_orig_pdf, f_dump_content_list, f_make_md_mode,
  307. server_url, **kwargs,
  308. )
  309. if __name__ == "__main__":
  310. # pdf_path = "../../demo/pdfs/demo3.pdf"
  311. pdf_path = "C:/Users/zhaoxiaomeng/Downloads/4546d0e2-ba60-40a5-a17e-b68555cec741.pdf"
  312. try:
  313. do_parse("./output", [Path(pdf_path).stem], [read_fn(Path(pdf_path))],["ch"],
  314. end_page_id=10,
  315. backend='vlm-huggingface'
  316. # backend = 'pipeline'
  317. )
  318. except Exception as e:
  319. logger.exception(e)