UNIPipe.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. import json
  2. from loguru import logger
  3. from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
  4. from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
  5. from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
  6. from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
  7. from magic_pdf.libs.commons import join_path
  8. from magic_pdf.pipe.AbsPipe import AbsPipe
  9. from magic_pdf.user_api import parse_union_pdf, parse_ocr_pdf
  10. class UNIPipe(AbsPipe):
  11. def __init__(self, pdf_bytes: bytes, jso_useful_key: dict, image_writer: AbsReaderWriter, is_debug: bool = False,
  12. <<<<<<< HEAD
  13. start_page_id=0, end_page_id=None, lang=None):
  14. self.pdf_type = jso_useful_key["_pdf_type"]
  15. super().__init__(pdf_bytes, jso_useful_key["model_list"], image_writer, is_debug, start_page_id, end_page_id, lang)
  16. =======
  17. start_page_id=0, end_page_id=None):
  18. self.pdf_type = jso_useful_key["_pdf_type"]
  19. super().__init__(pdf_bytes, jso_useful_key["model_list"], image_writer, is_debug, start_page_id, end_page_id)
  20. >>>>>>> 0140d7d271ac3b1561ca2272030e9e038b469999
  21. if len(self.model_list) == 0:
  22. self.input_model_is_empty = True
  23. else:
  24. self.input_model_is_empty = False
  25. def pipe_classify(self):
  26. self.pdf_type = AbsPipe.classify(self.pdf_bytes)
  27. def pipe_analyze(self):
  28. if self.pdf_type == self.PIP_TXT:
  29. self.model_list = doc_analyze(self.pdf_bytes, ocr=False,
  30. <<<<<<< HEAD
  31. start_page_id=self.start_page_id, end_page_id=self.end_page_id,
  32. lang=self.lang)
  33. elif self.pdf_type == self.PIP_OCR:
  34. self.model_list = doc_analyze(self.pdf_bytes, ocr=True,
  35. start_page_id=self.start_page_id, end_page_id=self.end_page_id,
  36. lang=self.lang)
  37. =======
  38. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  39. elif self.pdf_type == self.PIP_OCR:
  40. self.model_list = doc_analyze(self.pdf_bytes, ocr=True,
  41. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  42. >>>>>>> 0140d7d271ac3b1561ca2272030e9e038b469999
  43. def pipe_parse(self):
  44. if self.pdf_type == self.PIP_TXT:
  45. self.pdf_mid_data = parse_union_pdf(self.pdf_bytes, self.model_list, self.image_writer,
  46. is_debug=self.is_debug, input_model_is_empty=self.input_model_is_empty,
  47. <<<<<<< HEAD
  48. start_page_id=self.start_page_id, end_page_id=self.end_page_id,
  49. lang=self.lang)
  50. =======
  51. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  52. >>>>>>> 0140d7d271ac3b1561ca2272030e9e038b469999
  53. elif self.pdf_type == self.PIP_OCR:
  54. self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer,
  55. is_debug=self.is_debug,
  56. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  57. def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
  58. result = super().pipe_mk_uni_format(img_parent_path, drop_mode)
  59. logger.info("uni_pipe mk content list finished")
  60. return result
  61. def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD):
  62. result = super().pipe_mk_markdown(img_parent_path, drop_mode, md_make_mode)
  63. logger.info(f"uni_pipe mk {md_make_mode} finished")
  64. return result
  65. if __name__ == '__main__':
  66. # 测试
  67. drw = DiskReaderWriter(r"D:/project/20231108code-clean")
  68. pdf_file_path = r"linshixuqiu\19983-00.pdf"
  69. model_file_path = r"linshixuqiu\19983-00.json"
  70. pdf_bytes = drw.read(pdf_file_path, AbsReaderWriter.MODE_BIN)
  71. model_json_txt = drw.read(model_file_path, AbsReaderWriter.MODE_TXT)
  72. model_list = json.loads(model_json_txt)
  73. write_path = r"D:\project\20231108code-clean\linshixuqiu\19983-00"
  74. img_bucket_path = "imgs"
  75. img_writer = DiskReaderWriter(join_path(write_path, img_bucket_path))
  76. # pdf_type = UNIPipe.classify(pdf_bytes)
  77. # jso_useful_key = {
  78. # "_pdf_type": pdf_type,
  79. # "model_list": model_list
  80. # }
  81. jso_useful_key = {
  82. "_pdf_type": "",
  83. "model_list": model_list
  84. }
  85. pipe = UNIPipe(pdf_bytes, jso_useful_key, img_writer)
  86. pipe.pipe_classify()
  87. pipe.pipe_parse()
  88. md_content = pipe.pipe_mk_markdown(img_bucket_path)
  89. content_list = pipe.pipe_mk_uni_format(img_bucket_path)
  90. md_writer = DiskReaderWriter(write_path)
  91. md_writer.write(md_content, "19983-00.md", AbsReaderWriter.MODE_TXT)
  92. md_writer.write(json.dumps(pipe.pdf_mid_data, ensure_ascii=False, indent=4), "19983-00.json",
  93. AbsReaderWriter.MODE_TXT)
  94. md_writer.write(str(content_list), "19983-00.txt", AbsReaderWriter.MODE_TXT)