OCRPipe.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. from loguru import logger
  2. from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
  3. from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
  4. from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
  5. from magic_pdf.pipe.AbsPipe import AbsPipe
  6. from magic_pdf.user_api import parse_ocr_pdf
  7. class OCRPipe(AbsPipe):
  8. def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, is_debug: bool = False,
  9. <<<<<<< HEAD
  10. start_page_id=0, end_page_id=None, lang=None):
  11. super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id, lang)
  12. =======
  13. start_page_id=0, end_page_id=None):
  14. super().__init__(pdf_bytes, model_list, image_writer, is_debug, start_page_id, end_page_id)
  15. >>>>>>> 0140d7d271ac3b1561ca2272030e9e038b469999
  16. def pipe_classify(self):
  17. pass
  18. def pipe_analyze(self):
  19. self.model_list = doc_analyze(self.pdf_bytes, ocr=True,
  20. <<<<<<< HEAD
  21. start_page_id=self.start_page_id, end_page_id=self.end_page_id,
  22. lang=self.lang)
  23. =======
  24. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  25. >>>>>>> 0140d7d271ac3b1561ca2272030e9e038b469999
  26. def pipe_parse(self):
  27. self.pdf_mid_data = parse_ocr_pdf(self.pdf_bytes, self.model_list, self.image_writer, is_debug=self.is_debug,
  28. start_page_id=self.start_page_id, end_page_id=self.end_page_id)
  29. def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
  30. result = super().pipe_mk_uni_format(img_parent_path, drop_mode)
  31. logger.info("ocr_pipe mk content list finished")
  32. return result
  33. def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD):
  34. result = super().pipe_mk_markdown(img_parent_path, drop_mode, md_make_mode)
  35. logger.info(f"ocr_pipe mk {md_make_mode} finished")
  36. return result