app.py

import gradio as gr
import os
import sys
import subprocess
import tempfile
from pathlib import Path
import json
from loguru import logger
import shutil

# ============================================================================
# AUTOMATIC SETUP: GPU Support, MinerU & Model Downloads
# ============================================================================

def run_command(cmd, description="", show_output=False):
    """Run a shell command"""
    try:
        if description:
            print(f"[SETUP] {description}...")

        if show_output:
            result = subprocess.run(cmd, shell=True, text=True)
        else:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)

        if result.returncode == 0:
            if description:
                print(f"[SETUP] ✅ {description} completed")
            return True, result.stdout if hasattr(result, 'stdout') else ""
        else:
            if hasattr(result, 'stderr'):
                print(f"[SETUP] ⚠️ {result.stderr}")
            return False, result.stderr if hasattr(result, 'stderr') else ""
    except Exception as e:
        print(f"[SETUP] ❌ Error: {e}")
        return False, str(e)

def setup_environment():
    """Setup MinerU environment with GPU support"""
    print("=" * 70)
    print("🚀 MINERU OCR TOOL - SETUP WITH GPU & DOCX/PDF EXPORT")
    print("=" * 70)

    # Check GPU first
    try:
        import torch
        if torch.cuda.is_available():
            print(f"[SETUP] ✅ GPU Detected: {torch.cuda.get_device_name(0)}")
            print(f"[SETUP] ✅ CUDA Version: {torch.version.cuda}")
            print(f"[SETUP] ✅ GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
            gpu_available = True
        else:
            print("[SETUP] ⚠️ No GPU detected, will use CPU")
            gpu_available = False
    except ImportError:
        print("[SETUP] Installing PyTorch...")
        run_command(f"{sys.executable} -m pip install torch torchvision --upgrade", "Installing PyTorch")
        import torch
        gpu_available = torch.cuda.is_available()

    # Install MinerU from GitHub dev branch (better quality than PyPI)
    print("\n" + "=" * 70)
    print("[SETUP] Installing MinerU from GitHub dev branch...")
    print("[SETUP] (Better OCR quality than PyPI version)")
    print("=" * 70)

    # Uninstall old version
    print("[SETUP] Removing old MinerU installation...")
    os.system('pip uninstall -y mineru')

    # Install from GitHub dev branch
    print("[SETUP] Installing from GitHub (this may take a few minutes)...")
    os.system('pip install git+https://github.com/myhloli/Magic-PDF.git@dev')

    # Install additional packages
    print("\n" + "=" * 70)
    print("[SETUP] Installing additional packages...")
    print("=" * 70)

    packages = [
        "mineru-vl-utils",
        "gradio-pdf",
        "loguru",
        "pypandoc",
    ]

    # Add VLLM for GPU acceleration if GPU is available
    if gpu_available:
        print("[SETUP] Installing VLLM for GPU acceleration...")
        packages.append("vllm==0.10.1.1")

    for package in packages:
        run_command(
            f"{sys.executable} -m pip install '{package}' --upgrade",
            f"Installing {package}"
        )

    # Install pandoc system dependency
    print("\n" + "=" * 70)
    print("[SETUP] Installing Pandoc for document conversion...")
    print("=" * 70)
    try:
        import pypandoc
        # Download pandoc if not installed
        pypandoc.ensure_pandoc_installed()
        print("[SETUP] ✅ Pandoc installed successfully")
    except Exception as e:
        print(f"[SETUP] ⚠️ Pandoc installation warning: {e}")
        print("[SETUP] Document conversion may not work without pandoc")

    # Download models
    print("\n" + "=" * 70)
    print("[SETUP] Downloading MinerU models...")
    print("=" * 70)

    # Check if models already exist
    model_dir = Path.home() / ".cache" / "mineru"
    if model_dir.exists() and any(model_dir.rglob("*")):
        print("[SETUP] ✅ Models already downloaded")
    else:
        print("[SETUP] Downloading models (this may take 5-10 minutes)...")
        success, output = run_command(
            "mineru-models-download -s huggingface -m all",
            "Downloading models",
            show_output=True
        )
        if success:
            print("[SETUP] ✅ Models downloaded successfully")
        else:
            print("[SETUP] ⚠️ Models will be downloaded on first use")

    # Configure MinerU for GPU
    print("\n" + "=" * 70)
    print("[SETUP] Configuring MinerU...")
    print("=" * 70)

    config_path = Path.home() / "mineru.json"
    if config_path.exists():
        try:
            with open(config_path, 'r+') as file:
                config = json.load(file)

                # Set LaTeX delimiters
                delimiters = {
                    'display': {'left': '\\[', 'right': '\\]'},
                    'inline': {'left': '\\(', 'right': '\\)'}
                }
                config['latex-delimiter-config'] = delimiters

                # Enable GPU if available
                if gpu_available:
                    if 'device-mode' in config:
                        config['device-mode'] = 'cuda'
                    print("[SETUP] ✅ GPU mode enabled in config")

                file.seek(0)
                file.truncate()
                json.dump(config, file, indent=4)
                print("[SETUP] ✅ Configuration updated")
        except Exception as e:
            logger.warning(f"Could not update config: {e}")

    print("\n" + "=" * 70)
    print("✅ SETUP COMPLETE!")
    print("=" * 70 + "\n")

    return gpu_available

# Run setup
gpu_available = setup_environment()

# Import required modules after installation
try:
    import torch
    from gradio_pdf import PDF

    # GPU info for display
    if torch.cuda.is_available():
        gpu_info = f"🚀 GPU: {torch.cuda.get_device_name(0)} ({torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB VRAM)"
    else:
        gpu_info = "💻 CPU Mode"
except ImportError as e:
    print(f"[WARNING] Some imports failed: {e}")
    gpu_info = "💻 CPU Mode"
    gpu_available = False
    PDF = None

print(f"[STARTUP] Running with: {gpu_info}")

# Conversion functions with LaTeX support
def convert_markdown_to_docx(markdown_content, output_path):
    """Convert markdown (with LaTeX) to DOCX using pypandoc"""
    try:
        import pypandoc

        # Convert markdown with LaTeX formulas to DOCX
        # Pandoc will convert LaTeX math to Word equations
        pypandoc.convert_text(
            markdown_content,
            'docx',
            format='markdown+tex_math_dollars',  # Support $...$ and $$...$$ LaTeX
            outputfile=output_path,
            extra_args=[
                '--standalone',
                '--mathml',  # Convert LaTeX math to MathML for Word
            ]
        )
        print(f"[CONVERT] ✅ DOCX with LaTeX formulas created")
        return True, output_path
    except Exception as e:
        print(f"[CONVERT] Error converting to DOCX: {e}")
        # Fallback: try without LaTeX support
        try:
            pypandoc.convert_text(
                markdown_content,
                'docx',
                format='md',
                outputfile=output_path,
                extra_args=['--standalone']
            )
            print(f"[CONVERT] ⚠️ DOCX created without LaTeX support")
            return True, output_path
        except Exception as e2:
            print(f"[CONVERT] DOCX conversion failed: {e2}")
            return False, str(e2)

def convert_markdown_to_pdf(markdown_content, output_path):
    """Convert markdown (with LaTeX) to PDF using pypandoc"""
    try:
        import pypandoc

        # Convert markdown with LaTeX formulas to PDF
        pypandoc.convert_text(
            markdown_content,
            'pdf',
            format='markdown+tex_math_dollars',  # Support $...$ and $$...$$ LaTeX
            outputfile=output_path,
            extra_args=[
                '--pdf-engine=pdflatex',
                '--standalone',
                '-V', 'geometry:margin=1in',  # Better margins
            ]
        )
        print(f"[CONVERT] ✅ PDF with LaTeX formulas created")
        return True, output_path
    except Exception as e:
        print(f"[CONVERT] Error with pdflatex: {e}")
        # Fallback 1: Try xelatex
        try:
            pypandoc.convert_text(
                markdown_content,
                'pdf',
                format='markdown+tex_math_dollars',
                outputfile=output_path,
                extra_args=['--pdf-engine=xelatex', '--standalone']
            )
            print(f"[CONVERT] ✅ PDF created with xelatex")
            return True, output_path
        except Exception as e2:
            print(f"[CONVERT] xelatex failed: {e2}")
            # Fallback 2: Try without LaTeX engine (may not render formulas)
            try:
                pypandoc.convert_text(
                    markdown_content,
                    'pdf',
                    format='md',
                    outputfile=output_path,
                    extra_args=['--standalone']
                )
                print(f"[CONVERT] ⚠️ PDF created without LaTeX formula support")
                return True, output_path
            except Exception as e3:
                print(f"[CONVERT] PDF conversion failed completely: {e3}")
                return False, str(e3)

def process_with_mineru(input_path, output_base_dir, use_vllm=True):
    """Process file with MinerU CLI"""
    try:
        output_dir = Path(output_base_dir) / "output"
        output_dir.mkdir(parents=True, exist_ok=True)

        # Build command
        cmd = f'mineru -p "{input_path}" -o "{output_dir}"'

        print(f"[PROCESS] Running MinerU on: {Path(input_path).name}")
        print(f"[PROCESS] Command: {cmd}")

        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            timeout=600  # 10 minute timeout
        )

        if result.returncode == 0:
            print("[PROCESS] ✅ Processing completed")

            # Find output files
            md_files = list(output_dir.glob("**/*.md"))
            json_files = list(output_dir.glob("**/*.json"))

            md_content = ""
            json_content = ""

            # Read markdown output
            if md_files:
                print(f"[PROCESS] Found {len(md_files)} markdown file(s)")
                # Sort by size, get the largest (usually the main content)
                md_files.sort(key=lambda x: x.stat().st_size, reverse=True)
                with open(md_files[0], 'r', encoding='utf-8') as f:
                    md_content = f.read()

            # Read JSON output
            if json_files:
                print(f"[PROCESS] Found {len(json_files)} JSON file(s)")
                # Find the content.json or result.json file
                for jf in json_files:
                    if 'content' in jf.name.lower() or 'result' in jf.name.lower():
                        with open(jf, 'r', encoding='utf-8') as f:
                            json_content = f.read()
                        break

                # If no specific file found, use the largest one
                if not json_content and json_files:
                    json_files.sort(key=lambda x: x.stat().st_size, reverse=True)
                    with open(json_files[0], 'r', encoding='utf-8') as f:
                        json_content = f.read()

            if not md_content and not json_content:
                all_files = list(output_dir.glob("**/*"))
                print(f"[PROCESS] Found {len(all_files)} total files in output")
                return False, "No markdown or JSON output found", "", "", output_dir

            return True, "Success", md_content, json_content, output_dir

        else:
            error_msg = result.stderr if result.stderr else result.stdout
            print(f"[PROCESS] ❌ Error: {error_msg}")
            return False, error_msg, "", "", None

    except subprocess.TimeoutExpired:
        return False, "Processing timeout (>10 minutes)", "", "", None
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"[ERROR] {error_details}")
        return False, str(e), "", "", None

def download_as_docx(markdown_content, original_filename="document"):
    """Convert markdown to DOCX and return file path for download"""
    if not markdown_content or markdown_content.strip() == "":
        return None

    try:
        # Create temp file
        temp_dir = Path(tempfile.gettempdir()) / "mineru_gradio"
        temp_dir.mkdir(exist_ok=True)

        base_name = Path(original_filename).stem if original_filename else "document"
        output_path = temp_dir / f"{base_name}_extracted.docx"

        success, result = convert_markdown_to_docx(markdown_content, str(output_path))

        if success:
            print(f"[DOWNLOAD] DOCX created: {output_path}")
            return str(output_path)
        else:
            print(f"[DOWNLOAD] DOCX conversion failed: {result}")
            return None
    except Exception as e:
        print(f"[DOWNLOAD] Error creating DOCX: {e}")
        return None

def download_as_pdf(markdown_content, original_filename="document"):
    """Convert markdown to PDF and return file path for download"""
    if not markdown_content or markdown_content.strip() == "":
        return None

    try:
        # Create temp file
        temp_dir = Path(tempfile.gettempdir()) / "mineru_gradio"
        temp_dir.mkdir(exist_ok=True)

        base_name = Path(original_filename).stem if original_filename else "document"
        output_path = temp_dir / f"{base_name}_extracted.pdf"

        success, result = convert_markdown_to_pdf(markdown_content, str(output_path))

        if success:
            print(f"[DOWNLOAD] PDF created: {output_path}")
            return str(output_path)
        else:
            print(f"[DOWNLOAD] PDF conversion failed: {result}")
            return None
    except Exception as e:
        print(f"[DOWNLOAD] Error creating PDF: {e}")
        return None

# Store current filename for download functions
current_filename = {"name": "document"}

def process_file(file, use_gpu=True):
    """Process uploaded file"""
    if file is None:
        return (
            None,  # PDF preview
            "❌ No file uploaded. Please upload a PDF or image file.",
            "",  # Markdown
            "",  # JSON
            None,  # DOCX download
            None   # PDF download
        )

    try:
        file_path = Path(file.name)
        file_ext = file_path.suffix.lower()

        print(f"\n[PROCESS] ========================================")
        print(f"[PROCESS] Processing: {file_path.name}")
        print(f"[PROCESS] Type: {file_ext}")
        print(f"[PROCESS] Size: {file_path.stat().st_size / 1024:.1f} KB")
        print(f"[PROCESS] ========================================")

        if file_ext not in ['.pdf', '.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
            return (
                None,
                f"❌ Unsupported file type: {file_ext}\n\nSupported formats: PDF, PNG, JPG, JPEG, BMP, TIFF",
                "",
                "",
                None,
                None
            )

        # Store filename for download functions
        current_filename["name"] = file_path.name

        # Create persistent temp directory for this session
        temp_base = Path(tempfile.gettempdir()) / "mineru_gradio"
        temp_base.mkdir(exist_ok=True)

        # Convert images to PDF
        input_file = file_path
        pdf_preview_path = file_path if file_ext == '.pdf' else None

        if file_ext in ['.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif']:
            print("[PROCESS] Converting image to PDF...")
            from PIL import Image

            img = Image.open(file_path)
            if img.mode in ('RGBA', 'LA', 'P'):
                img = img.convert('RGB')

            # Create temp PDF
            temp_pdf = temp_base / f"{file_path.stem}.pdf"
            img.save(temp_pdf, "PDF", resolution=100.0)
            input_file = temp_pdf
            pdf_preview_path = temp_pdf
            print(f"[PROCESS] ✅ Converted to: {temp_pdf}")

        # Process with MinerU
        success, message, md_content, json_content, output_dir = process_with_mineru(
            str(input_file),
            str(temp_base),
            use_vllm=use_gpu and gpu_available
        )

        if success:
            # Count extracted elements
            pages = md_content.count('\n## ') if md_content else 0
            tables = md_content.count('|') // 4 if md_content else 0

            status = f"""✅ **Processing Complete!**

📄 **File:** {file_path.name}
📑 **Type:** {file_ext.upper()}
⚡ **Device:** {gpu_info}
📊 **Pages:** {pages if pages > 0 else 'N/A'}
📋 **Tables:** ~{tables} detected
⏱️ **Status:** Successfully extracted and parsed

---
**Ready!** View the extracted content in the tabs below or download as DOCX/PDF.
"""
            # Generate download files
            docx_file = download_as_docx(md_content, file_path.name) if md_content else None
            pdf_file = download_as_pdf(md_content, file_path.name) if md_content else None

            return (
                str(pdf_preview_path) if pdf_preview_path else None,
                status,
                md_content if md_content else "No markdown content generated",
                json_content if json_content else json.dumps({"status": "no content"}, indent=2),
                docx_file,
                pdf_file
            )
        else:
            error_status = f"""❌ **Processing Failed**

📄 **File:** {file_path.name}
⚡ **Device:** {gpu_info}

**Error Details:**
```
{message}
```

**Troubleshooting:**
- Ensure the file is not corrupted
- Try a smaller file first
- Check console logs for details
- For images, ensure they contain readable text
"""
            return (
                str(pdf_preview_path) if pdf_preview_path else None,
                error_status,
                "",
                "",
                None,
                None
            )

    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"[ERROR] {error_details}")
        return (
            None,
            f"❌ **Unexpected Error**\n\n```\n{str(e)}\n```\n\nSee console for full traceback.",
            "",
            "",
            None,
            None
        )

# Create Gradio Interface with improved layout
with gr.Blocks(
    title="MinerU OCR Tool",
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="cyan",
    ),
    css="""
    .gradio-container {
        max-width: 1400px !important;
    }
    .pdf-preview {
        height: 600px !important;
    }
    """
) as demo:

    # Header
    gr.Markdown(
        f"""
        # 🔮 MinerU - Advanced OCR & Document Parser

        Extract text, tables, and LaTeX formulas from PDFs and images with high precision.
        **GitHub Dev Branch** • Better quality than PyPI • **Export to DOCX/PDF with LaTeX support**

        <div style="padding: 10px; background: linear-gradient(90deg, #667eea 0%, #764ba2 100%); border-radius: 8px; color: white; text-align: center; margin: 10px 0;">
            <strong>{gpu_info}</strong>
        </div>
        """
    )

    with gr.Row():
        # Left column - Input and Preview
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Document")

            file_input = gr.File(
                label="Select PDF or Image",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif"],
                type="filepath",
                file_count="single"
            )

            with gr.Row():
                process_btn = gr.Button(
                    "🚀 Process Document",
                    variant="primary",
                    size="lg",
                    scale=3
                )

                use_gpu_checkbox = gr.Checkbox(
                    label="GPU",
                    value=gpu_available,
                    interactive=gpu_available,
                    scale=1,
                    info="Use GPU acceleration" if gpu_available else "No GPU"
                )

            gr.Markdown("---")

            gr.Markdown("### 👁️ Document Preview")

            if PDF is not None:
                pdf_preview = PDF(
                    label="PDF Preview",
                    height=600,
                    elem_classes=["pdf-preview"]
                )
            else:
                pdf_preview = gr.File(label="File Path", visible=False)

            gr.Markdown(
                f"""
                ---
                ### 📋 Supported Formats
                - **PDF**: Multi-page documents
                - **Images**: PNG, JPG, JPEG, BMP, TIFF

                ### ✨ Features
                - 🌍 OCR for 84+ languages
                - 📊 Table extraction
                - 🔢 Formula recognition
                - 📐 Layout preservation
                - {f"⚡ GPU acceleration (3-10x faster)" if gpu_available else "💻 CPU processing"}

                ### ⏱️ Processing Time
                - **First document:** 5-10 min (model download)
                - **Subsequent:** {f"~10-30s with GPU" if gpu_available else "~1-2min with CPU"}
                """
            )

        # Right column - Results
        with gr.Column(scale=1):
            gr.Markdown("### 📊 Processing Status")

            status_output = gr.Textbox(
                label="Status",
                lines=8,
                interactive=False,
                show_copy_button=False,
                placeholder="Upload a file and click 'Process Document' to begin..."
            )

            gr.Markdown("### 📄 Extracted Content")

            with gr.Tabs():
                with gr.Tab("📝 Markdown"):
                    gr.Markdown("**Human-readable format** - Easy to read and edit")
                    markdown_output = gr.Textbox(
                        label="Markdown Output",
                        lines=18,
                        max_lines=40,
                        interactive=False,
                        show_copy_button=True,
                        placeholder="Markdown content will appear here after processing..."
                    )

                    # Download buttons for converted formats
                    gr.Markdown("### 📥 Download Extracted Content")
                    gr.Markdown("**Includes LaTeX formulas** converted to Word equations (DOCX) or rendered math (PDF)")
                    with gr.Row():
                        docx_download = gr.File(
                            label="📄 Download as DOCX (with LaTeX)",
                            interactive=False,
                            visible=True
                        )
                        pdf_download = gr.File(
                            label="📕 Download as PDF (with LaTeX)",
                            interactive=False,
                            visible=True
                        )

                with gr.Tab("📋 JSON"):
                    gr.Markdown("**Structured data** - Machine-readable format with metadata")
                    json_output = gr.Textbox(
                        label="JSON Output",
                        lines=20,
                        max_lines=40,
                        interactive=False,
                        show_copy_button=True,
                        placeholder="JSON data will appear here after processing..."
                    )

    # Connect button
    process_btn.click(
        fn=process_file,
        inputs=[file_input, use_gpu_checkbox],
        outputs=[pdf_preview, status_output, markdown_output, json_output, docx_download, pdf_download]
    )

    # Footer
    gr.Markdown(
        """
        ---
        <div style="text-align: center; padding: 20px; background: #f5f5f5; border-radius: 8px;">
            <p style="margin: 5px 0;">
                <strong>Powered by</strong>
                <a href="https://github.com/opendatalab/MinerU" target="_blank">MinerU (Magic-PDF)</a>
            </p>
            <p style="margin: 5px 0; font-size: 0.9em;">
                📚 <a href="https://opendatalab.github.io/MinerU/" target="_blank">Documentation</a> |
                💬 <a href="https://discord.gg/Tdedn9GTXq" target="_blank">Discord</a> |
                🌐 <a href="https://mineru.net" target="_blank">Official Demo</a>
            </p>
            <p style="margin: 5px 0; font-size: 0.8em; color: #666;">
                MinerU converts PDFs to machine-readable formats • Supports 84+ languages • Open Source
            </p>
        </div>
        """
    )

# Launch
if __name__ == "__main__":
    print("\n" + "=" * 70)
    print("🚀 Starting MinerU OCR Gradio Interface")
    print("=" * 70)
    print(f"Device: {gpu_info}")
    print(f"PDF Preview: {'Enabled' if PDF is not None else 'Disabled (install gradio-pdf)'}")
    print("=" * 70 + "\n")

    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        show_api=False
    )