user_api.py 3.94 KB
Newer Older
1
"""用户输入: model数组,每个元素代表一个页面 pdf在s3的路径 截图保存的s3位置.
赵小蒙's avatar
赵小蒙 committed
2
3
4
5
6
7
8

然后:
    1)根据s3路径,调用spark集群的api,拿到ak,sk,endpoint,构造出s3PDFReader
    2)根据用户输入的s3地址,调用spark集群的api,拿到ak,sk,endpoint,构造出s3ImageWriter

其余部分至于构造s3cli, 获取ak,sk都在code-clean里写代码完成。不要反向依赖!!!
"""
9

赵小蒙's avatar
赵小蒙 committed
10
11
from loguru import logger

12
from magic_pdf.data.data_reader_writer import DataWriter
13
from magic_pdf.data.dataset import Dataset
赵小蒙's avatar
赵小蒙 committed
14
from magic_pdf.libs.version import __version__
赵小蒙's avatar
赵小蒙 committed
15
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
赵小蒙's avatar
赵小蒙 committed
16
17
from magic_pdf.pdf_parse_by_ocr import parse_pdf_by_ocr
from magic_pdf.pdf_parse_by_txt import parse_pdf_by_txt
赵小蒙's avatar
赵小蒙 committed
18

19
20
PARSE_TYPE_TXT = 'txt'
PARSE_TYPE_OCR = 'ocr'
kernel.h@qq.com's avatar
kernel.h@qq.com committed
21

22

23
24
25
26
27
28
29
30
31
32
33
def parse_txt_pdf(
    dataset: Dataset,
    model_list: list,
    imageWriter: DataWriter,
    is_debug=False,
    start_page_id=0,
    end_page_id=None,
    lang=None,
    *args,
    **kwargs
):
34
    """解析文本类pdf."""
赵小蒙's avatar
赵小蒙 committed
35
    pdf_info_dict = parse_pdf_by_txt(
36
37
        dataset,
        model_list,
赵小蒙's avatar
赵小蒙 committed
38
        imageWriter,
39
40
        start_page_id=start_page_id,
        end_page_id=end_page_id,
赵小蒙's avatar
赵小蒙 committed
41
        debug_mode=is_debug,
42
        lang=lang,
赵小蒙's avatar
赵小蒙 committed
43
44
    )

45
    pdf_info_dict['_parse_type'] = PARSE_TYPE_TXT
赵小蒙's avatar
赵小蒙 committed
46

47
    pdf_info_dict['_version_name'] = __version__
赵小蒙's avatar
赵小蒙 committed
48

49
    if lang is not None:
50
        pdf_info_dict['_lang'] = lang
51

赵小蒙's avatar
赵小蒙 committed
52
53
54
    return pdf_info_dict


55
56
57
58
59
60
61
62
63
64
65
def parse_ocr_pdf(
    dataset: Dataset,
    model_list: list,
    imageWriter: DataWriter,
    is_debug=False,
    start_page_id=0,
    end_page_id=None,
    lang=None,
    *args,
    **kwargs
):
66
    """解析ocr类pdf."""
赵小蒙's avatar
赵小蒙 committed
67
    pdf_info_dict = parse_pdf_by_ocr(
68
69
        dataset,
        model_list,
赵小蒙's avatar
赵小蒙 committed
70
        imageWriter,
71
72
        start_page_id=start_page_id,
        end_page_id=end_page_id,
赵小蒙's avatar
赵小蒙 committed
73
        debug_mode=is_debug,
74
        lang=lang,
赵小蒙's avatar
赵小蒙 committed
75
76
    )

77
    pdf_info_dict['_parse_type'] = PARSE_TYPE_OCR
赵小蒙's avatar
赵小蒙 committed
78

79
    pdf_info_dict['_version_name'] = __version__
赵小蒙's avatar
赵小蒙 committed
80

81
    if lang is not None:
82
        pdf_info_dict['_lang'] = lang
83

赵小蒙's avatar
赵小蒙 committed
84
85
86
    return pdf_info_dict


87
88
89
90
91
92
93
94
95
96
97
def parse_union_pdf(
    dataset: Dataset,
    model_list: list,
    imageWriter: DataWriter,
    is_debug=False,
    start_page_id=0,
    end_page_id=None,
    lang=None,
    *args,
    **kwargs
):
98
    """ocr和文本混合的pdf,全部解析出来."""
赵小蒙's avatar
赵小蒙 committed
99
100
101
102

    def parse_pdf(method):
        try:
            return method(
103
104
                dataset,
                model_list,
赵小蒙's avatar
赵小蒙 committed
105
                imageWriter,
106
107
                start_page_id=start_page_id,
                end_page_id=end_page_id,
赵小蒙's avatar
赵小蒙 committed
108
                debug_mode=is_debug,
109
                lang=lang,
赵小蒙's avatar
赵小蒙 committed
110
111
            )
        except Exception as e:
112
            logger.exception(e)
赵小蒙's avatar
赵小蒙 committed
113
114
115
            return None

    pdf_info_dict = parse_pdf(parse_pdf_by_txt)
116
117
    if pdf_info_dict is None or pdf_info_dict.get('_need_drop', False):
        logger.warning('parse_pdf_by_txt drop or error, switch to parse_pdf_by_ocr')
118
        if len(model_list) == 0:
119
120
121
            layout_model = kwargs.get('layout_model', None)
            formula_enable = kwargs.get('formula_enable', None)
            table_enable = kwargs.get('table_enable', None)
122
123
            infer_res = doc_analyze(
                dataset,
124
125
126
127
128
129
130
131
                ocr=True,
                start_page_id=start_page_id,
                end_page_id=end_page_id,
                lang=lang,
                layout_model=layout_model,
                formula_enable=formula_enable,
                table_enable=table_enable,
            )
132
            model_list = infer_res.get_infer_res()
赵小蒙's avatar
赵小蒙 committed
133
134
        pdf_info_dict = parse_pdf(parse_pdf_by_ocr)
        if pdf_info_dict is None:
135
            raise Exception('Both parse_pdf_by_txt and parse_pdf_by_ocr failed.')
赵小蒙's avatar
赵小蒙 committed
136
        else:
137
            pdf_info_dict['_parse_type'] = PARSE_TYPE_OCR
赵小蒙's avatar
赵小蒙 committed
138
    else:
139
        pdf_info_dict['_parse_type'] = PARSE_TYPE_TXT
赵小蒙's avatar
赵小蒙 committed
140

141
    pdf_info_dict['_version_name'] = __version__
赵小蒙's avatar
赵小蒙 committed
142

143
    if lang is not None:
144
        pdf_info_dict['_lang'] = lang
145

赵小蒙's avatar
赵小蒙 committed
146
    return pdf_info_dict