Update app.py

app.py

@@ -1,28 +1,22 @@
 # Copyright (c) Opendatalab. All rights reserved.
 
 import base64
-import json
 import os
+import json
+import re
 import time
 import zipfile
 from pathlib import Path
-import re
-import uuid
-import pymupdf
 
-# os.system('pip install -U magic-pdf==0.10.5')
 os.system('pip uninstall -y magic-pdf')
+os.system('pip uninstall -y mineru')
 os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
-# os.system('pip install git+https://github.com/myhloli/Magic-PDF.git@dev')
 
-os.system('wget https://github.com/opendatalab/MinerU/raw/dev/scripts/download_models_hf.py -O download_models_hf.py')
-os.system('python download_models_hf.py')
+os.system('mineru-models-download -s huggingface -m pipeline')
 
-with open('/home/user/magic-pdf.json', 'r') as file:
+with open('/home/user/mineru.json', 'r') as file:
     config = json.load(file)
 
-config['device-mode'] = "cuda"
-
 delimiters = {
     'display': {'left': '\\[', 'right': '\\]'},
     'inline': {'left': '\\(', 'right': '\\)'}
@@ -34,49 +28,39 @@ if os.getenv('apikey'):
     config['llm-aided-config']['title_aided']['api_key'] = os.getenv('apikey')
     config['llm-aided-config']['title_aided']['enable'] = True
 
-with open('/home/user/magic-pdf.json', 'w') as file:
+with open('/home/user/mineru.json', 'w') as file:
     json.dump(config, file, indent=4)
 
-# os.system('cp -r paddleocr /home/user/.paddleocr')
 from gradio_pdf import PDF
-
 import gradio as gr
 from loguru import logger
 
-from magic_pdf.data.data_reader_writer import FileBasedDataReader
-from magic_pdf.libs.hash_utils import compute_sha256
-from magic_pdf.tools.common import do_parse, prepare_env
+from mineru.cli.common import prepare_env, do_parse, read_fn
+from mineru.utils.hash_utils import str_sha256
 
+os.environ['MINERU_MODEL_SOURCE'] = 'local'
 
-def read_fn(path):
-    disk_rw = FileBasedDataReader(os.path.dirname(path))
-    return disk_rw.read(os.path.basename(path))
 
-
-def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_enable, table_enable, language):
+def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, formula_enable, table_enable, language):
     os.makedirs(output_dir, exist_ok=True)
 
     try:
-        file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
+        file_name = f'{str(Path(doc_path).stem)}_{time.strftime("%y%m%d_%H%M%S")}'
         pdf_data = read_fn(doc_path)
         if is_ocr:
-            parse_method = "ocr"
+            parse_method = 'ocr'
         else:
-            parse_method = "auto"
+            parse_method = 'auto'
         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,
+            output_dir=output_dir,
+            pdf_file_names=[file_name],
+            pdf_bytes_list=[pdf_data],
+            p_lang_list=[language],
+            parse_method=parse_method,
             end_page_id=end_page_id,
-            layout_model=layout_mode,
-            formula_enable=formula_enable,
-            table_enable=table_enable,
-            lang=language,
-            f_dump_orig_pdf=False,
+            p_formula_enable=formula_enable,
+            p_table_enable=table_enable,
         )
         return local_md_dir, file_name
     except Exception as e:
@@ -128,25 +112,24 @@ def replace_image_with_base64(markdown_text, image_dir_path):
     return re.sub(pattern, replace, markdown_text)
 
 
-def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
+def to_markdown(file_path, end_pages, is_ocr, formula_enable, table_enable, language):
     file_path = to_pdf(file_path)
     if end_pages > 20:
         end_pages = 20
     # 获取识别的md文件以及压缩包文件路径
-    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr,
-                                        layout_mode, formula_enable, table_enable, language)
-    archive_zip_path = os.path.join("./output", compute_sha256(local_md_dir) + ".zip")
+    local_md_dir, file_name = parse_pdf(file_path, './output', end_pages - 1, is_ocr, formula_enable, table_enable, language)
+    archive_zip_path = os.path.join('./output', str_sha256(local_md_dir) + '.zip')
     zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
     if zip_archive_success == 0:
-        logger.info("压缩成功")
+        logger.info('压缩成功')
     else:
-        logger.error("压缩失败")
-    md_path = os.path.join(local_md_dir, file_name + ".md")
+        logger.error('压缩失败')
+    md_path = os.path.join(local_md_dir, file_name + '.md')
     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路径
-    new_pdf_path = os.path.join(local_md_dir, file_name + "_layout.pdf")
+    new_pdf_path = os.path.join(local_md_dir, file_name + '_layout.pdf')
 
     return md_content, txt_content, archive_zip_path, new_pdf_path
 
@@ -155,24 +138,6 @@ latex_delimiters = [{"left": "\\[", "right": "\\]", "display": True},
                     {"left": "\\(", "right": "\\)", "display": False}]
 
 
-def init_model():
-    from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
-    try:
-        model_manager = ModelSingleton()
-        txt_model = model_manager.get_model(False, False)
-        logger.info(f"txt_model init final")
-        ocr_model = model_manager.get_model(True, False)
-        logger.info(f"ocr_model init final")
-        return 0
-    except Exception as e:
-        logger.exception(e)
-        return -1
-
-
-model_init = init_model()
-logger.info(f"model_init: {model_init}")
-
-
 with open("header.html", "r") as file:
     header = file.read()
 
@@ -201,67 +166,72 @@ all_lang = []
 all_lang.extend([*other_lang, *add_lang])
 
 
+def safe_stem(file_path):
+    stem = Path(file_path).stem
+    # 只保留字母、数字、下划线和点，其他字符替换为下划线
+    return re.sub(r'[^\w.]', '_', stem)
+
+
 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中
-            # 生成唯一的文件名
-            unique_filename = f"{uuid.uuid4()}.pdf"
 
-            # 构建完整的文件路径
-            tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
+    if file_path is None:
+        return None
+
+    pdf_bytes = read_fn(file_path)
+
+    # unique_filename = f'{uuid.uuid4()}.pdf'
+    unique_filename = f'{safe_stem(file_path)}.pdf'
 
-            # 将字节数据写入文件
-            with open(tmp_file_path, 'wb') as tmp_pdf_file:
-                tmp_pdf_file.write(pdf_bytes)
+    # 构建完整的文件路径
+    tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
 
-            return tmp_file_path
+    # 将字节数据写入文件
+    with open(tmp_file_path, 'wb') as tmp_pdf_file:
+        tmp_pdf_file.write(pdf_bytes)
 
+    return tmp_file_path
+    
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     with gr.Blocks() as demo:
         gr.HTML(header)
         with gr.Row():
             with gr.Column(variant='panel', scale=5):
-                file = gr.File(label="Please upload a PDF or image", file_types=[".pdf", ".png", ".jpeg", ".jpg"])
-                max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
                 with gr.Row():
-                    layout_mode = gr.Dropdown([
-                        # "layoutlmv3", 
-                        "doclayout_yolo"
-                    ],label="Layout model", value="doclayout_yolo")
-                    language = gr.Dropdown(all_lang, label="Language", value='ch')
+                    file = gr.File(label='Please upload a PDF or image', file_types=['.pdf', '.png', '.jpeg', '.jpg'])
+                with gr.Row(equal_height=True):
+                    with gr.Column(scale=4):
+                        max_pages = gr.Slider(1, 20, 10, step=1, label='Max convert pages')
+                    with gr.Column(scale=1):
+                        language = gr.Dropdown(all_lang, label='Language', value='ch')
                 with gr.Row():
-                    formula_enable = gr.Checkbox(label="Enable formula recognition", value=True)
-                    is_ocr = gr.Checkbox(label="Force enable OCR", value=False)
-                    table_enable = gr.Checkbox(label="Enable table recognition(test)", value=True)
+                    is_ocr = gr.Checkbox(label='Force enable OCR', value=False)
+                    formula_enable = gr.Checkbox(label='Enable formula recognition', value=True)
+                    table_enable = gr.Checkbox(label='Enable table recognition(test)', value=True)
                 with gr.Row():
-                    change_bu = gr.Button("Convert")
-                    clear_bu = gr.ClearButton(value="Clear")
+                    change_bu = gr.Button('Convert')
+                    clear_bu = gr.ClearButton(value='Clear')
                 pdf_show = PDF(label='PDF preview', interactive=False, visible=True, height=800)
-                with gr.Accordion("Examples:"):
-                    example_root = os.path.join(os.path.dirname(__file__), "examples")
+                with gr.Accordion('Examples:'):
+                    example_root = os.path.join(os.path.dirname(__file__), 'examples')
                     gr.Examples(
                         examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
-                                  _.endswith("pdf")],
+                                  _.endswith('pdf')],
                         inputs=file
                     )
 
             with gr.Column(variant='panel', scale=5):
-                output_file = gr.File(label="convert result", interactive=False)
+                output_file = gr.File(label='convert result', interactive=False)
                 with gr.Tabs():
-                    with gr.Tab("Markdown rendering"):
-                        md = gr.Markdown(label="Markdown rendering", height=1100, show_copy_button=True,
+                    with gr.Tab('Markdown rendering'):
+                        md = gr.Markdown(label='Markdown rendering', height=1100, show_copy_button=True,
                                          latex_delimiters=latex_delimiters,
                                          line_breaks=True)
-                    with gr.Tab("Markdown text"):
+                    with gr.Tab('Markdown text'):
                         md_text = gr.TextArea(lines=45, show_copy_button=True)
         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],
-                        outputs=[md, md_text, output_file, pdf_show], api_name=False)
+        change_bu.click(fn=to_markdown, inputs=[file, max_pages, is_ocr, formula_enable, table_enable, language],
+                        outputs=[md, md_text, output_file, pdf_show])
         clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr])
 
     demo.launch(ssr_mode=True)