Update app.py

app.py

@@ -1,65 +1,194 @@
-import gradio as gr
-import os
-from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
-from magic_pdf.pipe.UNIPipe import UNIPipe
-from magic_pdf.tools.common import parse_pdf_by_unipipe
-
-def convert_pdf(pdf_file, enable_formula, enable_table, max_pages=100):
-    try:
-        # 获取上传文件路径
-        file_path = pdf_file.name
-        parent_path = os.path.dirname(file_path)
-        
-        # 初始化 MinerU 管道
-        image_writer = DiskReaderWriter(parent_path)
-        jso_useful_key = {
-            "model_list": [],
-            "enable_formula_recognition": enable_formula,
-            "enable_table_recognition": enable_table
-        }
-        pipe = UNIPipe(file_path, jso_useful_key, image_writer=image_writer)
-        
-        # 执行 PDF 解析
-        pipe.pipe_classify()
-        pipe.pipe_parse()
-        
-        # 转换为 Markdown
-        md_content = parse_pdf_by_unipipe(file_path, pipe.pipe_mkユニ, max_pages=max_pages)
-        
-        # 保存 Markdown 到临时文件
-        output_dir = "output"
-        os.makedirs(output_dir, exist_ok=True)
-        output_file = os.path.join(output_dir, f"{os.path.basename(file_path)}.md")
-        with open(output_file, "w", encoding="utf-8") as f:
-            f.write(md_content)
-        
-        return md_content, output_file
-    except Exception as e:
-        return f"Error: {str(e)}", None
-
-# 定义 Gradio 界面
-with gr.Blocks(title="MinerU PDF to Markdown Converter") as iface:
-    gr.Markdown("# MinerU PDF to Markdown Converter")
-    gr.Markdown("Upload a PDF file, configure options, and convert to Markdown.")
-    
-    with gr.Row():
-        with gr.Column():
-            pdf_input = gr.File(label="Upload PDF File", file_types=[".pdf"])
-            formula_checkbox = gr.Checkbox(label="Enable formula recognition", value=False)
-            table_checkbox = gr.Checkbox(label="Enable table recognition", value=False)
-            max_pages = gr.Number(label="Max Pages to Process", value=100, precision=0)
-            convert_button = gr.Button("Convert to Markdown")
-        
-        with gr.Column():
-            markdown_output = gr.Markdown(label="Markdown Output")
-            download_button = gr.File(label="Download Markdown File")
-    
-    # 绑定按钮事件
-    convert_button.click(
-        fn=convert_pdf,
-        inputs=[pdf_input, formula_checkbox, table_checkbox, max_pages],
-        outputs=[markdown_output, download_button]
-    )
-
-# 启动 Gradio 应用
-iface.launch()
+import gradio as gr
+import fitz  # PyMuPDF
+from pymupdf4llm import get_markdown
+import os
+import tempfile
+import aiofiles
+from typing import Optional
+import markdown2
+from pathlib import Path
+
+class MinerUProcessor:
+    def __init__(self):
+        self.supported_formats = ['.pdf', '.png', '.jpg', '.jpeg', '.bmp', '.tiff']
+    
+    async def process_file(self, file_path: str, enable_formula: bool, enable_table: bool) -> str:
+        """处理文件并返回markdown内容"""
+        try:
+            file_ext = os.path.splitext(file_path)[1].lower()
+            
+            if file_ext == '.pdf':
+                # 处理PDF文件
+                markdown_text = get_markdown(
+                    file_path,
+                    write_images=False,
+                    use_latex=enable_formula,
+                    use_table=enable_table
+                )
+            else:
+                # 处理图片文件（转换为PDF再处理）
+                pdf_path = await self.image_to_pdf(file_path)
+                markdown_text = get_markdown(
+                    pdf_path,
+                    write_images=False,
+                    use_latex=enable_formula,
+                    use_table=enable_table
+                )
+                # 清理临时PDF文件
+                if os.path.exists(pdf_path):
+                    os.remove(pdf_path)
+            
+            return markdown_text
+            
+        except Exception as e:
+            return f"处理文件时出错: {str(e)}"
+    
+    async def image_to_pdf(self, image_path: str) -> str:
+        """将图片转换为PDF"""
+        doc = fitz.open()
+        img = fitz.open(image_path)
+        rect = img[0].rect
+        pdf_bytes = img.convert_to_pdf()
+        img.close()
+        
+        img_pdf = fitz.open("pdf", pdf_bytes)
+        page = doc.new_page(width=rect.width, height=rect.height)
+        page.show_pdf_page(rect, img_pdf, 0)
+        
+        # 保存临时PDF文件
+        temp_pdf = tempfile.NamedTemporaryFile(suffix='.pdf', delete=False)
+        doc.save(temp_pdf.name)
+        doc.close()
+        img_pdf.close()
+        
+        return temp_pdf.name
+
+async def process_document(
+    file: Optional[gr.components.File] = None,
+    enable_formula: bool = True,
+    enable_table: bool = True
+):
+    """处理文档的主函数"""
+    if file is None:
+        return "请上传文件", None, None
+    
+    try:
+        processor = MinerUProcessor()
+        
+        # 保存上传的文件
+        file_path = file.name
+        file_ext = os.path.splitext(file_path)[1].lower()
+        
+        if file_ext not in processor.supported_formats:
+            return f"不支持的文件格式: {file_ext}", None, None
+        
+        # 处理文件
+        markdown_content = await processor.process_file(file_path, enable_formula, enable_table)
+        
+        # 生成HTML预览
+        html_preview = markdown2.markdown(markdown_content)
+        
+        # 生成下载文件名
+        original_name = os.path.basename(file_path)
+        base_name = os.path.splitext(original_name)[0]
+        download_filename = f"{base_name}.md"
+        
+        return markdown_content, html_preview, download_filename
+        
+    except Exception as e:
+        return f"处理过程中出错: {str(e)}", None, None
+
+def create_download_file(markdown_content: str, filename: str):
+    """创建下载文件"""
+    if not markdown_content or markdown_content.startswith("处理过程中出错") or markdown_content.startswith("请上传文件"):
+        return None
+    
+    # 创建临时文件
+    temp_file = tempfile.NamedTemporaryFile(suffix='.md', delete=False, mode='w', encoding='utf-8')
+    temp_file.write(markdown_content)
+    temp_file.close()
+    
+    return temp_file.name
+
+# 创建Gradio界面
+with gr.Blocks(title="MinerU Document Processor", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 📄 MinerU Document Processor
+    
+    将PDF和图片文档转换为Markdown格式，支持公式和表格识别。
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            file_input = gr.File(
+                label="上传文档",
+                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".tiff"],
+                type="filepath"
+            )
+            
+            with gr.Group():
+                gr.Markdown("### 识别选项")
+                enable_formula = gr.Checkbox(
+                    label="Enable formula recognition", 
+                    value=True,
+                    info="是否启用公式识别"
+                )
+                enable_table = gr.Checkbox(
+                    label="Enable table recognition", 
+                    value=True,
+                    info="是否启用表格识别"
+                )
+            
+            process_btn = gr.Button("🚀 开始处理", variant="primary")
+            
+            download_btn = gr.DownloadButton(
+                "📥 下载Markdown文件",
+                visible=False
+            )
+        
+        with gr.Column(scale=2):
+            with gr.Tab("Markdown预览"):
+                markdown_output = gr.Code(
+                    label="转换结果 (Markdown)",
+                    language="markdown",
+                    lines=20,
+                    interactive=False
+                )
+            
+            with gr.Tab("HTML预览"):
+                html_output = gr.HTML(
+                    label="HTML预览"
+                )
+    
+    # 设置事件处理
+    process_btn.click(
+        fn=process_document,
+        inputs=[file_input, enable_formula, enable_table],
+        outputs=[markdown_output, html_output, download_btn]
+    )
+    
+    # 下载处理
+    download_btn.click(
+        fn=create_download_file,
+        inputs=[markdown_output, download_btn],
+        outputs=download_btn
+    )
+    
+    # 示例
+    gr.Examples(
+        examples=[
+            ["https://example.com/sample.pdf", True, True],
+            ["https://example.com/sample.png", False, True]
+        ],
+        inputs=[file_input, enable_formula, enable_table],
+        outputs=[markdown_output, html_output, download_btn],
+        fn=process_document,
+        cache_examples=False
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )