ocr_demo.py 3.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. import json
  2. import os
  3. from loguru import logger
  4. from pathlib import Path
  5. from magic_pdf.pipeline_ocr import ocr_parse_pdf_core
  6. from magic_pdf.spark.s3 import get_s3_config
  7. from demo.demo_commons import get_json_from_local_or_s3
  8. from magic_pdf.dict2md.ocr_mkcontent import (
  9. ocr_mk_mm_markdown_with_para,
  10. make_standard_format_with_para
  11. )
  12. from magic_pdf.libs.commons import join_path, read_file
  13. from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
  14. def save_markdown(markdown_text, input_filepath):
  15. # 获取输入文件的目录
  16. directory = os.path.dirname(input_filepath)
  17. # 获取输入文件的文件名(不带扩展名)
  18. base_name = os.path.basename(input_filepath)
  19. file_name_without_ext = os.path.splitext(base_name)[0]
  20. # 定义输出文件的路径
  21. output_filepath = os.path.join(directory, f"{file_name_without_ext}.md")
  22. # 将Markdown文本写入.md文件
  23. with open(output_filepath, 'w', encoding='utf-8') as file:
  24. file.write(markdown_text)
  25. def read_json_file(file_path):
  26. with open(file_path, 'r') as f:
  27. data = json.load(f)
  28. return data
  29. def ocr_local_parse(ocr_pdf_path, ocr_json_file_path):
  30. try:
  31. ocr_pdf_model_info = read_json_file(ocr_json_file_path)
  32. pth = Path(ocr_json_file_path)
  33. book_name = pth.name
  34. pdf_bytes = read_file(ocr_pdf_path, None)
  35. ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info)
  36. except Exception as e:
  37. logger.exception(e)
  38. def ocr_online_parse(book_name, start_page_id=0, debug_mode=True):
  39. try:
  40. json_object = get_json_from_local_or_s3(book_name)
  41. # logger.info(json_object)
  42. s3_pdf_path = json_object["file_location"]
  43. s3_config = get_s3_config(s3_pdf_path)
  44. pdf_bytes = read_file(s3_pdf_path, s3_config)
  45. ocr_pdf_model_info = json_object.get("doc_layout_result")
  46. ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info)
  47. except Exception as e:
  48. logger.exception(e)
  49. def ocr_parse_core(book_name, pdf_bytes, ocr_pdf_model_info, start_page_id=0):
  50. save_tmp_path = os.path.join(os.path.dirname(__file__), "../..", "tmp", "unittest")
  51. save_path = join_path(save_tmp_path, "md")
  52. save_path_with_bookname = os.path.join(save_path, book_name)
  53. text_content_save_path = f"{save_path_with_bookname}/book.md"
  54. pdf_info_dict, parse_time = ocr_parse_pdf_core(pdf_bytes, ocr_pdf_model_info, book_name, start_page_id=start_page_id, debug_mode=True)
  55. parent_dir = os.path.dirname(text_content_save_path)
  56. if not os.path.exists(parent_dir):
  57. os.makedirs(parent_dir)
  58. # markdown_content = mk_nlp_markdown(pdf_info_dict)
  59. markdown_content = ocr_mk_mm_markdown_with_para(pdf_info_dict)
  60. # markdown_pagination = ocr_mk_mm_markdown_with_para_and_pagination(pdf_info_dict)
  61. with open(text_content_save_path, "w", encoding="utf-8") as f:
  62. f.write(markdown_content)
  63. standard_format = make_standard_format_with_para(pdf_info_dict)
  64. standard_format_save_path = f"{save_path_with_bookname}/standard_format.txt"
  65. with open(standard_format_save_path, "w", encoding="utf-8") as f:
  66. # 将standard_format dump成json文本并保存
  67. f.write(json.dumps(standard_format, ensure_ascii=False))
  68. if __name__ == '__main__':
  69. pdf_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.pdf"
  70. json_file_path = r"/home/cxu/workspace/Magic-PDF/ocr_demo/j.1540-627x.2006.00176.x.json"
  71. # ocr_local_parse(pdf_path, json_file_path)
  72. book_name = "科数网/edu_00011318"
  73. ocr_online_parse(book_name)
  74. pass