import gradio as gr
from mineru_vl_utils.mineru_client import MinerUClient
from PIL import Image
import fitz  # PyMuPDF pour lire les PDFs
import os

# Init client
model_path = "opendatalab/MinerU2.5-2509-1.2B"
client = MinerUClient(backend="transformers", model_path=model_path)

def extract_from_file(file):
    # Vérifier si PDF ou image
    ext = os.path.splitext(file.name)[-1].lower()
    images = []

    if ext == ".pdf":
        doc = fitz.open(file.name)
        for page in doc:
            pix = page.get_pixmap()
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            images.append(img)
    else:
        images.append(Image.open(file.name))

    results = []
    for img in images:
        blocks = client.two_step_extract(img)
        text_blocks = [b.text for b in blocks if hasattr(b, "text")]
        results.append("\n".join(text_blocks))

    return "\n\n--- PAGE ---\n\n".join(results)

demo = gr.Interface(
    fn=extract_from_file,
    inputs=gr.File(type="filepath", label="Upload PDF or Image"),
    outputs=gr.Textbox(label="Extracted Text", lines=20),
    title="MinerU2.5 Document Extractor",
    description="Upload a PDF or Image to extract structured text using MinerU2.5."
)

demo.launch()