reorganized structure

app.py

@@ -1,252 +1,78 @@
-# Copyright (c) Opendatalab. All rights reserved.
-
-import base64
-import json
+# app.py
 import os
-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 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')
-
-with open('/home/user/magic-pdf.json', 'r') as file:
-    data = json.load(file)
-
-data['device-mode'] = "cuda"
-if os.getenv('apikey'):
-    data['llm-aided-config']['title_aided']['api_key'] = os.getenv('apikey')
-    data['llm-aided-config']['title_aided']['enable'] = True
-
-with open('/home/user/magic-pdf.json', 'w') as file:
-    json.dump(data, 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
-
-
-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):
-    os.makedirs(output_dir, exist_ok=True)
-
-    try:
-        file_name = f"{str(Path(doc_path).stem)}_{time.time()}"
-        pdf_data = read_fn(doc_path)
-        if is_ocr:
-            parse_method = "ocr"
-        else:
-            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,
-            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,
-        )
-        return local_md_dir, file_name
-    except Exception as e:
-        logger.exception(e)
-
-
-def compress_directory_to_zip(directory_path, output_zip_path):
-    """
-    压缩指定目录到一个 ZIP 文件。
-
-    :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):
-    with open(image_path, "rb") as image_file:
-        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)
-        return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
-
-    # 应用替换
-    return re.sub(pattern, replace, markdown_text)
-
-
-def to_markdown(file_path, end_pages, is_ocr, layout_mode, 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")
-    zip_archive_success = compress_directory_to_zip(local_md_dir, archive_zip_path)
-    if zip_archive_success == 0:
-        logger.info("压缩成功")
-    else:
-        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")
-
-    return md_content, txt_content, archive_zip_path, new_pdf_path
-
-
-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}")
+from gradio_pdf import PDF
+from model import model_initialized
+from pdf_processor import to_pdf, to_markdown
+from config import config
+import logging
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
 
+# Load header HTML content
 with open("header.html", "r") as file:
     header = file.read()
 
-
-latin_lang = [
-        'af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr',
-        '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'
-]
+# Language options (you may also move these to config.yaml if preferred)
+latin_lang = ['af', 'az', 'bs', 'cs', 'cy', 'da', 'de', 'es', 'et', 'fr', 'ga', 'hr', '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 = [
-        'ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava',
-        'dar', 'inh', 'che', 'lbe', 'lez', 'tab'
-]
-devanagari_lang = [
-        'hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom',
-        'sa', 'bgc'
-]
+cyrillic_lang = ['ru', 'rs_cyrillic', 'be', 'bg', 'uk', 'mn', 'abq', 'ady', 'kbd', 'ava', 'dar', 'inh', 'che', 'lbe', 'lez', 'tab']
+devanagari_lang = ['hi', 'mr', 'ne', 'bh', 'mai', 'ang', 'bho', 'mah', 'sck', 'new', 'gom', 'sa', 'bgc']
 other_lang = ['ch', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']
 
-all_lang = ['', 'auto']
-all_lang.extend([*other_lang, *latin_lang, *arabic_lang, *cyrillic_lang, *devanagari_lang])
-
-
-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)
-
-            # 将字节数据写入文件
-            with open(tmp_file_path, 'wb') as tmp_pdf_file:
-                tmp_pdf_file.write(pdf_bytes)
-
-            return tmp_file_path
-
+all_lang = ['', 'auto'] + other_lang + latin_lang + arabic_lang + cyrillic_lang + devanagari_lang
+
+# Utility function to ensure input file is a PDF
+def file_to_pdf(file_obj):
+    if file_obj is not None:
+        return to_pdf(file_obj.name)
+    return None
+
+with gr.Blocks() as demo:
+    gr.HTML(header)
+    with gr.Row():
+        with gr.Column(variant='panel', scale=5):
+            file_input = gr.File(label="Please upload a PDF or image", file_types=[".pdf", ".png", ".jpeg", ".jpg"])
+            max_pages = gr.Slider(1, 20, config.get("max_pages_default", 10), step=1, label='Max convert pages')
+            with gr.Row():
+                layout_mode = gr.Dropdown(
+                    ["layoutlmv3", "doclayout_yolo"],
+                    label="Layout model",
+                    value=config.get("layout_model_default", "doclayout_yolo")
+                )
+                language = gr.Dropdown(
+                    all_lang,
+                    label="Language",
+                    value=config.get("language_default", "auto")
+                )
+            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", value=True)
+            with gr.Row():
+                convert_button = gr.Button("Convert")
+                clear_button = gr.ClearButton(value="Clear")
+            pdf_display = PDF(label='PDF preview', interactive=False, visible=True, height=800)
+            with gr.Accordion("Examples:"):
+                example_root = os.path.join(os.path.dirname(__file__), "examples")
+                examples = [os.path.join(example_root, f) for f in os.listdir(example_root) if f.endswith("pdf")]
+                gr.Examples(examples=examples, inputs=file_input)
+        with gr.Column(variant='panel', scale=5):
+            output_file = gr.File(label="Convert result", interactive=False)
+            with gr.Tabs():
+                with gr.Tab("Markdown rendering"):
+                    md_render = gr.Markdown(label="Markdown rendering", height=1100, show_copy_button=True, line_breaks=True)
+                with gr.Tab("Markdown text"):
+                    md_text = gr.TextArea(lines=45, show_copy_button=True)
+    
+    file_input.change(fn=file_to_pdf, inputs=file_input, outputs=pdf_display)
+    convert_button.click(
+        fn=to_markdown,
+        inputs=[file_input, max_pages, is_ocr, layout_mode, formula_enable, table_enable, language],
+        outputs=[md_render, md_text, output_file, pdf_display]
+    )
+    clear_button.add([file_input, md_render, pdf_display, md_text, output_file, is_ocr])
 
 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='auto')
-                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)
-                with gr.Row():
-                    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")
-                    gr.Examples(
-                        examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
-                                  _.endswith("pdf")],
-                        inputs=file
-                    )
-
-            with gr.Column(variant='panel', scale=5):
-                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,
-                                         latex_delimiters=latex_delimiters, line_breaks=True)
-                    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)
-        clear_bu.add([file, md, pdf_show, md_text, output_file, is_ocr])
-
     demo.launch(ssr_mode=True)

config.py

@@ -0,0 +1,9 @@
+# config.py
+import yaml
+
+def load_config(config_path="config.yaml"):
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+config = load_config()

config.yaml

@@ -0,0 +1,12 @@
+# config.yaml
+
+# Device and processing settings
+device_mode: "cuda"
+magic_pdf_json: "/home/user/magic-pdf.json"
+paddleocr_dir: "/home/user/.paddleocr"
+output_dir: "./output"
+
+# Default UI values
+layout_model_default: "doclayout_yolo"
+language_default: "auto"
+max_pages_default: 10

download_models_hf.py

@@ -0,0 +1,66 @@
+import json
+import os
+
+import requests
+from huggingface_hub import snapshot_download
+
+
+def download_json(url):
+    # 下载JSON文件
+    response = requests.get(url)
+    response.raise_for_status()  # 检查请求是否成功
+    return response.json()
+
+
+def download_and_modify_json(url, local_filename, modifications):
+    if os.path.exists(local_filename):
+        data = json.load(open(local_filename))
+        config_version = data.get('config_version', '0.0.0')
+        if config_version < '1.1.1':
+            data = download_json(url)
+    else:
+        data = download_json(url)
+
+    # 修改内容
+    for key, value in modifications.items():
+        data[key] = value
+
+    # 保存修改后的内容
+    with open(local_filename, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=4)
+
+
+if __name__ == '__main__':
+
+    mineru_patterns = [
+        "models/Layout/LayoutLMv3/*",
+        "models/Layout/YOLO/*",
+        "models/MFD/YOLO/*",
+        "models/MFR/unimernet_small_2501/*",
+        "models/TabRec/TableMaster/*",
+        "models/TabRec/StructEqTable/*",
+    ]
+    model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns)
+
+    layoutreader_pattern = [
+        "*.json",
+        "*.safetensors",
+    ]
+    layoutreader_model_dir = snapshot_download('hantian/layoutreader', allow_patterns=layoutreader_pattern)
+
+    model_dir = model_dir + '/models'
+    print(f'model_dir is: {model_dir}')
+    print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
+
+    json_url = 'https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json'
+    config_file_name = 'magic-pdf.json'
+    home_dir = os.path.expanduser('~')
+    config_file = os.path.join(home_dir, config_file_name)
+
+    json_mods = {
+        'models-dir': model_dir,
+        'layoutreader-model-dir': layoutreader_model_dir,
+    }
+
+    download_and_modify_json(json_url, config_file, json_mods)
+    print(f'The configuration file has been configured successfully, the path is: {config_file}')

model.py

@@ -0,0 +1,20 @@
+# model.py
+import logging
+from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
+
+def init_model():
+    """
+    Initialize the document analysis models.
+    """
+    try:
+        model_manager = ModelSingleton()
+        _ = model_manager.get_model(False, False)
+        logging.info("Text model initialized successfully")
+        _ = model_manager.get_model(True, False)
+        logging.info("OCR model initialized successfully")
+        return True
+    except Exception as e:
+        logging.exception("Model initialization failed: %s", e)
+        return False
+
+model_initialized = init_model()

pdf_processor.py

@@ -0,0 +1,79 @@
+# pdf_processor.py
+import os
+import time
+from pathlib import Path
+import uuid
+import logging
+import pymupdf
+
+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 utils import compress_directory_to_zip, replace_image_with_base64
+
+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):
+    """
+    Parse the PDF using the specified parameters.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    try:
+        file_name = f"{Path(doc_path).stem}_{int(time.time())}"
+        pdf_data = read_fn(doc_path)
+        parse_method = "ocr" if is_ocr else "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,
+            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,
+        )
+        return local_md_dir, file_name
+    except Exception as e:
+        logging.exception("Error in parse_pdf: %s", e)
+        raise
+
+def to_pdf(file_path):
+    """
+    Ensures the file is in PDF format. Converts if necessary.
+    """
+    with pymupdf.open(file_path) as f:
+        if f.is_pdf:
+            return file_path
+        else:
+            pdf_bytes = f.convert_to_pdf()
+            unique_filename = f"{uuid.uuid4()}.pdf"
+            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
+
+def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language, output_dir="./output"):
+    """
+    Converts the PDF to markdown and compresses the result.
+    """
+    file_path = to_pdf(file_path)
+    end_pages = min(end_pages, 20)
+    local_md_dir, file_name = parse_pdf(file_path, output_dir, end_pages - 1, is_ocr, layout_mode, formula_enable, table_enable, language)
+    archive_zip_path = os.path.join(output_dir, compute_sha256(local_md_dir) + ".zip")
+    if compress_directory_to_zip(local_md_dir, archive_zip_path):
+        logging.info("Compression successful")
+    else:
+        logging.error("Compression failed")
+    md_path = os.path.join(local_md_dir, f"{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)
+    new_pdf_path = os.path.join(local_md_dir, f"{file_name}_layout.pdf")
+    return md_content, txt_content, archive_zip_path, new_pdf_path

utils.py

@@ -0,0 +1,40 @@
+# utils.py
+import os
+import zipfile
+import base64
+import re
+import logging
+
+def compress_directory_to_zip(directory_path, output_zip_path):
+    """
+    Compresses the specified directory into a ZIP file.
+    """
+    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)
+                    zipf.write(file_path, arcname)
+        return True
+    except Exception as e:
+        logging.exception("Error compressing directory: %s", e)
+        return False
+
+def image_to_base64(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
+def replace_image_with_base64(markdown_text, image_dir_path):
+    """
+    Finds markdown image tags and replaces file paths with embedded base64 data.
+    """
+    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)
+        return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
+    
+    return re.sub(pattern, replace, markdown_text)