File size: 4,939 Bytes
fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 0ebf67e fbfc939 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | 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)
|