app.py 8.51 KB
Newer Older
1
2
3
4
# Copyright (c) Opendatalab. All rights reserved.

import base64
import os
5
import re
6
import time
7
import uuid
8
9
10
import zipfile
from pathlib import Path

11
import gradio as gr
12
import pymupdf
13
from gradio_pdf import PDF
14
15
from loguru import logger

16
from magic_pdf.data.data_reader_writer import FileBasedDataReader
17
18
19
20
21
from magic_pdf.libs.hash_utils import compute_sha256
from magic_pdf.tools.common import do_parse, prepare_env


def read_fn(path):
22
    disk_rw = FileBasedDataReader(os.path.dirname(path))
23
    return disk_rw.read(os.path.basename(path))
24
25


26
def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_enable, table_enable, language):
27
28
29
    os.makedirs(output_dir, exist_ok=True)

    try:
30
        file_name = f'{str(Path(doc_path).stem)}_{time.time()}'
31
32
        pdf_data = read_fn(doc_path)
        if is_ocr:
33
            parse_method = 'ocr'
34
        else:
35
            parse_method = 'auto'
36
37
38
39
40
41
42
43
44
        local_image_dir, local_md_dir = prepare_env(output_dir, file_name, parse_method)
        do_parse(
            output_dir,
            file_name,
            pdf_data,
            [],
            parse_method,
            False,
            end_page_id=end_page_id,
45
46
47
48
            layout_model=layout_mode,
            formula_enable=formula_enable,
            table_enable=table_enable,
            lang=language,
49
50
51
52
53
54
55
        )
        return local_md_dir, file_name
    except Exception as e:
        logger.exception(e)


def compress_directory_to_zip(directory_path, output_zip_path):
56
    """压缩指定目录到一个 ZIP 文件。
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

    :param directory_path: 要压缩的目录路径
    :param output_zip_path: 输出的 ZIP 文件路径
    """
    try:
        with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:

            # 遍历目录中的所有文件和子目录
            for root, dirs, files in os.walk(directory_path):
                for file in files:
                    # 构建完整的文件路径
                    file_path = os.path.join(root, file)
                    # 计算相对路径
                    arcname = os.path.relpath(file_path, directory_path)
                    # 添加文件到 ZIP 文件
                    zipf.write(file_path, arcname)
        return 0
    except Exception as e:
        logger.exception(e)
        return -1


def image_to_base64(image_path):
80
    with open(image_path, 'rb') as image_file:
81
82
83
84
85
86
87
88
89
90
91
92
        return base64.b64encode(image_file.read()).decode('utf-8')


def replace_image_with_base64(markdown_text, image_dir_path):
    # 匹配Markdown中的图片标签
    pattern = r'\!\[(?:[^\]]*)\]\(([^)]+)\)'

    # 替换图片链接
    def replace(match):
        relative_path = match.group(1)
        full_path = os.path.join(image_dir_path, relative_path)
        base64_image = image_to_base64(full_path)
93
        return f'![{relative_path}](data:image/jpeg;base64,{base64_image})'
94
95
96
97
98

    # 应用替换
    return re.sub(pattern, replace, markdown_text)


99
def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
100
    file_path = to_pdf(file_path)
101
    # 获取识别的md文件以及压缩包文件路径
102
103
    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
                                        layout_mode, formula_enable, table_enable, language)
104
    archive_zip_path = os.path.join('./output', compute_sha256(local_md_dir) + '.zip')
105
106
    zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
    if zip_archive_success == 0:
107
        logger.info('压缩成功')
108
    else:
109
110
        logger.error('压缩失败')
    md_path = os.path.join(local_md_dir, file_name + '.md')
111
112
113
114
    with open(md_path, 'r', encoding='utf-8') as f:
        txt_content = f.read()
    md_content = replace_image_with_base64(txt_content, local_md_dir)
    # 返回转换后的PDF路径
115
    new_pdf_path = os.path.join(local_md_dir, file_name + '_layout.pdf')
116
117
118
119

    return md_content, txt_content, archive_zip_path, new_pdf_path


120
121
latex_delimiters = [{'left': '$$', 'right': '$$', 'display': True},
                    {'left': '$', 'right': '$', 'display': False}]
122
123
124
125
126
127


def init_model():
    from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
    try:
        model_manager = ModelSingleton()
128
129
130
131
        txt_model = model_manager.get_model(False, False)  # noqa: F841
        logger.info('txt_model init final')
        ocr_model = model_manager.get_model(True, False)  # noqa: F841
        logger.info('ocr_model init final')
132
133
134
135
136
137
138
        return 0
    except Exception as e:
        logger.exception(e)
        return -1


model_init = init_model()
139
logger.info(f'model_init: {model_init}')
140
141


142
with open('header.html', 'r') as file:
143
144
145
    header = file.read()


146
latin_lang = [
147
        'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr',  # noqa: E126
148
149
150
151
152
153
        'hu', 'id', 'is', 'it', 'ku', 'la', 'lt', 'lv', 'mi', 'ms', 'mt', 'nl',
        'no', 'oc', 'pi', 'pl', 'pt', 'ro', 'rs_latin', 'sk', 'sl', 'sq', 'sv',
        'sw', 'tl', 'tr', 'uz', 'vi', 'french', 'german'
]
arabic_lang = ['ar', 'fa', 'ug', 'ur']
cyrillic_lang = [
154
        'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',  # noqa: E126
155
156
157
        'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
]
devanagari_lang = [
158
        'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',  # noqa: E126
159
160
161
162
        'sa', 'bgc'
]
other_lang = ['ch', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']

163
all_lang = ['', 'auto']
164
165
166
all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])


167
168
169
170
171
172
173
174
def to_pdf(file_path):
    with pymupdf.open(file_path) as f:
        if f.is_pdf:
            return file_path
        else:
            pdf_bytes = f.convert_to_pdf()
            # 将pdfbytes 写入到uuid.pdf中
            # 生成唯一的文件名
175
            unique_filename = f'{uuid.uuid4()}.pdf'
176
177
178
179
180
181
182
183
184
185

            # 构建完整的文件路径
            tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)

            # 将字节数据写入文件
            with open(tmp_file_path, 'wb') as tmp_pdf_file:
                tmp_pdf_file.write(pdf_bytes)

            return tmp_file_path

186

187
if __name__ == '__main__':
188
189
190
191
    with gr.Blocks() as demo:
        gr.HTML(header)
        with gr.Row():
            with gr.Column(variant='panel', scale=5):
192
                file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'])
193
                max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
194
                with gr.Row():
195
                    layout_mode = gr.Dropdown(['layoutlmv3', 'doclayout_yolo'], label='Layout model', value='doclayout_yolo')
196
                    language = gr.Dropdown(all_lang, label='Language', value='')
197
                with gr.Row():
198
199
                    formula_enable = gr.Checkbox(label='Enable formula recognition', value=True)
                    is_ocr = gr.Checkbox(label='Force enable OCR', value=False)
200
                    table_enable = gr.Checkbox(label='Enable table recognition(test)', value=True)
201
                with gr.Row():
202
203
                    change_bu = gr.Button('Convert')
                    clear_bu = gr.ClearButton(value='Clear')
204
                pdf_show = PDF(label='PDF preview', interactive=False, visible=True, height=800)
205
206
                with gr.Accordion('Examples:'):
                    example_root = os.path.join(os.path.dirname(__file__), 'examples')
207
208
                    gr.Examples(
                        examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
209
                                  _.endswith('pdf')],
210
                        inputs=file
211
                    )
212
213

            with gr.Column(variant='panel', scale=5):
214
                output_file = gr.File(label='convert result', interactive=False)
215
                with gr.Tabs():
216
                    with gr.Tab('Markdown rendering'):
217
                        md = gr.Markdown(label='Markdown rendering', height=1100, show_copy_button=True,
218
                                         latex_delimiters=latex_delimiters, line_breaks=True)
219
                    with gr.Tab('Markdown text'):
220
                        md_text = gr.TextArea(lines=45, show_copy_button=True)
221
222
        file.change(fn=to_pdf, inputs=file, outputs=pdf_show)
        change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
223
                        outputs=[md, md_text, output_file, pdf_show])
224
        clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr, table_enable, language])
225

226
    demo.launch(server_name='0.0.0.0')