common.py 14 KB

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