demo.py 916 B

12345678910111213141516171819202122
  1. import os
  2. import json
  3. from magic_pdf.pipe.UNIPipe import UNIPipe
  4. from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
  5. current_script_dir = os.path.dirname(os.path.abspath(__file__))
  6. demo_name = "demo1"
  7. pdf_path = os.path.join(current_script_dir, f"{demo_name}.pdf")
  8. model_path = os.path.join(current_script_dir, f"{demo_name}.json")
  9. pdf_bytes = open(pdf_path, "rb").read()
  10. model_json = json.loads(open(model_path, "r", encoding="utf-8").read())
  11. jso_useful_key = {"_pdf_type": "", "model_list": model_json}
  12. local_image_dir = os.path.join(current_script_dir, 'images')
  13. image_dir = str(os.path.basename(local_image_dir))
  14. image_writer = DiskReaderWriter(local_image_dir)
  15. pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer)
  16. pipe.pipe_classify()
  17. pipe.pipe_parse()
  18. md_content = pipe.pipe_mk_markdown(image_dir, drop_mode="none")
  19. with open(f"{demo_name}.md", "w", encoding="utf-8") as f:
  20. f.write(md_content)