src/pino/ocr_client.py

from __future__ import annotations

"""
Client for the Unlimited-OCR Hugging Face Space.

Uploads an image or PDF file and polls for extracted text.
"""

import json
import logging
import time
from pathlib import Path
from typing import Any

import requests

logger = logging.getLogger("pino.ocr_client")

SPACE_API = "https://akhaliq-unlimited-ocr.hf.space/gradio_api"


def upload_file(file_path: str | Path, token: str | None = None) -> dict[str, Any]:
    """Upload a file to the Space and return the Gradio FileData payload."""
    path = Path(file_path)
    url = f"{SPACE_API}/upload"
    headers = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    with path.open("rb") as f:
        response = requests.post(url, files={"files": f}, headers=headers, timeout=60)
    response.raise_for_status()
    data = response.json()
    # The upload endpoint returns a list of FileData objects.
    if isinstance(data, list):
        return data[0]
    return data


def call_ocr(file_data: dict[str, Any], endpoint: str = "run_ocr", token: str | None = None) -> str:
    """Call the OCR endpoint and return the event_id to poll."""
    url = f"{SPACE_API}/call/{endpoint}"
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"

    # Ensure file_data is a proper Gradio FileData dict.
    if isinstance(file_data, str):
        file_payload = {
            "path": file_data,
            "url": None,
            "size": None,
            "orig_name": None,
            "mime_type": None,
            "is_stream": False,
            "meta": {"_type": "gradio.FileData"},
        }
    else:
        file_payload = dict(file_data)
        if "meta" not in file_payload:
            file_payload["meta"] = {"_type": "gradio.FileData"}

    payload = {"data": [file_payload, "gundam", "document parsing"]}
    response = requests.post(url, json=payload, headers=headers, timeout=60)
    response.raise_for_status()
    data = response.json()
    logger.debug("OCR call response: %s", data)
    return str(data["event_id"])


def poll_result(event_id: str, endpoint: str = "run_ocr", token: str | None = None, timeout: float = 180.0) -> dict[str, Any]:
    """Poll the result endpoint until completion or timeout.

    Gradio streaming endpoints emit Server-Sent Events. We accumulate ``text``
    chunks until ``done`` is true, then return the concatenated result.
    """
    url = f"{SPACE_API}/call/{endpoint}/{event_id}"
    headers = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    deadline = time.time() + timeout
    collected_text_parts: list[str] = []
    while time.time() < deadline:
        response = requests.get(url, headers=headers, timeout=60)
        response.raise_for_status()
        done = False
        for line in response.text.splitlines():
            line = line.strip()
            if not line or not line.startswith("data:"):
                continue
            try:
                payload = json.loads(line[5:].strip())
            except json.JSONDecodeError:
                continue
            if isinstance(payload, dict):
                if payload.get("text"):
                    collected_text_parts.append(str(payload["text"]))
                if payload.get("done"):
                    done = True
            elif isinstance(payload, list) and len(payload) > 0:
                return {"result": payload}
        if done:
            return {"result": "".join(collected_text_parts)}
        time.sleep(1.0)
    raise TimeoutError(f"OCR polling timed out for event_id {event_id}")


def extract_text(file_path: str | Path, token: str | None = None) -> str:
    """End-to-end: upload a file, run OCR, and return the extracted text."""
    file_data = upload_file(file_path, token=token)
    event_id = call_ocr(file_data, endpoint="run_ocr", token=token)
    result = poll_result(event_id, endpoint="run_ocr", token=token)
    # The OCR result is typically a list of strings or a single string.
    if isinstance(result, dict) and "result" in result:
        data = result["result"]
    else:
        data = result
    if isinstance(data, list):
        return "\n".join(str(item) for item in data)
    return str(data)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Extract text from an image/PDF using Unlimited-OCR")
    parser.add_argument("file", help="Path to image or PDF file")
    parser.add_argument("--token", default=None, help="Hugging Face token")
    parser.add_argument("--output", default=None, help="Optional output text file")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    text = extract_text(args.file, token=args.token)
    if args.output:
        Path(args.output).write_text(text, encoding="utf-8")
    else:
        print(text)