demo_test.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. import json
  2. import os
  3. import sys
  4. from pathlib import Path
  5. import click
  6. from magic_pdf.pipeline import (
  7. meta_scan,
  8. classify_by_type,
  9. parse_pdf,
  10. pdf_intermediate_dict_to_markdown,
  11. save_tables_to_s3,
  12. )
  13. from magic_pdf.libs.commons import join_path, read_file, json_dump_path
  14. from app.common.s3 import get_s3_config
  15. from loguru import logger
  16. local_json_path = "Z:/format.json"
  17. local_jsonl_path = "Z:/format.jsonl"
  18. def get_json_from_local_or_s3(book_name=None):
  19. if book_name is None:
  20. with open(local_json_path, "r", encoding="utf-8") as json_file:
  21. json_line = json_file.read()
  22. json_object = json.loads(json_line)
  23. else:
  24. # error_log_path & json_dump_path
  25. # 可配置从上述两个地址获取源json
  26. json_path = join_path(json_dump_path, book_name + ".json")
  27. s3_config = get_s3_config(json_path)
  28. file_content = read_file(json_path, s3_config)
  29. json_str = file_content.decode("utf-8")
  30. logger.info(json_str)
  31. json_object = json.loads(json_str)
  32. return json_object
  33. def write_json_to_local(jso, book_name=None):
  34. if book_name is None:
  35. with open(local_json_path, "w", encoding="utf-8") as file:
  36. file.write(json.dumps(jso, ensure_ascii=False))
  37. else:
  38. pass
  39. def demo_parse_pdf(book_name=None, start_page_id=0, debug_mode=True):
  40. json_object = get_json_from_local_or_s3(book_name)
  41. jso = parse_pdf(json_object, start_page_id=start_page_id, debug_mode=debug_mode)
  42. logger.info(f"pdf_parse_time: {jso['parse_time']}")
  43. write_json_to_local(jso, book_name)
  44. jso_md = pdf_intermediate_dict_to_markdown(jso, debug_mode=debug_mode)
  45. md_content = jso_md.get("content")
  46. if book_name is not None:
  47. save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
  48. markdown_save_path = join_path(save_tmp_path, "md", book_name + ".md")
  49. with open(markdown_save_path, "w", encoding="utf-8") as f:
  50. f.write(md_content)
  51. else:
  52. logger.info(md_content)
  53. def demo_save_tables(book_name=None, start_page_id=0, debug_mode=True):
  54. json_object = get_json_from_local_or_s3(book_name)
  55. jso = parse_pdf(json_object, start_page_id=start_page_id, debug_mode=debug_mode)
  56. logger.info(f"pdf_parse_time: {jso['parse_time']}")
  57. write_json_to_local(jso, book_name)
  58. save_tables_to_s3(jso, debug_mode=debug_mode)
  59. def demo_classify_by_type(book_name=None, debug_mode=True):
  60. json_object = get_json_from_local_or_s3(book_name)
  61. jso = classify_by_type(json_object, debug_mode=debug_mode)
  62. logger.info(json.dumps(jso, ensure_ascii=False))
  63. logger.info(f"classify_time: {jso['classify_time']}")
  64. write_json_to_local(jso, book_name)
  65. def demo_meta_scan(book_name=None, debug_mode=True):
  66. json_object = get_json_from_local_or_s3(book_name)
  67. # doc_layout_check=False
  68. jso = meta_scan(json_object, doc_layout_check=True)
  69. logger.info(json.dumps(jso, ensure_ascii=False))
  70. logger.info(f"meta_scan_time: {jso['meta_scan_time']}")
  71. write_json_to_local(jso, book_name)
  72. def demo_meta_scan_from_jsonl():
  73. with open(local_jsonl_path, "r", encoding="utf-8") as jsonl_file:
  74. for line in jsonl_file:
  75. jso = json.loads(line)
  76. jso = meta_scan(jso)
  77. logger.info(f"pdf_path: {jso['content']['pdf_path']}")
  78. logger.info(f"read_file_time: {jso['read_file_time']}")
  79. logger.info(f"meta_scan_time: {jso['meta_scan_time']}")
  80. def demo_test5():
  81. with open(local_json_path, "r", encoding="utf-8") as json_file:
  82. json_line = json_file.read()
  83. jso = json.loads(json_line)
  84. img_list_len = len(jso["content"]["image_info_per_page"])
  85. logger.info(f"img_list_len: {img_list_len}")
  86. def read_more_para_test_samples(type="scihub"):
  87. # 读取多段落测试样本
  88. curr_dir = Path(__file__).parent
  89. files_path = ""
  90. if type == "gift":
  91. relative_path = "../tests/assets/more_para_test_samples/gift_files.txt"
  92. files_path = os.path.join(curr_dir, relative_path)
  93. if type == "scihub":
  94. relative_path = "../tests/assets/more_para_test_samples/scihub_files.txt"
  95. files_path = os.path.join(curr_dir, relative_path)
  96. if type == "zlib":
  97. relative_path = "../tests/assets/more_para_test_samples/zlib_files.txt"
  98. files_path = os.path.join(curr_dir, relative_path)
  99. # check if file exists
  100. if not os.path.exists(files_path):
  101. print("File not exist!")
  102. sys.exit(0)
  103. with open(files_path, "r", encoding="utf-8") as f:
  104. lines = f.readlines()
  105. # print("lines", lines)
  106. return lines
  107. def batch_test_more_para(type="scihub"):
  108. # 批量测试多段落
  109. para_test_files = read_more_para_test_samples(type)
  110. for file in para_test_files:
  111. file = file.strip()
  112. print(file)
  113. demo_parse_pdf(book_name=file)
  114. @click.command()
  115. @click.option("--book-name", help="s3上pdf文件的路径")
  116. def main(book_name: str):
  117. demo_parse_pdf(book_name, start_page_id=0)
  118. if __name__ == "__main__":
  119. main()