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**
{gpu_info}
"""
)
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(
"""
---
"""
)
# 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
)