Upload gradio_app.py

gradio_app.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import json
+import shutil
+import uuid
+from pathlib import Path
+
+import spaces
+import fitz
+import gradio as gr
+
+from app import pdf_engine as editor
+
+
+def upload_pdf_gradio(pdf_file):
+    if not pdf_file:
+        raise gr.Error("Please upload a PDF.")
+
+    source = Path(pdf_file)
+    if source.suffix.lower() != ".pdf":
+        raise gr.Error("Please upload a PDF file.")
+
+    try:
+        raw = source.read_bytes()
+    except Exception as exc:
+        raise gr.Error(f"Could not read PDF: {exc}")
+
+    if len(raw) > editor.MAX_UPLOAD_BYTES:
+        raise gr.Error(
+            f"PDF must be smaller than "
+            f"{editor.MAX_UPLOAD_BYTES // (1024 * 1024)} MB."
+        )
+
+    if not raw.startswith(b"%PDF-"):
+        raise gr.Error("The uploaded file is not a valid PDF.")
+
+    try:
+        pdf = fitz.open(stream=raw, filetype="pdf")
+    except Exception as exc:
+        raise gr.Error("Could not read this PDF.") from exc
+
+    if pdf.needs_pass:
+        pdf.close()
+        raise gr.Error("Password-protected PDFs are not supported.")
+
+    if len(pdf) == 0 or len(pdf) > editor.MAX_PAGES:
+        count = len(pdf)
+        pdf.close()
+        raise gr.Error(
+            f"PDF must have between 1 and {editor.MAX_PAGES} pages. "
+            f"This file has {count} pages."
+        )
+
+    document_id = uuid.uuid4().hex
+    target = editor.DATA_DIR / document_id
+
+    try:
+        target.mkdir(parents=True, exist_ok=False)
+        (target / "original.pdf").write_bytes(raw)
+
+        objects = editor.extract_native_objects(pdf, target)
+
+        pages = [
+            {
+                "number": i,
+                "width": float(page.rect.width),
+                "height": float(page.rect.height),
+                "hasText": any(
+                    obj["page"] == i
+                    and obj.get("type", "text") == "text"
+                    for obj in objects
+                ),
+                "rotation": int(page.rotation),
+                "analyzed": False,
+                "regions": [],
+            }
+            for i, page in enumerate(pdf)
+        ]
+
+        model = {
+            "id": document_id,
+            "filename": source.name,
+            "pages": pages,
+            "objects": objects,
+        }
+
+        (target / "document.json").write_text(
+            json.dumps(model, ensure_ascii=False, indent=2),
+            encoding="utf-8",
+        )
+
+    except Exception:
+        shutil.rmtree(target, ignore_errors=True)
+        raise
+    finally:
+        pdf.close()
+
+    with editor._document_storage_lock:
+        editor.purge_document_storage(keep_id=document_id)
+
+    return (
+        document_id,
+        0,
+        json.dumps(
+            [obj for obj in objects if obj["page"] == 0],
+            ensure_ascii=False,
+            indent=2,
+        ),
+        "[]",
+        render_page(document_id, 0),
+        f"### PDF loaded\n\n**File:** `{source.name}`\n\n"
+        f"**Pages:** `{len(pages)}`",
+    )
+
+
+def render_page(document_id: str, page_number: int):
+    if not document_id:
+        return None
+
+    folder = editor.document_dir(document_id)
+    pdf_path = folder / "original.pdf"
+
+    if not pdf_path.exists():
+        return None
+
+    try:
+        page_number = int(page_number)
+        pdf = fitz.open(pdf_path)
+
+        if page_number < 0 or page_number >= len(pdf):
+            pdf.close()
+            return None
+
+        page = pdf[page_number]
+        pixmap = page.get_pixmap(
+            matrix=fitz.Matrix(1.5, 1.5),
+            alpha=False,
+        )
+
+        image_path = folder / f"gradio-page-{page_number}.png"
+        pixmap.save(image_path)
+        pdf.close()
+
+        return str(image_path)
+
+    except Exception as exc:
+        raise gr.Error(f"Could not render page: {exc}")
+
+
+def change_page(document_id: str, page_number: int):
+    if not document_id:
+        return None
+
+    try:
+        return render_page(document_id, int(page_number))
+    except Exception as exc:
+        raise gr.Error(f"Could not change page: {exc}")
+
+
+@spaces.GPU(duration=120)
+def analyze_page_gradio(document_id: str, page_number: int):
+    if not document_id:
+        raise gr.Error("Upload a PDF first.")
+
+    try:
+        page_number = int(page_number)
+
+        result = editor.analyze_page(
+            document_id=document_id,
+            page_number=page_number,
+            refresh=False,
+        )
+
+        preview = render_page(document_id, page_number)
+
+        objects_json = json.dumps(
+            result["objects"],
+            indent=2,
+            ensure_ascii=False,
+        )
+
+        regions_json = json.dumps(
+            result["regions"],
+            indent=2,
+            ensure_ascii=False,
+        )
+
+        status = (
+            "### Analysis complete\n\n"
+            f"- OCR model: `{result['ocrModel']}`\n"
+            f"- Layout model: `{result['layoutModel']}`\n"
+            f"- OCR detections: `{result['ocrDetections']}`\n"
+            f"- Font model: `{result['fontModel']}`"
+        )
+
+        if result.get("fontModelError"):
+            status += f"\n\n> FontID: `{result['fontModelError']}`"
+
+        return preview, objects_json, regions_json, status
+
+    except editor.HTTPException as exc:
+        raise gr.Error(exc.detail)
+    except Exception as exc:
+        raise gr.Error(f"Page analysis failed: {exc}")
+
+
+def export_pdf_gradio(
+    document_id: str,
+    objects_json: str,
+    regions_json: str,
+):
+    if not document_id:
+        raise gr.Error("Upload a PDF first.")
+
+    try:
+        objects_data = json.loads(objects_json or "[]")
+        regions_data = json.loads(regions_json or "[]")
+
+        request = editor.ExportRequest(
+            objects=objects_data,
+            tableGroups=regions_data,
+        )
+
+        output_path = editor.export_pdf(
+            document_id=document_id,
+            request=request,
+        )
+
+        return str(output_path)
+
+    except json.JSONDecodeError as exc:
+        raise gr.Error(f"Invalid editor JSON: {exc}")
+    except editor.HTTPException as exc:
+        raise gr.Error(exc.detail)
+    except Exception as exc:
+        raise gr.Error(f"PDF export failed: {exc}")
+
+
+with gr.Blocks(title="Paperplane PDF Editor") as demo:
+    gr.Markdown(
+        """
+        # 📄 Paperplane PDF Editor
+
+        Upload a PDF, analyze its pages with OCR and layout detection,
+        inspect the detected objects, and export an edited PDF.
+        """
+    )
+
+    document_id = gr.State("")
+
+    with gr.Row():
+        pdf_upload = gr.File(
+            label="PDF",
+            file_types=[".pdf"],
+            type="filepath",
+        )
+        upload_button = gr.Button(
+            "Open PDF",
+            variant="primary",
+        )
+
+    status = gr.Markdown("Upload a PDF to begin.")
+
+    with gr.Row():
+        page_number = gr.Number(
+            label="Page",
+            value=0,
+            precision=0,
+            minimum=0,
+        )
+        analyze_button = gr.Button(
+            "Analyze Page",
+            variant="primary",
+        )
+
+    page_preview = gr.Image(
+        label="Page Preview",
+        type="filepath",
+    )
+
+    with gr.Row():
+        objects_editor = gr.Code(
+            label="Detected / Editable Objects",
+            language="json",
+            lines=20,
+        )
+        regions_editor = gr.Code(
+            label="Layout Regions",
+            language="json",
+            lines=20,
+        )
+
+    export_button = gr.Button(
+        "Export PDF",
+        variant="primary",
+    )
+
+    exported_pdf = gr.File(label="Edited PDF")
+
+    upload_button.click(
+        fn=upload_pdf_gradio,
+        inputs=pdf_upload,
+        outputs=[
+            document_id,
+            page_number,
+            objects_editor,
+            regions_editor,
+            page_preview,
+            status,
+        ],
+    )
+
+    page_number.change(
+        fn=change_page,
+        inputs=[document_id, page_number],
+        outputs=page_preview,
+    )
+
+    analyze_button.click(
+        fn=analyze_page_gradio,
+        inputs=[document_id, page_number],
+        outputs=[
+            page_preview,
+            objects_editor,
+            regions_editor,
+            status,
+        ],
+    )
+
+    export_button.click(
+        fn=export_pdf_gradio,
+        inputs=[
+            document_id,
+            objects_editor,
+            regions_editor,
+        ],
+        outputs=exported_pdf,
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)