common.py 4.79 KB
Newer Older
icecraft's avatar
icecraft committed
1
import copy
2
3
4
import json as json_parse
import os

icecraft's avatar
icecraft committed
5
6
import click
from loguru import logger
7
8
9
10

import magic_pdf.model as model_config
from magic_pdf.libs.draw_bbox import (draw_layout_bbox, draw_span_bbox,
                                      drow_model_bbox)
icecraft's avatar
icecraft committed
11
12
13
from magic_pdf.libs.MakeContentConfig import DropMode, MakeMode
from magic_pdf.pipe.OCRPipe import OCRPipe
from magic_pdf.pipe.TXTPipe import TXTPipe
14
from magic_pdf.pipe.UNIPipe import UNIPipe
icecraft's avatar
icecraft committed
15
from magic_pdf.rw.AbsReaderWriter import AbsReaderWriter
16
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
icecraft's avatar
icecraft committed
17
18
19
20
21


def prepare_env(output_dir, pdf_file_name, method):
    local_parent_dir = os.path.join(output_dir, pdf_file_name, method)

22
    local_image_dir = os.path.join(str(local_parent_dir), 'images')
icecraft's avatar
icecraft committed
23
24
25
26
27
28
29
30
31
32
33
34
    local_md_dir = local_parent_dir
    os.makedirs(local_image_dir, exist_ok=True)
    os.makedirs(local_md_dir, exist_ok=True)
    return local_image_dir, local_md_dir


def do_parse(
    output_dir,
    pdf_file_name,
    pdf_bytes,
    model_list,
    parse_method,
35
    debug_able,
icecraft's avatar
icecraft committed
36
37
38
39
40
41
42
43
    f_draw_span_bbox=True,
    f_draw_layout_bbox=True,
    f_dump_md=True,
    f_dump_middle_json=True,
    f_dump_model_json=True,
    f_dump_orig_pdf=True,
    f_dump_content_list=False,
    f_make_md_mode=MakeMode.MM_MD,
44
    f_draw_model_bbox=False,
45
46
    start_page_id=0,
    end_page_id=None,
47
    lang=None,
icecraft's avatar
icecraft committed
48
):
49
    if debug_able:
50
        logger.warning('debug mode is on')
51
52
53
        f_dump_content_list = True
        f_draw_model_bbox = True

icecraft's avatar
icecraft committed
54
    orig_model_list = copy.deepcopy(model_list)
55
56
    local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name,
                                                parse_method)
57

58
59
    image_writer, md_writer = DiskReaderWriter(
        local_image_dir), DiskReaderWriter(local_md_dir)
icecraft's avatar
icecraft committed
60
61
    image_dir = str(os.path.basename(local_image_dir))

62
63
    if parse_method == 'auto':
        jso_useful_key = {'_pdf_type': '', 'model_list': model_list}
64
        pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer, is_debug=True,
65
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang)
66
    elif parse_method == 'txt':
67
        pipe = TXTPipe(pdf_bytes, model_list, image_writer, is_debug=True,
68
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang)
69
    elif parse_method == 'ocr':
70
        pipe = OCRPipe(pdf_bytes, model_list, image_writer, is_debug=True,
71
                       start_page_id=start_page_id, end_page_id=end_page_id, lang=lang)
icecraft's avatar
icecraft committed
72
    else:
73
        logger.error('unknown parse method')
icecraft's avatar
icecraft committed
74
75
76
77
78
79
80
81
82
        exit(1)

    pipe.pipe_classify()

    if len(model_list) == 0:
        if model_config.__use_inside_model__:
            pipe.pipe_analyze()
            orig_model_list = copy.deepcopy(pipe.model_list)
        else:
83
            logger.error('need model list input')
icecraft's avatar
icecraft committed
84
85
86
            exit(2)

    pipe.pipe_parse()
87
    pdf_info = pipe.pdf_mid_data['pdf_info']
icecraft's avatar
icecraft committed
88
    if f_draw_layout_bbox:
89
        draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
icecraft's avatar
icecraft committed
90
    if f_draw_span_bbox:
91
        draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, pdf_file_name)
92
    if f_draw_model_bbox:
93
        drow_model_bbox(copy.deepcopy(orig_model_list), pdf_bytes, local_md_dir, pdf_file_name)
icecraft's avatar
icecraft committed
94

95
96
97
    md_content = pipe.pipe_mk_markdown(image_dir,
                                       drop_mode=DropMode.NONE,
                                       md_make_mode=f_make_md_mode)
icecraft's avatar
icecraft committed
98
99
100
    if f_dump_md:
        md_writer.write(
            content=md_content,
101
            path=f'{pdf_file_name}.md',
icecraft's avatar
icecraft committed
102
103
104
105
106
            mode=AbsReaderWriter.MODE_TXT,
        )

    if f_dump_middle_json:
        md_writer.write(
107
108
109
110
            content=json_parse.dumps(pipe.pdf_mid_data,
                                     ensure_ascii=False,
                                     indent=4),
            path=f'{pdf_file_name}_middle.json',
icecraft's avatar
icecraft committed
111
112
113
114
115
            mode=AbsReaderWriter.MODE_TXT,
        )

    if f_dump_model_json:
        md_writer.write(
116
117
118
119
            content=json_parse.dumps(orig_model_list,
                                     ensure_ascii=False,
                                     indent=4),
            path=f'{pdf_file_name}_model.json',
icecraft's avatar
icecraft committed
120
121
122
123
124
125
            mode=AbsReaderWriter.MODE_TXT,
        )

    if f_dump_orig_pdf:
        md_writer.write(
            content=pdf_bytes,
126
            path=f'{pdf_file_name}_origin.pdf',
icecraft's avatar
icecraft committed
127
128
129
130
131
132
            mode=AbsReaderWriter.MODE_BIN,
        )

    content_list = pipe.pipe_mk_uni_format(image_dir, drop_mode=DropMode.NONE)
    if f_dump_content_list:
        md_writer.write(
133
134
135
136
            content=json_parse.dumps(content_list,
                                     ensure_ascii=False,
                                     indent=4),
            path=f'{pdf_file_name}_content_list.json',
icecraft's avatar
icecraft committed
137
138
139
            mode=AbsReaderWriter.MODE_TXT,
        )

140
    logger.info(f'local output dir is {local_md_dir}')
141

icecraft's avatar
icecraft committed
142

143
parse_pdf_methods = click.Choice(['ocr', 'txt', 'auto'])