backend/gradio_ui/ocr_tab.py

# ocr_tab.py
"""Structured document OCR panel."""

from __future__ import annotations

import base64
import io
import json
from typing import Any, Optional, Tuple

import gradio as gr
from PIL import Image

from core.bootstrap import get_chart_extractor_service
from gradio_ui.renderers import render_structured_ocr


def _preview_image(preview_b64: Optional[str]):
    if not preview_b64:
        return None
    try:
        raw = base64.b64decode(preview_b64)
        return Image.open(io.BytesIO(raw))
    except Exception:
        return None


def run_structured_ocr(
    file_path: str | None, page_num: int
) -> Tuple[Optional[Any], str, gr.update, str, str]:
    if not file_path:
        return (
            None,
            "<p class='panel-hint'>Upload a document to extract structured OCR.</p>",
            gr.update(),
            "",
            "0 fields / rows",
        )

    with open(file_path, "rb") as handle:
        content = handle.read()
    filename = file_path.rsplit("\\", 1)[-1].rsplit("/", 1)[-1]

    extractor = get_chart_extractor_service()
    structured = extractor.extract_structured(content, filename)
    preview_b64 = extractor.preview_page(content, filename, page_num=page_num)

    pages = structured.get("pages") or []
    page_choices = [p.get("page_number", idx + 1) for idx, p in enumerate(pages)]
    if not page_choices:
        page_choices = [1]
    if page_num not in page_choices:
        page_num = page_choices[0]

    html = render_structured_ocr(structured, page_num=page_num)
    image = _preview_image(preview_b64)
    field_count = 0
    if pages:
        target = next((p for p in pages if p.get("page_number") == page_num), pages[0])
        for section in target.get("sections") or []:
            field_count += len(section.get("fields") or {})
            field_count += len(section.get("rows") or [])
    meta = f'<span class="ocr-meta">{field_count} fields / rows</span>'
    return image, html, gr.update(choices=page_choices, value=page_num), json.dumps(structured), meta


def on_page_change(structured_json: str, page_num: int) -> str:
    if not structured_json:
        return "<p class='panel-hint'>No structured data.</p>"
    return render_structured_ocr(json.loads(structured_json), page_num=page_num)


def build_ocr_panel() -> dict:
    structured_state = gr.State("")
    with gr.Column(elem_classes=["fs-ocr-panel"]):
        ocr_file = gr.File(
            label="Upload PDF or image for OCR · PDF, PNG, JPG, TIFF",
            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".webp"],
            type="filepath",
            elem_classes=["fs-upload-zone", "fs-ocr-upload"],
            file_count="single",
            height=140,
        )
        with gr.Row(elem_classes=["fs-ocr-toolbar"]):
            page_selector = gr.Radio(
                label="Page",
                choices=[1],
                value=1,
                interactive=True,
                elem_classes=["fs-ocr-pages"],
            )
            ocr_meta = gr.HTML('<span class="ocr-meta">0 fields / rows</span>', padding=False)
            ocr_btn = gr.Button("Extract", elem_classes=["fs-gold-btn"], scale=0, min_width=100)
        with gr.Row(elem_classes=["fs-ocr-split"]):
            with gr.Column(elem_classes=["fs-ocr-pane", "fs-ocr-pane-left"]):
                gr.HTML('<div class="fs-pane-title">ORIGINAL DOCUMENT</div>', padding=False)
                preview = gr.Image(show_label=False, height=420, container=False)
            with gr.Column(elem_classes=["fs-ocr-pane", "fs-ocr-pane-right"]):
                gr.HTML(
                    """
                    <div class="fs-pane-title fs-pane-title-gold">STRUCTURED EXTRACTION</div>
                    <p class="fs-pane-sub">Labels and values taken from the document</p>
                    """,
                    padding=False,
                )
                structured_html = gr.HTML(
                    "<p class='panel-hint'>Upload a file and click Extract.</p>",
                    padding=False,
                )

    ocr_btn.click(
        run_structured_ocr,
        inputs=[ocr_file, page_selector],
        outputs=[preview, structured_html, page_selector, structured_state, ocr_meta],
        queue=True,
    )
    page_selector.change(
        on_page_change,
        inputs=[structured_state, page_selector],
        outputs=[structured_html],
    )

    return {
        "ocr_file": ocr_file,
        "ocr_btn": ocr_btn,
        "preview": preview,
        "page_selector": page_selector,
        "structured_html": structured_html,
    }