MinerU/magic_pdf/pipe/AbsPipe.py

from abc import ABC, abstractmethod

from magic_pdf.dict2md.mkcontent import mk_universal_format, mk_mm_markdown
from magic_pdf.dict2md.ocr_mkcontent import make_standard_format_with_para, ocr_mk_mm_markdown_with_para
from magic_pdf.filter.pdf_classify_by_type import classify
from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
from magic_pdf.libs.drop_reason import DropReason
from magic_pdf.libs.json_compressor import JsonCompressor


class AbsPipe(ABC):
    """
    txt和ocr处理的抽象类
    """
    PIP_OCR = "ocr"
    PIP_TXT = "txt"

    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: AbsReaderWriter, img_parent_path:str, is_debug:bool=False):
        self.pdf_bytes = pdf_bytes
        self.model_list = model_list
        self.image_writer = image_writer
        self.img_parent_path = img_parent_path
        self.pdf_mid_data = None # 未压缩
        self.is_debug = is_debug
    
    def get_compress_pdf_mid_data(self):
        return JsonCompressor.compress_json(self.pdf_mid_data)

    @abstractmethod
    def pipe_classify(self):
        """
        有状态的分类
        """
        raise NotImplementedError

    @abstractmethod
    def pipe_parse(self):
        """
        有状态的解析
        """
        raise NotImplementedError

    @abstractmethod
    def pipe_mk_uni_format(self):
        """
        有状态的组装统一格式
        """
        raise NotImplementedError

    @abstractmethod
    def pipe_mk_markdown(self):
        """
        有状态的组装markdown
        """
        raise NotImplementedError

    @staticmethod
    def classify(pdf_bytes: bytes) -> str:
        """
        根据pdf的元数据，判断是否是文本pdf，还是ocr pdf
        """
        pdf_meta = pdf_meta_scan(pdf_bytes)
        if pdf_meta.get("_need_drop", False):  # 如果返回了需要丢弃的标志，则抛出异常
            raise Exception(f"pdf meta_scan need_drop,reason is {pdf_meta['_drop_reason']}")
        else:
            is_encrypted = pdf_meta["is_encrypted"]
            is_needs_password = pdf_meta["is_needs_password"]
            if is_encrypted or is_needs_password:  # 加密的，需要密码的，没有页面的，都不处理
                raise Exception(f"pdf meta_scan need_drop,reason is {DropReason.ENCRYPTED}")
            else:
                is_text_pdf, results = classify(
                    pdf_meta["total_page"],
                    pdf_meta["page_width_pts"],
                    pdf_meta["page_height_pts"],
                    pdf_meta["image_info_per_page"],
                    pdf_meta["text_len_per_page"],
                    pdf_meta["imgs_per_page"],
                    pdf_meta["text_layout_per_page"],
                )
                if is_text_pdf:
                    return AbsPipe.PIP_TXT
                else:
                    return AbsPipe.PIP_OCR

    @staticmethod
    def mk_uni_format(compressed_pdf_mid_data: str, img_buket_path: str) -> list:
        """
        根据pdf类型，生成统一格式content_list
        """
        pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data)
        parse_type = pdf_mid_data["_parse_type"]
        pdf_info_list = pdf_mid_data["pdf_info"]
        if parse_type == AbsPipe.PIP_TXT:
            content_list = mk_universal_format(pdf_info_list, img_buket_path)
        elif parse_type == AbsPipe.PIP_OCR:
            content_list = make_standard_format_with_para(pdf_info_list, img_buket_path)
        return content_list

    @staticmethod
    def mk_markdown(compressed_pdf_mid_data: str, img_buket_path: str) -> list:
        """
        根据pdf类型，markdown
        """
        pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data)
        parse_type = pdf_mid_data["_parse_type"]
        pdf_info_list = pdf_mid_data["pdf_info"]
        if parse_type == AbsPipe.PIP_TXT:
            # content_list = mk_universal_format(pdf_info_list, img_buket_path)
            # md_content = mk_mm_markdown(content_list)
            md_content = ocr_mk_mm_markdown_with_para(pdf_info_list, img_buket_path)
        elif parse_type == AbsPipe.PIP_OCR:
            md_content = ocr_mk_mm_markdown_with_para(pdf_info_list, img_buket_path)
        return md_content