ocr_workbench/export.py

from __future__ import annotations

import base64
import binascii
import json
import mimetypes
import re
import shutil
from pathlib import Path
from typing import Any

from .documents import DocumentPage


_IMAGE_EXTENSION = {
    "image/png": ".png",
    "image/jpeg": ".jpg",
    "image/webp": ".webp",
}


def _decode_image(payload: dict[str, Any], output_stem: Path) -> Path | None:
    encoded = payload.get("annotated_image_base64")
    if not encoded:
        return None
    mime = str(payload.get("annotated_image_mime") or "image/png")
    extension = _IMAGE_EXTENSION.get(mime) or mimetypes.guess_extension(mime) or ".png"
    output = output_stem.with_suffix(extension)
    try:
        output.write_bytes(base64.b64decode(encoded, validate=True))
    except (binascii.Error, ValueError) as exc:
        raise ValueError(f"Worker returned an invalid base64 image: {exc}") from exc
    return output


def _plain_from_markdown(markdown: str) -> str:
    text = re.sub(r"```[^\n]*\n(.*?)```", r"\1", markdown, flags=re.DOTALL)
    text = re.sub(r"!\[[^\]]*\]\([^)]*\)", "", text)
    text = re.sub(r"\[([^\]]+)\]\([^)]*\)", r"\1", text)
    text = re.sub(r"<[^>]+>", "", text)
    text = re.sub(r"^[#>*+-]+\s*", "", text, flags=re.MULTILINE)
    text = re.sub(r"[`*_~]", "", text)
    return text.strip()


def build_exports(
    *,
    run_dir: Path,
    model_id: str,
    model_label: str,
    pages: list[DocumentPage],
    responses: list[dict[str, Any]],
) -> tuple[list[tuple[str, str]], str, str, str, str]:
    output_dir = run_dir / "results"
    annotated_dir = output_dir / "annotated"
    annotated_dir.mkdir(parents=True, exist_ok=True)

    annotated_gallery: list[tuple[str, str]] = []
    markdown_parts = [f"# OCR result — {model_label}"]
    text_parts: list[str] = []
    raw_pages: list[dict[str, Any]] = []

    for page, response in zip(pages, responses, strict=True):
        image = _decode_image(
            response,
            annotated_dir / f"page_{page.page_number:04d}",
        )
        if image is None:
            image = annotated_dir / f"page_{page.page_number:04d}.png"
            shutil.copy2(page.image_path, image)
        annotated_gallery.append((str(image), f"Page {page.page_number}"))

        markdown = str(response.get("markdown") or "").strip()
        text = str(response.get("text") or "").strip()
        if not text and markdown:
            text = _plain_from_markdown(markdown)
        warnings = [str(item) for item in (response.get("warnings", []) or [])]
        warning_markdown = ""
        warning_text = ""
        if warnings:
            warning_markdown = "\n\n> ⚠ " + "\n> ⚠ ".join(warnings)
            warning_text = "\n\n[Warnings]\n- " + "\n- ".join(warnings)
        markdown_parts.append(
            f"\n## Page {page.page_number}\n\n"
            f"{markdown or text or '(empty result)'}{warning_markdown}"
        )
        text_parts.append(
            f"===== Page {page.page_number} =====\n"
            f"{text or '(empty result)'}{warning_text}"
        )

        raw_page = dict(response)
        raw_page.pop("annotated_image_base64", None)
        raw_page["page_number"] = page.page_number
        raw_pages.append(raw_page)

    markdown_document = "\n".join(markdown_parts).strip() + "\n"
    text_document = "\n\n".join(text_parts).strip() + "\n"
    raw_payload = {
        "schema_version": "1.0",
        "model": model_id,
        "model_label": model_label,
        "pages": raw_pages,
    }
    json_document = json.dumps(raw_payload, ensure_ascii=False, indent=2)

    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "result.md").write_text(markdown_document, encoding="utf-8")
    (output_dir / "result.txt").write_text(text_document, encoding="utf-8")
    (output_dir / "result.json").write_text(json_document, encoding="utf-8")

    archive_base = run_dir.parent / f"{run_dir.name}_ocr_result"
    archive = shutil.make_archive(str(archive_base), "zip", run_dir)
    return annotated_gallery, markdown_document, text_document, json_document, archive