AbsPipe.py 4.27 KB
Newer Older
赵小蒙's avatar
赵小蒙 committed
1
2
from abc import ABC, abstractmethod

3
4
from magic_pdf.config.drop_reason import DropReason
from magic_pdf.config.make_content_config import DropMode, MakeMode
5
from magic_pdf.data.data_reader_writer import DataWriter
赵小蒙's avatar
赵小蒙 committed
6
from magic_pdf.dict2md.ocr_mkcontent import union_make
赵小蒙's avatar
赵小蒙 committed
7
8
9
10
11
12
from magic_pdf.filter.pdf_classify_by_type import classify
from magic_pdf.filter.pdf_meta_scan import pdf_meta_scan
from magic_pdf.libs.json_compressor import JsonCompressor


class AbsPipe(ABC):
13
14
15
    """txt和ocr处理的抽象类."""
    PIP_OCR = 'ocr'
    PIP_TXT = 'txt'
赵小蒙's avatar
赵小蒙 committed
16

17
    def __init__(self, pdf_bytes: bytes, model_list: list, image_writer: DataWriter, is_debug: bool = False,
18
                 start_page_id=0, end_page_id=None, lang=None, layout_model=None, formula_enable=None, table_enable=None):
赵小蒙's avatar
赵小蒙 committed
19
20
21
        self.pdf_bytes = pdf_bytes
        self.model_list = model_list
        self.image_writer = image_writer
22
        self.pdf_mid_data = None  # 未压缩
kernel.h@qq.com's avatar
kernel.h@qq.com committed
23
        self.is_debug = is_debug
24
25
        self.start_page_id = start_page_id
        self.end_page_id = end_page_id
26
        self.lang = lang
27
28
29
        self.layout_model = layout_model
        self.formula_enable = formula_enable
        self.table_enable = table_enable
30

kernel.h@qq.com's avatar
kernel.h@qq.com committed
31
32
    def get_compress_pdf_mid_data(self):
        return JsonCompressor.compress_json(self.pdf_mid_data)
赵小蒙's avatar
赵小蒙 committed
33
34
35

    @abstractmethod
    def pipe_classify(self):
36
        """有状态的分类."""
赵小蒙's avatar
赵小蒙 committed
37
38
        raise NotImplementedError

39
40
    @abstractmethod
    def pipe_analyze(self):
41
        """有状态的跑模型分析."""
42
43
        raise NotImplementedError

赵小蒙's avatar
赵小蒙 committed
44
45
    @abstractmethod
    def pipe_parse(self):
46
        """有状态的解析."""
赵小蒙's avatar
赵小蒙 committed
47
48
        raise NotImplementedError

49
50
51
    def pipe_mk_uni_format(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF):
        content_list = AbsPipe.mk_uni_format(self.get_compress_pdf_mid_data(), img_parent_path, drop_mode)
        return content_list
赵小蒙's avatar
赵小蒙 committed
52

53
54
55
    def pipe_mk_markdown(self, img_parent_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD):
        md_content = AbsPipe.mk_markdown(self.get_compress_pdf_mid_data(), img_parent_path, drop_mode, md_make_mode)
        return md_content
赵小蒙's avatar
赵小蒙 committed
56
57
58

    @staticmethod
    def classify(pdf_bytes: bytes) -> str:
59
        """根据pdf的元数据,判断是文本pdf,还是ocr pdf."""
赵小蒙's avatar
赵小蒙 committed
60
        pdf_meta = pdf_meta_scan(pdf_bytes)
61
        if pdf_meta.get('_need_drop', False):  # 如果返回了需要丢弃的标志,则抛出异常
赵小蒙's avatar
赵小蒙 committed
62
63
            raise Exception(f"pdf meta_scan need_drop,reason is {pdf_meta['_drop_reason']}")
        else:
64
65
            is_encrypted = pdf_meta['is_encrypted']
            is_needs_password = pdf_meta['is_needs_password']
赵小蒙's avatar
赵小蒙 committed
66
            if is_encrypted or is_needs_password:  # 加密的,需要密码的,没有页面的,都不处理
67
                raise Exception(f'pdf meta_scan need_drop,reason is {DropReason.ENCRYPTED}')
赵小蒙's avatar
赵小蒙 committed
68
69
            else:
                is_text_pdf, results = classify(
70
71
72
73
74
75
76
77
                    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'],
                    pdf_meta['invalid_chars'],
赵小蒙's avatar
赵小蒙 committed
78
79
                )
                if is_text_pdf:
kernel.h@qq.com's avatar
kernel.h@qq.com committed
80
                    return AbsPipe.PIP_TXT
赵小蒙's avatar
赵小蒙 committed
81
                else:
kernel.h@qq.com's avatar
kernel.h@qq.com committed
82
                    return AbsPipe.PIP_OCR
赵小蒙's avatar
赵小蒙 committed
83
84

    @staticmethod
85
    def mk_uni_format(compressed_pdf_mid_data: str, img_buket_path: str, drop_mode=DropMode.WHOLE_PDF) -> list:
86
        """根据pdf类型,生成统一格式content_list."""
赵小蒙's avatar
赵小蒙 committed
87
        pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data)
88
        pdf_info_list = pdf_mid_data['pdf_info']
89
        content_list = union_make(pdf_info_list, MakeMode.STANDARD_FORMAT, drop_mode, img_buket_path)
赵小蒙's avatar
赵小蒙 committed
90
91
92
        return content_list

    @staticmethod
93
    def mk_markdown(compressed_pdf_mid_data: str, img_buket_path: str, drop_mode=DropMode.WHOLE_PDF, md_make_mode=MakeMode.MM_MD) -> list:
94
        """根据pdf类型,markdown."""
赵小蒙's avatar
赵小蒙 committed
95
        pdf_mid_data = JsonCompressor.decompress_json(compressed_pdf_mid_data)
96
        pdf_info_list = pdf_mid_data['pdf_info']
97
        md_content = union_make(pdf_info_list, md_make_mode, drop_mode, img_buket_path)
赵小蒙's avatar
赵小蒙 committed
98
        return md_content